diff --git a/CMakeLists.txt b/CMakeLists.txt index cbce88c..99df21b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -181,12 +181,9 @@ option(NEMO_SPEECH_CUBLAS_SHIM "Build the in-tree drop-in cuBLAS shim (nativ # 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/app/serve.cpp b/app/serve.cpp index f855754..a82dc02 100644 --- a/app/serve.cpp +++ b/app/serve.cpp @@ -228,6 +228,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) @@ -503,8 +506,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; - tts_config.runtime.sampling_backend = - nemo_speech::tts::MagpieBackendPreference::Cuda; } } tts_config.runtime.magpie_model = magpie_path; @@ -518,6 +519,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( @@ -618,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/app/synthesize.cpp b/app/synthesize.cpp index a820c87..c104d74 100644 --- a/app/synthesize.cpp +++ b/app/synthesize.cpp @@ -214,7 +214,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; - parsed.runtime.sampling_backend = nemo_speech::tts::MagpieBackendPreference::Cuda; 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 35319f0..70fb142 100644 --- a/config/server.example.yaml +++ b/config/server.example.yaml @@ -70,7 +70,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 @@ -84,7 +84,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 31148f2..3726661 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/conversion/tts.py b/conversion/tts.py index bee7b17..b58f859 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"] @@ -79,24 +80,69 @@ 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]: + 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)) + 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()] - - 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)) - n_codebooks = int( - len([k for k in sd if k.startswith("audio_embeddings.") and k.endswith(".weight")]) + 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: + raise ValueError( + "final_proj rows do not match stacked audio layout: " + f"rows={sd['final_proj.weight'].shape[0]} expected={expected_logits}" + ) + 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 ) - n_codebooks //= frame_stacking inf = cfg.get("inference_parameters", {}) @@ -105,9 +151,11 @@ 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, + "stacked_audio_codebooks": n_stacked_codebooks, "audio_codebook_size": codebook_size, "audio_vocab_size": audio_vocab, "frame_stacking_factor": frame_stacking, @@ -127,6 +175,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)) @@ -135,6 +184,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 +358,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/conversion/tts_tokenizer_profiles.py b/conversion/tts_tokenizer_profiles.py new file mode 100644 index 0000000..5176514 --- /dev/null +++ b/conversion/tts_tokenizer_profiles.py @@ -0,0 +1,164 @@ +#!/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/configuration.md b/docs/tts/configuration.md index 17839ad..38482cc 100644 --- a/docs/tts/configuration.md +++ b/docs/tts/configuration.md @@ -41,17 +41,24 @@ when present. The older split layout (`classify/tokenize_and_classify.far`, The optional `riva_server` supports `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` -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`. Mandarin uses bundled Jieba and +`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 compiled-in TTS languages and the -dotted voice names accepted by synthesis requests. +`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`. @@ -137,7 +144,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 text language code | | `tts.voice-name` | - | - | default voice name or speaker index | @@ -160,7 +167,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 | @@ -177,7 +184,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/docs/tts/models.md b/docs/tts/models.md index 5ce2050..e086e90 100644 --- a/docs/tts/models.md +++ b/docs/tts/models.md @@ -16,12 +16,50 @@ options are omitted. Hugging Face: [nvidia/magpie_tts_multilingual_357m](https://huggingface.co/nvidia/magpie_tts_multilingual_357m) +```bash +# Download the v2602 GGUF and the original archive containing its tokenizer. +hf download nvidia/magpie_tts_multilingual_357m \ + --include magpie_tts_multilingual_357m.v2602.f16.gguf \ + --include magpie_tts_multilingual_357m.nemo \ + --local-dir models/magpie-tts + +# Extract the tokenizer assets loaded by the runtime. +mkdir -p models/magpie-tts/extracted +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. The built-in pull extracts only the required, pinned tokenizer members and verifies each one. For a custom Magpie checkpoint, extract its `.nemo` archive and pass that directory as `--tokenizer-dir` or `--tts.tokenizer-model-dir`. 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/ggml-patches/0007-magpietts-nanocodec.patch b/ggml-patches/0007-magpietts-nanocodec.patch index 13dc370..23f8c8b 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,55 @@ 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..a347447c 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, +@@ -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); +- 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 +801,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 +595,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 +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)) { @@ -551,8 +604,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 +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)) { @@ -569,7 +626,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-fused-attention-extensions.patch b/ggml-patches/0014-cuda-fused-attention-extensions.patch new file mode 100644 index 0000000..4d9de3f --- /dev/null +++ b/ggml-patches/0014-cuda-fused-attention-extensions.patch @@ -0,0 +1,2626 @@ +diff --git a/include/ggml.h b/include/ggml.h +index fb823571..d403756e 100644 +--- a/include/ggml.h ++++ b/include/ggml.h +@@ -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) +- // mask [kv_len] or [kv_len, batch] additive (0 / -inf) key mask, +- // or NULL shared or per-stream columns ++ // mask [kv_len], [kv_len, q_len], additive (0 / -inf) mask, ++ // or [kv_len, batch], or NULL offline per-query or streaming ++ // per-stream columns + // 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 +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. ++ // 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( ++ struct ggml_context * ctx, ++ struct ggml_tensor * q, ++ struct ggml_tensor * k, ++ struct ggml_tensor * v, ++ struct ggml_tensor * p, ++ struct ggml_tensor * bias_u, ++ struct ggml_tensor * bias_v, ++ struct ggml_tensor * mask, ++ struct ggml_tensor * kv_cache, ++ struct ggml_tensor * slot_ids, ++ struct ggml_tensor * cache_state, ++ int64_t cache_len, ++ float scale, ++ bool merge_heads); ++ ++ // 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, ++ 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 * cache_state, ++ 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-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 + } + + 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 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 ++// 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 attention_warp_all_sum(float value) { ++#pragma unroll ++ for (int offset = 16; offset > 0; offset >>= 1) { ++ value += __shfl_xor_sync(0xffffffff, value, offset); ++ } ++ return value; ++} ++ ++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 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) { ++ if (j < cache_len) { ++ int physical_j = ring_head + j; ++ if (physical_j >= cache_len) { ++ physical_j -= cache_len; ++ } ++ return attention_load4(cache_head + (size_t) physical_j * cache_sj + d4); ++ } ++ return attention_load4(chunk_head + (size_t) (j - cache_len) * chunk_sj + d4); ++ } ++ return attention_load4(chunk_head + (size_t) j * chunk_sj + d4); ++} ++ ++template ++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) { ++ if (j < cache_len) { ++ int physical_j = ring_head + j; ++ if (physical_j >= cache_len) { ++ physical_j -= cache_len; ++ } ++ return cache_head[(size_t) physical_j * cache_sj + d]; ++ } ++ return (float) chunk_head[(size_t) (j - cache_len) * chunk_sj + d]; ++ } ++ 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_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, ++ int cache_len, int chunk_len, int n_feat, int d_k, ++ long cache_ss, long cache_sp, ++ long k_sj, long k_sh, long k_sb, long v_sj, long v_sh, long v_sb) { ++ const int feature = (int) blockIdx.x * blockDim.x + threadIdx.x; ++ const int b = blockIdx.y; ++ const int plane = blockIdx.z; ++ if (feature >= n_feat) { ++ return; ++ } ++ ++ const int slot = slot_ids[b]; ++ const int ring_head = ring_heads[b]; ++ float * dst = arena + (size_t) plane * cache_sp + (size_t) slot * cache_ss + feature; ++ const int h = feature / d_k; ++ const int d = feature - h * d_k; ++ const T * src = (plane == 0 ? K : V) + ++ (size_t) b * (plane == 0 ? k_sb : v_sb) + ++ (size_t) h * (plane == 0 ? k_sh : v_sh) + d; ++ const long src_sj = plane == 0 ? k_sj : v_sj; ++ const int append = min(cache_len, chunk_len); ++ for (int j = 0; j < append; ++j) { ++ const int src_j = chunk_len - append + j; ++ int physical_j = ring_head + j; ++ if (physical_j >= cache_len) { ++ physical_j -= cache_len; ++ } ++ dst[(size_t) physical_j * n_feat] = (float) src[(size_t) src_j * src_sj]; ++ } ++} ++ ++template ++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 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[]; ++ 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 ++ ? Kcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128 ++ : nullptr; ++ 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; ++ ++ 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 = 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] * 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]; ++} ++ ++// 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, ++ 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 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[]; ++ 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 ++ ? Kcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128 ++ : nullptr; ++ 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; ++ ++ 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 = 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 = ++ attention_load_kv( ++ Vh, Vch, j, cache_len, ring_head, v_sj, cache_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]; ++} ++ ++// 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, ++ 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, ++ 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 cache_sj, long cache_ss, ++ long p_sr, long p_sh, ++ long o_si, long o_sh, long o_sb, long m_sb) { ++ constexpr int warps = 8; ++ constexpr int keys_per_warp = 9; ++ constexpr int cache_len = 70; ++ ++ extern __shared__ float sh[]; ++ float * den = sh; // [warp, query] ++ float * part = den + warps * 2; // [warp, query, d] ++ ++ const int tid = threadIdx.x; ++ const int lane = tid & 31; ++ const int warp = tid >> 5; ++ const int h = blockIdx.x; ++ const int b = blockIdx.y; ++ const int d4 = lane * 4; ++ const int j0 = warp * keys_per_warp; ++ ++ 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 = slot_ids[b]; ++ const int ring_head = ring_heads[b]; ++ const float * Kch = ++ Kcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128; ++ const float * Vch = ++ Vcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128; ++ const T * Ph = Ppos + (size_t) h * p_sh; ++ const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; ++ ++ 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( ++ q14.x + bu4.x, q14.y + bu4.y, q14.z + bu4.z, q14.w + bu4.w); ++ const float4 qp0 = make_float4( ++ q04.x + bv4.x, q04.y + bv4.y, q04.z + bv4.z, q04.w + bv4.w); ++ const float4 qp1 = make_float4( ++ q14.x + bv4.x, q14.y + bv4.y, q14.z + bv4.z, q14.w + bv4.w); ++ ++ float scores0[keys_per_warp]; ++ float scores1[keys_per_warp]; ++ float local_max0 = -INFINITY; ++ float local_max1 = -INFINITY; ++ 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 = 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 = attention_load4(Kch + (size_t) physical_j * cache_sj + d4); ++ } else { ++ k4 = attention_load4(Kh + (size_t) (j - cache_len) * k_sj + d4); ++ } ++ float score0 = ++ k4.x * qc0.x + p04.x * qp0.x + ++ k4.y * qc0.y + p04.y * qp0.y + ++ k4.z * qc0.z + p04.z * qp0.z + ++ k4.w * qc0.w + p04.w * qp0.w; ++ float score1 = ++ k4.x * qc1.x + p14.x * qp1.x + ++ 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 = 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; ++ scores0[c] = score0; ++ scores1[c] = score1; ++ local_max0 = fmaxf(local_max0, score0); ++ local_max1 = fmaxf(local_max1, score1); ++ p14 = p04; ++ } ++ if (lane == 0) { ++ den[warp * 2] = local_max0; ++ den[warp * 2 + 1] = local_max1; ++ } ++ __syncthreads(); ++ ++ float maximum0 = den[0]; ++ float maximum1 = den[1]; ++#pragma unroll ++ for (int w = 1; w < warps; ++w) { ++ maximum0 = fmaxf(maximum0, den[w * 2]); ++ maximum1 = fmaxf(maximum1, den[w * 2 + 1]); ++ } ++ // Every warp must finish reading the maxima before lane zero reuses den. ++ __syncthreads(); ++ float local_den0 = 0.0f; ++ float local_den1 = 0.0f; ++#pragma unroll ++ for (int c = 0; c < keys_per_warp; ++c) { ++ const float weight0 = __expf(scores0[c] - maximum0); ++ const float weight1 = __expf(scores1[c] - maximum1); ++ scores0[c] = weight0; ++ scores1[c] = weight1; ++ local_den0 += weight0; ++ local_den1 += weight1; ++ } ++ if (lane == 0) { ++ den[warp * 2] = local_den0; ++ den[warp * 2 + 1] = local_den1; ++ } ++ ++ float4 acc0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); ++ float4 acc1 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); ++#pragma unroll ++ for (int c = 0; c < keys_per_warp; ++c) { ++ const int j = j0 + c; ++ float4 v4; ++ if (c < 7 || j < cache_len) { ++ int physical_j = ring_head + j; ++ if (physical_j >= cache_len) { ++ physical_j -= cache_len; ++ } ++ v4 = attention_load4(Vch + (size_t) physical_j * cache_sj + d4); ++ } else { ++ 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); ++ acc0.z = fmaf(scores0[c], v4.z, acc0.z); ++ acc0.w = fmaf(scores0[c], v4.w, acc0.w); ++ acc1.x = fmaf(scores1[c], v4.x, acc1.x); ++ acc1.y = fmaf(scores1[c], v4.y, acc1.y); ++ acc1.z = fmaf(scores1[c], v4.z, acc1.z); ++ acc1.w = fmaf(scores1[c], v4.w, acc1.w); ++ } ++ *reinterpret_cast( ++ part + ((warp * 2) * RELPOS_ATTN_DK_128) + d4) = acc0; ++ *reinterpret_cast( ++ part + ((warp * 2 + 1) * RELPOS_ATTN_DK_128) + d4) = acc1; ++ __syncthreads(); ++ ++ float total_den0 = den[0]; ++ float total_den1 = den[1]; ++#pragma unroll ++ for (int w = 1; w < warps; ++w) { ++ total_den0 += den[w * 2]; ++ total_den1 += den[w * 2 + 1]; ++ } ++ const int i = tid / RELPOS_ATTN_DK_128; ++ const int d = tid % RELPOS_ATTN_DK_128; ++ float context = part[i * RELPOS_ATTN_DK_128 + d]; ++#pragma unroll ++ for (int w = 1; w < warps; ++w) { ++ context += part[((w * 2 + i) * RELPOS_ATTN_DK_128) + d]; ++ } ++ const float denominator = i == 0 ? total_den0 : total_den1; ++ ctx[(size_t) b * o_sb + (size_t) i * o_si + (size_t) h * o_sh + d] = ++ context / denominator; ++} ++ ++// 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, ++ 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, ++ 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 cache_sj, long cache_ss, ++ long p_sr, long p_sh, ++ long o_si, long o_sh, long o_sb, long m_sb) { ++ constexpr int warps = 6; ++ constexpr int queries = 4; ++ constexpr int keys_per_warp = 10; ++ constexpr int cache_len = 56; ++ ++ extern __shared__ float sh[]; ++ float * den = sh; // [warp, query] ++ float * part = den + warps * queries; // [warp, query, d] ++ ++ const int tid = threadIdx.x; ++ const int lane = tid & 31; ++ const int warp = tid >> 5; ++ const int h = blockIdx.x; ++ const int b = blockIdx.y; ++ const int d4 = lane * 4; ++ const int j0 = warp * keys_per_warp; ++ ++ const float * Qh0 = Q + (size_t) b * q_sb + (size_t) h * q_sh; ++ const float * Qh1 = Qh0 + q_sq; ++ const float * Qh2 = Qh1 + q_sq; ++ const float * Qh3 = Qh2 + 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 = slot_ids[b]; ++ const int ring_head = ring_heads[b]; ++ const float * Kch = ++ Kcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128; ++ const float * Vch = ++ Vcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128; ++ const T * Ph = Ppos + (size_t) h * p_sh; ++ const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; ++ ++ 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); ++ const float4 qc1 = RELPOS_ADD_BIAS(q14, bu4); ++ const float4 qc2 = RELPOS_ADD_BIAS(q24, bu4); ++ const float4 qc3 = RELPOS_ADD_BIAS(q34, bu4); ++ const float4 qp0 = RELPOS_ADD_BIAS(q04, bv4); ++ const float4 qp1 = RELPOS_ADD_BIAS(q14, bv4); ++ const float4 qp2 = RELPOS_ADD_BIAS(q24, bv4); ++ const float4 qp3 = RELPOS_ADD_BIAS(q34, bv4); ++#undef RELPOS_ADD_BIAS ++ ++ float scores0[keys_per_warp]; ++ float scores1[keys_per_warp]; ++ float scores2[keys_per_warp]; ++ float scores3[keys_per_warp]; ++ float local_max0 = -INFINITY; ++ float local_max1 = -INFINITY; ++ float local_max2 = -INFINITY; ++ float local_max3 = -INFINITY; ++#pragma unroll ++ for (int c = 0; c < keys_per_warp; ++c) { ++ const int j = j0 + c; ++ const float4 k4 = attention_load_kv4( ++ Kh, Kch, j, cache_len, ring_head, k_sj, cache_sj, 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 = 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; ++ score1 = score1 * scale + additive_mask; ++ score2 = score2 * scale + additive_mask; ++ score3 = score3 * scale + additive_mask; ++ scores0[c] = score0; ++ scores1[c] = score1; ++ scores2[c] = score2; ++ scores3[c] = score3; ++ local_max0 = fmaxf(local_max0, score0); ++ local_max1 = fmaxf(local_max1, score1); ++ local_max2 = fmaxf(local_max2, score2); ++ local_max3 = fmaxf(local_max3, score3); ++ } ++ if (lane == 0) { ++ den[warp * queries + 0] = local_max0; ++ den[warp * queries + 1] = local_max1; ++ den[warp * queries + 2] = local_max2; ++ den[warp * queries + 3] = local_max3; ++ } ++ __syncthreads(); ++ ++ float maximum0 = den[0]; ++ float maximum1 = den[1]; ++ float maximum2 = den[2]; ++ float maximum3 = den[3]; ++#pragma unroll ++ for (int w = 1; w < warps; ++w) { ++ maximum0 = fmaxf(maximum0, den[w * queries + 0]); ++ maximum1 = fmaxf(maximum1, den[w * queries + 1]); ++ maximum2 = fmaxf(maximum2, den[w * queries + 2]); ++ maximum3 = fmaxf(maximum3, den[w * queries + 3]); ++ } ++ __syncthreads(); ++ ++ float local_den0 = 0.0f; ++ float local_den1 = 0.0f; ++ float local_den2 = 0.0f; ++ float local_den3 = 0.0f; ++#pragma unroll ++ for (int c = 0; c < keys_per_warp; ++c) { ++ scores0[c] = __expf(scores0[c] - maximum0); ++ scores1[c] = __expf(scores1[c] - maximum1); ++ scores2[c] = __expf(scores2[c] - maximum2); ++ scores3[c] = __expf(scores3[c] - maximum3); ++ local_den0 += scores0[c]; ++ local_den1 += scores1[c]; ++ local_den2 += scores2[c]; ++ local_den3 += scores3[c]; ++ } ++ if (lane == 0) { ++ den[warp * queries + 0] = local_den0; ++ den[warp * queries + 1] = local_den1; ++ den[warp * queries + 2] = local_den2; ++ den[warp * queries + 3] = local_den3; ++ } ++ ++ float4 acc0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); ++ float4 acc1 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); ++ float4 acc2 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); ++ float4 acc3 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); ++#pragma unroll ++ for (int c = 0; c < keys_per_warp; ++c) { ++ const int j = j0 + c; ++ 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 { \ ++ (acc).x = fmaf((weight), (value).x, (acc).x); \ ++ (acc).y = fmaf((weight), (value).y, (acc).y); \ ++ (acc).z = fmaf((weight), (value).z, (acc).z); \ ++ (acc).w = fmaf((weight), (value).w, (acc).w); \ ++ } while (0) ++ RELPOS_ACCUM(acc0, scores0[c], v4); ++ RELPOS_ACCUM(acc1, scores1[c], v4); ++ RELPOS_ACCUM(acc2, scores2[c], v4); ++ RELPOS_ACCUM(acc3, scores3[c], v4); ++#undef RELPOS_ACCUM ++ } ++ *reinterpret_cast( ++ part + ((warp * queries + 0) * RELPOS_ATTN_DK_128) + d4) = acc0; ++ *reinterpret_cast( ++ part + ((warp * queries + 1) * RELPOS_ATTN_DK_128) + d4) = acc1; ++ *reinterpret_cast( ++ part + ((warp * queries + 2) * RELPOS_ATTN_DK_128) + d4) = acc2; ++ *reinterpret_cast( ++ part + ((warp * queries + 3) * RELPOS_ATTN_DK_128) + d4) = acc3; ++ __syncthreads(); ++ ++ for (int index = tid; index < queries * RELPOS_ATTN_DK_128; index += blockDim.x) { ++ const int query = index / RELPOS_ATTN_DK_128; ++ const int d = index % RELPOS_ATTN_DK_128; ++ float denominator = den[query]; ++ float context = part[query * RELPOS_ATTN_DK_128 + d]; ++#pragma unroll ++ for (int w = 1; w < warps; ++w) { ++ denominator += den[w * queries + query]; ++ context += part[(w * queries + query) * RELPOS_ATTN_DK_128 + d]; ++ } ++ ctx[(size_t) b * o_sb + (size_t) query * o_si + (size_t) h * o_sh + d] = ++ context / denominator; ++ } ++} ++ ++// 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, ++ 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 q_len, int kv_len, 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 cache_sj, long cache_ss, ++ long p_sr, long p_sh, ++ long o_si, long o_sh, long o_sb, long m_sb) { ++ constexpr int warps = 8; ++ constexpr int queries = 2; ++ constexpr int keys_per_warp = 11; ++ ++ extern __shared__ float sh[]; ++ float * den = sh; ++ float * part = den + warps * queries; ++ ++ const int tid = threadIdx.x; ++ const int lane = tid & 31; ++ const int warp = tid >> 5; ++ const int h = blockIdx.x; ++ const int b = blockIdx.y; ++ const int query0 = blockIdx.z * queries; ++ const bool has_query1 = query0 + 1 < q_len; ++ const int d4 = lane * 4; ++ const int j0 = warp * keys_per_warp; ++ ++ const float * Qh0 = ++ Q + (size_t) b * q_sb + (size_t) query0 * q_sq + (size_t) h * q_sh; ++ const float * Qh1 = has_query1 ? Qh0 + q_sq : Qh0; ++ 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 = slot_ids[b]; ++ const int ring_head = ring_heads[b]; ++ const float * Kch = ++ Kcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128; ++ const float * Vch = ++ Vcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128; ++ const T * Ph = Ppos + (size_t) h * p_sh; ++ const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; ++ ++ 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); ++ const float4 qc1 = RELPOS_ADD_BIAS(q14, bu4); ++ const float4 qp0 = RELPOS_ADD_BIAS(q04, bv4); ++ const float4 qp1 = RELPOS_ADD_BIAS(q14, bv4); ++#undef RELPOS_ADD_BIAS ++ ++ float scores0[keys_per_warp]; ++ float scores1[keys_per_warp]; ++ float local_max0 = -INFINITY; ++ float local_max1 = -INFINITY; ++#pragma unroll ++ for (int c = 0; c < keys_per_warp; ++c) { ++ const int j = j0 + c; ++ float score0 = -INFINITY; ++ float score1 = -INFINITY; ++ if (j < kv_len) { ++ 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 = 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 = attention_warp_all_sum(RELPOS_SCORE(k4, qc0, p04, qp0)); ++ if (has_query1) { ++ 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; ++ score0 = score0 * scale + additive_mask; ++ if (has_query1) ++ score1 = score1 * scale + additive_mask; ++ } ++ scores0[c] = score0; ++ scores1[c] = score1; ++ local_max0 = fmaxf(local_max0, score0); ++ local_max1 = fmaxf(local_max1, score1); ++ } ++ if (lane == 0) { ++ den[warp * queries] = local_max0; ++ den[warp * queries + 1] = local_max1; ++ } ++ __syncthreads(); ++ ++ float maximum0 = den[0]; ++ float maximum1 = den[1]; ++#pragma unroll ++ for (int w = 1; w < warps; ++w) { ++ maximum0 = fmaxf(maximum0, den[w * queries]); ++ maximum1 = fmaxf(maximum1, den[w * queries + 1]); ++ } ++ __syncthreads(); ++ ++ float local_den0 = 0.0f; ++ float local_den1 = 0.0f; ++#pragma unroll ++ for (int c = 0; c < keys_per_warp; ++c) { ++ const int j = j0 + c; ++ if (j < kv_len) { ++ scores0[c] = __expf(scores0[c] - maximum0); ++ scores1[c] = has_query1 ? __expf(scores1[c] - maximum1) : 0.0f; ++ local_den0 += scores0[c]; ++ local_den1 += scores1[c]; ++ } else { ++ scores0[c] = 0.0f; ++ scores1[c] = 0.0f; ++ } ++ } ++ if (lane == 0) { ++ den[warp * queries] = local_den0; ++ den[warp * queries + 1] = local_den1; ++ } ++ ++ float4 acc0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); ++ float4 acc1 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); ++#pragma unroll ++ for (int c = 0; c < keys_per_warp; ++c) { ++ const int j = j0 + c; ++ if (j >= kv_len) ++ continue; ++ 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 { \ ++ (acc).x = fmaf((weight), (value).x, (acc).x); \ ++ (acc).y = fmaf((weight), (value).y, (acc).y); \ ++ (acc).z = fmaf((weight), (value).z, (acc).z); \ ++ (acc).w = fmaf((weight), (value).w, (acc).w); \ ++ } while (0) ++ RELPOS_ACCUM(acc0, scores0[c], v4); ++ RELPOS_ACCUM(acc1, scores1[c], v4); ++#undef RELPOS_ACCUM ++ } ++ *reinterpret_cast( ++ part + ((warp * queries) * RELPOS_ATTN_DK_128) + d4) = acc0; ++ *reinterpret_cast( ++ part + ((warp * queries + 1) * RELPOS_ATTN_DK_128) + d4) = acc1; ++ __syncthreads(); ++ ++ const int query = tid / RELPOS_ATTN_DK_128; ++ if (query == 0 || has_query1) { ++ const int d = tid % RELPOS_ATTN_DK_128; ++ float denominator = den[query]; ++ float context = part[query * RELPOS_ATTN_DK_128 + d]; ++#pragma unroll ++ for (int w = 1; w < warps; ++w) { ++ denominator += den[w * queries + query]; ++ context += part[(w * queries + query) * RELPOS_ATTN_DK_128 + d]; ++ } ++ ctx[(size_t) b * o_sb + (size_t) (query0 + query) * o_si + ++ (size_t) h * o_sh + d] = context / denominator; ++ } ++} ++ ++template ++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 ++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, ++ long cache_sj, long cache_ss, ++ 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] ++ 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; ++ 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 = 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] + (relative ? bu[h * dk + d] : 0.0f); ++ Qv[d] = Qhi[d] + (relative ? bv[h * dk + d] : 0.0f); ++ __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 float4 k4 = ++ 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 = 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) { ++ 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. ++ 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 = attention_load_kv4( ++ Kh, Kch, j, cache_len, ring_head, k_sj, cache_sj, 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 = 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 = key_begin + d; j < kv; j += dk) { ++ 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); ++ } ++ } ++ __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) { ++ 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; ++ } ++ 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] * 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; ++} ++ ++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 * 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); ++ GGML_ASSERT(!relative || k->type == p->type); ++ 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 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_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); ++ 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(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(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(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); ++ if (cached) { ++ // Streaming: shared across the batch or one key-mask column per ++ // stream. ++ GGML_ASSERT(mask->ne[1] == 1 || mask->ne[1] == batch); ++ GGML_ASSERT(mask->ne[2] == 1 && mask->ne[3] == 1); ++ } else { ++ // Offline: shared/per-batch key vector or one column per query. ++ 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 && 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)); ++ ++ // 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); ++ ++ 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 *) cache_state->data : nullptr; ++ const int32_t * active_lengths = ++ cached && cache_state->ne[1] == 2 ++ ? active_ring_heads + cache_state->ne[0] ++ : nullptr; ++ ++ auto update_cache = [&]() { ++ if (!cached) { ++ return; ++ } ++ const dim3 update_grid((d_k * n_head + 255) / 256, batch, 2); ++ if (k->type == GGML_TYPE_F16) { ++ 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_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); ++ } ++ }; ++ ++ 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 && 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; ++ 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; ++ 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; ++ if (use_register_q4) { ++ constexpr int register_warps = 6; ++ constexpr int register_queries = 4; ++ const size_t register_shmem = ++ ((size_t) register_warps * register_queries + ++ (size_t) register_warps * register_queries * RELPOS_ATTN_DK_128) * sizeof(float); ++ GGML_ASSERT(register_shmem <= max_shmem); ++ const dim3 register_grid(n_head, batch, 1); ++ if (k->type == GGML_TYPE_F16) { ++ fused_relpos_attn_q4_register_kernel<<< ++ register_grid, 192, register_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, cache_k, cache_v, active_slots, active_ring_heads, 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_q4_register_kernel<<< ++ register_grid, 192, register_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, cache_k, cache_v, active_slots, active_ring_heads, 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); ++ } ++ update_cache(); ++ return; ++ } ++ if (use_register_q2) { ++ constexpr int register_warps = 8; ++ const size_t register_shmem = ++ ((size_t) register_warps * 2 + ++ (size_t) register_warps * 2 * RELPOS_ATTN_DK_128) * sizeof(float); ++ GGML_ASSERT(register_shmem <= max_shmem); ++ const dim3 register_grid(n_head, batch, 1); ++ if (k->type == GGML_TYPE_F16) { ++ fused_relpos_attn_q2_register_kernel<<< ++ register_grid, 256, register_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, cache_k, cache_v, active_slots, active_ring_heads, 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_q2_register_kernel<<< ++ register_grid, 256, register_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, cache_k, cache_v, active_slots, active_ring_heads, 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); ++ } ++ update_cache(); ++ return; ++ } ++ if (use_register_common) { ++ constexpr int register_warps = 8; ++ constexpr int register_queries = 2; ++ const size_t register_shmem = ++ ((size_t) register_warps * register_queries + ++ (size_t) register_warps * register_queries * RELPOS_ATTN_DK_128) * sizeof(float); ++ GGML_ASSERT(register_shmem <= max_shmem); ++ const dim3 register_grid(n_head, batch, (q_len + register_queries - 1) / register_queries); ++ if (k->type == GGML_TYPE_F16) { ++ fused_relpos_attn_common_q2_register_kernel<<< ++ register_grid, 256, register_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, cache_k, cache_v, active_slots, active_ring_heads, ++ q_len, 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_common_q2_register_kernel<<< ++ register_grid, 256, register_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, cache_k, cache_v, active_slots, active_ring_heads, ++ q_len, 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); ++ } ++ 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); ++ 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 ++ ? relpos_attn_warp_128_max_blocks_per_sm(device, warp_shmem) ++ : relpos_attn_warp_128_max_blocks_per_sm(device, warp_shmem); ++ } else { ++ max_single_blocks_per_sm = cached ++ ? 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); ++ ++ 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<<< ++ 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<<< ++ 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, ++ (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); ++ } ++ }; ++ if (k->type == GGML_TYPE_F16) { ++ if (cached) { ++ launch_tuned(half{}, std::true_type{}); ++ } else { ++ launch_tuned(half{}, std::false_type{}); ++ } ++ } else { ++ if (cached) { ++ launch_tuned(float{}, std::true_type{}); ++ } else { ++ launch_tuned(float{}, std::false_type{}); ++ } ++ } ++ update_cache(); ++ 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_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, ++ 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); ++ }; ++ if (k->type == GGML_TYPE_F16) { ++ if (cached) { ++ launch_generic(half{}, std::true_type{}); ++ } else { ++ launch_generic(half{}, std::false_type{}); ++ } ++ } 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, +- 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); +- } +-} +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..4f5aeb09 100644 +--- a/src/ggml.c ++++ b/src/ggml.c +@@ -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_attention + +-struct ggml_tensor * ggml_fused_relpos_attn( ++static struct ggml_tensor * ggml_fused_attention_impl( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, +@@ -5386,6 +5386,10 @@ struct ggml_tensor * ggml_fused_relpos_attn( + struct ggml_tensor * bias_u, + struct ggml_tensor * bias_v, + struct ggml_tensor * mask, ++ struct ggml_tensor * kv_cache, ++ struct ggml_tensor * slot_ids, ++ struct ggml_tensor * cache_state, ++ int64_t cache_len, + float scale, + bool merge_heads) { + // Q/K/V/P may be arbitrary-strided views; the CUDA op derives addressing +@@ -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); + // 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(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)); ++ 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)); ++ GGML_ASSERT(kv_cache->ne[0] == 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] == q->ne[3]); ++ 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); ++ } + if (mask) { + GGML_ASSERT(ggml_is_contiguous(mask)); + GGML_ASSERT(mask->ne[0] == kv_len); +- // One shared key mask, or one column per batch item (the cache-aware +- // streaming layout, where each stream's history has its own validity). +- GGML_ASSERT(mask->ne[1] == 1 || mask->ne[1] == q->ne[3]); +- GGML_ASSERT(mask->ne[2] == 1 && mask->ne[3] == 1); ++ if (kv_cache) { ++ // Cache-aware streaming: one shared key mask, or one column per ++ // batch item whose history has its own validity. ++ GGML_ASSERT(mask->ne[1] == 1 || mask->ne[1] == q->ne[3]); ++ GGML_ASSERT(mask->ne[2] == 1 && mask->ne[3] == 1); ++ } else { ++ // Offline: a shared/per-batch key vector or a [key,query] ++ // local-attention mask, broadcast over heads. ++ GGML_ASSERT(mask->ne[1] == 1 || mask->ne[1] == q_len); ++ GGML_ASSERT(mask->ne[2] == 1 && (mask->ne[3] == 1 || mask->ne[3] == q->ne[3])); ++ } + } + + // Output mirrors q logically: [d_k, q_len, n_head, batch]. With +@@ -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_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] = cache_state; + + return result; + } + ++struct ggml_tensor * ggml_fused_relpos_attn( ++ struct ggml_context * ctx, ++ struct ggml_tensor * q, ++ struct ggml_tensor * k, ++ struct ggml_tensor * v, ++ struct ggml_tensor * p, ++ struct ggml_tensor * bias_u, ++ struct ggml_tensor * bias_v, ++ struct ggml_tensor * mask, ++ float scale, ++ bool merge_heads) { ++ return ggml_fused_attention_impl( ++ ctx, q, k, v, p, bias_u, bias_v, mask, NULL, NULL, NULL, 0, scale, merge_heads); ++} ++ ++struct ggml_tensor * ggml_fused_relpos_attn_cached( ++ struct ggml_context * ctx, ++ struct ggml_tensor * q, ++ struct ggml_tensor * k, ++ struct ggml_tensor * v, ++ struct ggml_tensor * p, ++ struct ggml_tensor * bias_u, ++ struct ggml_tensor * bias_v, ++ struct ggml_tensor * mask, ++ struct ggml_tensor * kv_cache, ++ struct ggml_tensor * slot_ids, ++ struct ggml_tensor * cache_state, ++ int64_t cache_len, ++ float scale, ++ bool merge_heads) { ++ 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); ++} ++ ++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 * cache_state, ++ int64_t cache_len, ++ float scale, ++ bool merge_heads) { ++ return ggml_fused_attention_impl( ++ ctx, q, k, v, NULL, NULL, NULL, mask, kv_cache, slot_ids, cache_state, cache_len, scale, ++ merge_heads); ++} ++ + void ggml_flash_attn_ext_set_prec( + struct ggml_tensor * a, + enum ggml_prec prec) { diff --git a/ggml-patches/0014-cuda-relpos-extensions.patch b/ggml-patches/0014-cuda-relpos-extensions.patch deleted file mode 100644 index cbf2af9..0000000 --- a/ggml-patches/0014-cuda-relpos-extensions.patch +++ /dev/null @@ -1,1481 +0,0 @@ -diff --git a/include/ggml.h b/include/ggml.h -index fb823571..5cfa65e0 100644 ---- a/include/ggml.h -+++ b/include/ggml.h -@@ -2436,8 +2436,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) -- // mask [kv_len] or [kv_len, batch] additive (0 / -inf) key mask, -- // or NULL shared or per-stream columns -+ // mask [kv_len], [kv_len, q_len], additive (0 / -inf) mask, -+ // or [kv_len, batch], or NULL offline per-query or streaming -+ // per-stream columns - // 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" { - 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. -+ GGML_API struct ggml_tensor * ggml_fused_relpos_attn_cached( -+ struct ggml_context * ctx, -+ struct ggml_tensor * q, -+ struct ggml_tensor * k, -+ struct ggml_tensor * v, -+ struct ggml_tensor * p, -+ struct ggml_tensor * bias_u, -+ struct ggml_tensor * bias_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 ---- 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" - - #include -+#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. -+// 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 -+// 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) { -+#pragma unroll -+ for (int offset = 16; offset > 0; offset >>= 1) { -+ value += __shfl_xor_sync(0xffffffff, value, offset); -+ } -+ 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])); - } - -+template -+static __device__ __forceinline__ float4 relpos_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) { -+ if (j < cache_len) { -+ int physical_j = ring_head + j; -+ if (physical_j >= cache_len) { -+ physical_j -= cache_len; -+ } -+ return relpos_load4(cache_head + (size_t) physical_j * cache_sj + d4); -+ } -+ return relpos_load4(chunk_head + (size_t) (j - cache_len) * chunk_sj + d4); -+ } -+ return relpos_load4(chunk_head + (size_t) j * chunk_sj + d4); -+} -+ -+template -+static __device__ __forceinline__ float relpos_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) { -+ if (j < cache_len) { -+ int physical_j = ring_head + j; -+ if (physical_j >= cache_len) { -+ physical_j -= cache_len; -+ } -+ return cache_head[(size_t) physical_j * cache_sj + d]; -+ } -+ return (float) chunk_head[(size_t) (j - cache_len) * chunk_sj + d]; -+ } -+ return (float) chunk_head[(size_t) j * chunk_sj + d]; -+} -+ -+// 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( -+ float * __restrict__ arena, const T * __restrict__ K, -+ const T * __restrict__ V, const int32_t * __restrict__ slot_ids, -+ const int32_t * __restrict__ ring_heads, -+ int cache_len, int chunk_len, int n_feat, int d_k, -+ long cache_ss, long cache_sp, -+ long k_sj, long k_sh, long k_sb, long v_sj, long v_sh, long v_sb) { -+ const int feature = (int) blockIdx.x * blockDim.x + threadIdx.x; -+ const int b = blockIdx.y; -+ const int plane = blockIdx.z; -+ if (feature >= n_feat) { -+ return; -+ } -+ -+ const int slot = slot_ids[b]; -+ const int ring_head = ring_heads[b]; -+ float * dst = arena + (size_t) plane * cache_sp + (size_t) slot * cache_ss + feature; -+ const int h = feature / d_k; -+ const int d = feature - h * d_k; -+ const T * src = (plane == 0 ? K : V) + -+ (size_t) b * (plane == 0 ? k_sb : v_sb) + -+ (size_t) h * (plane == 0 ? k_sh : v_sh) + d; -+ const long src_sj = plane == 0 ? k_sj : v_sj; -+ const int append = min(cache_len, chunk_len); -+ for (int j = 0; j < append; ++j) { -+ const int src_j = chunk_len - append + j; -+ int physical_j = ring_head + j; -+ if (physical_j >= cache_len) { -+ physical_j -= cache_len; -+ } -+ dst[(size_t) physical_j * n_feat] = (float) src[(size_t) src_j * src_sj]; -+ } -+} -+ - // 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, -+ 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 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; -+ const int slot = Cached ? slot_ids[b] : 0; -+ const int ring_head = Cached ? ring_heads[b] : 0; -+ const float * Kch = Cached -+ ? Kcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128 -+ : nullptr; -+ 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( -+ 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]; -+ context += -+ 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] = 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 -+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, -+ 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 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; -+ const int slot = Cached ? slot_ids[b] : 0; -+ const int ring_head = Cached ? ring_heads[b] : 0; -+ const float * Kch = Cached -+ ? Kcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128 -+ : nullptr; -+ 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( -+ 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 float value = -+ relpos_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 -+__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, -+ 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, -+ 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 cache_sj, long cache_ss, -+ long p_sr, long p_sh, -+ long o_si, long o_sh, long o_sb, long m_sb) { -+ constexpr int warps = 8; -+ constexpr int keys_per_warp = 9; -+ constexpr int cache_len = 70; -+ -+ extern __shared__ float sh[]; -+ float * den = sh; // [warp, query] -+ float * part = den + warps * 2; // [warp, query, d] -+ -+ const int tid = threadIdx.x; -+ const int lane = tid & 31; -+ const int warp = tid >> 5; -+ const int h = blockIdx.x; -+ const int b = blockIdx.y; -+ const int d4 = lane * 4; -+ const int j0 = warp * keys_per_warp; -+ -+ 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 = slot_ids[b]; -+ const int ring_head = ring_heads[b]; -+ const float * Kch = -+ Kcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128; -+ const float * Vch = -+ Vcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128; -+ 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 qc0 = make_float4( -+ q04.x + bu4.x, q04.y + bu4.y, q04.z + bu4.z, q04.w + bu4.w); -+ const float4 qc1 = make_float4( -+ q14.x + bu4.x, q14.y + bu4.y, q14.z + bu4.z, q14.w + bu4.w); -+ const float4 qp0 = make_float4( -+ q04.x + bv4.x, q04.y + bv4.y, q04.z + bv4.z, q04.w + bv4.w); -+ const float4 qp1 = make_float4( -+ q14.x + bv4.x, q14.y + bv4.y, q14.z + bv4.z, q14.w + bv4.w); -+ -+ float scores0[keys_per_warp]; -+ float scores1[keys_per_warp]; -+ float local_max0 = -INFINITY; -+ float local_max1 = -INFINITY; -+ float4 p14 = relpos_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); -+ 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); -+ } else { -+ k4 = relpos_load4(Kh + (size_t) (j - cache_len) * k_sj + d4); -+ } -+ float score0 = -+ k4.x * qc0.x + p04.x * qp0.x + -+ k4.y * qc0.y + p04.y * qp0.y + -+ k4.z * qc0.z + p04.z * qp0.z + -+ k4.w * qc0.w + p04.w * qp0.w; -+ float score1 = -+ k4.x * qc1.x + p14.x * qp1.x + -+ 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); -+ const float additive_mask = Mb ? Mb[j] : 0.0f; -+ score0 = score0 * scale + additive_mask; -+ score1 = score1 * scale + additive_mask; -+ scores0[c] = score0; -+ scores1[c] = score1; -+ local_max0 = fmaxf(local_max0, score0); -+ local_max1 = fmaxf(local_max1, score1); -+ p14 = p04; -+ } -+ if (lane == 0) { -+ den[warp * 2] = local_max0; -+ den[warp * 2 + 1] = local_max1; -+ } -+ __syncthreads(); -+ -+ float maximum0 = den[0]; -+ float maximum1 = den[1]; -+#pragma unroll -+ for (int w = 1; w < warps; ++w) { -+ maximum0 = fmaxf(maximum0, den[w * 2]); -+ maximum1 = fmaxf(maximum1, den[w * 2 + 1]); -+ } -+ // Every warp must finish reading the maxima before lane zero reuses den. -+ __syncthreads(); -+ float local_den0 = 0.0f; -+ float local_den1 = 0.0f; -+#pragma unroll -+ for (int c = 0; c < keys_per_warp; ++c) { -+ const float weight0 = __expf(scores0[c] - maximum0); -+ const float weight1 = __expf(scores1[c] - maximum1); -+ scores0[c] = weight0; -+ scores1[c] = weight1; -+ local_den0 += weight0; -+ local_den1 += weight1; -+ } -+ if (lane == 0) { -+ den[warp * 2] = local_den0; -+ den[warp * 2 + 1] = local_den1; -+ } -+ -+ float4 acc0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); -+ float4 acc1 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); -+#pragma unroll -+ for (int c = 0; c < keys_per_warp; ++c) { -+ const int j = j0 + c; -+ float4 v4; -+ if (c < 7 || j < cache_len) { -+ int physical_j = ring_head + j; -+ if (physical_j >= cache_len) { -+ physical_j -= cache_len; -+ } -+ v4 = relpos_load4(Vch + (size_t) physical_j * cache_sj + d4); -+ } else { -+ v4 = relpos_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); -+ acc0.z = fmaf(scores0[c], v4.z, acc0.z); -+ acc0.w = fmaf(scores0[c], v4.w, acc0.w); -+ acc1.x = fmaf(scores1[c], v4.x, acc1.x); -+ acc1.y = fmaf(scores1[c], v4.y, acc1.y); -+ acc1.z = fmaf(scores1[c], v4.z, acc1.z); -+ acc1.w = fmaf(scores1[c], v4.w, acc1.w); -+ } -+ *reinterpret_cast( -+ part + ((warp * 2) * RELPOS_ATTN_DK_128) + d4) = acc0; -+ *reinterpret_cast( -+ part + ((warp * 2 + 1) * RELPOS_ATTN_DK_128) + d4) = acc1; -+ __syncthreads(); -+ -+ float total_den0 = den[0]; -+ float total_den1 = den[1]; -+#pragma unroll -+ for (int w = 1; w < warps; ++w) { -+ total_den0 += den[w * 2]; -+ total_den1 += den[w * 2 + 1]; -+ } -+ const int i = tid / RELPOS_ATTN_DK_128; -+ const int d = tid % RELPOS_ATTN_DK_128; -+ float context = part[i * RELPOS_ATTN_DK_128 + d]; -+#pragma unroll -+ for (int w = 1; w < warps; ++w) { -+ context += part[((w * 2 + i) * RELPOS_ATTN_DK_128) + d]; -+ } -+ const float denominator = i == 0 ? total_den0 : total_den1; -+ ctx[(size_t) b * o_sb + (size_t) i * o_si + (size_t) h * o_sh + d] = -+ 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. -+template -+__global__ __launch_bounds__(192, 3) void fused_relpos_attn_q4_register_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, -+ 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 cache_sj, long cache_ss, -+ long p_sr, long p_sh, -+ long o_si, long o_sh, long o_sb, long m_sb) { -+ constexpr int warps = 6; -+ constexpr int queries = 4; -+ constexpr int keys_per_warp = 10; -+ constexpr int cache_len = 56; -+ -+ extern __shared__ float sh[]; -+ float * den = sh; // [warp, query] -+ float * part = den + warps * queries; // [warp, query, d] -+ -+ const int tid = threadIdx.x; -+ const int lane = tid & 31; -+ const int warp = tid >> 5; -+ const int h = blockIdx.x; -+ const int b = blockIdx.y; -+ const int d4 = lane * 4; -+ const int j0 = warp * keys_per_warp; -+ -+ const float * Qh0 = Q + (size_t) b * q_sb + (size_t) h * q_sh; -+ const float * Qh1 = Qh0 + q_sq; -+ const float * Qh2 = Qh1 + q_sq; -+ const float * Qh3 = Qh2 + 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 = slot_ids[b]; -+ const int ring_head = ring_heads[b]; -+ const float * Kch = -+ Kcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128; -+ const float * Vch = -+ Vcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128; -+ 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); -+#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); -+ const float4 qc1 = RELPOS_ADD_BIAS(q14, bu4); -+ const float4 qc2 = RELPOS_ADD_BIAS(q24, bu4); -+ const float4 qc3 = RELPOS_ADD_BIAS(q34, bu4); -+ const float4 qp0 = RELPOS_ADD_BIAS(q04, bv4); -+ const float4 qp1 = RELPOS_ADD_BIAS(q14, bv4); -+ const float4 qp2 = RELPOS_ADD_BIAS(q24, bv4); -+ const float4 qp3 = RELPOS_ADD_BIAS(q34, bv4); -+#undef RELPOS_ADD_BIAS -+ -+ float scores0[keys_per_warp]; -+ float scores1[keys_per_warp]; -+ float scores2[keys_per_warp]; -+ float scores3[keys_per_warp]; -+ float local_max0 = -INFINITY; -+ float local_max1 = -INFINITY; -+ float local_max2 = -INFINITY; -+ float local_max3 = -INFINITY; -+#pragma unroll -+ for (int c = 0; c < keys_per_warp; ++c) { -+ const int j = j0 + c; -+ const float4 k4 = relpos_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); -+#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)); -+#undef RELPOS_SCORE -+ const float additive_mask = Mb ? Mb[j] : 0.0f; -+ score0 = score0 * scale + additive_mask; -+ score1 = score1 * scale + additive_mask; -+ score2 = score2 * scale + additive_mask; -+ score3 = score3 * scale + additive_mask; -+ scores0[c] = score0; -+ scores1[c] = score1; -+ scores2[c] = score2; -+ scores3[c] = score3; -+ local_max0 = fmaxf(local_max0, score0); -+ local_max1 = fmaxf(local_max1, score1); -+ local_max2 = fmaxf(local_max2, score2); -+ local_max3 = fmaxf(local_max3, score3); -+ } -+ if (lane == 0) { -+ den[warp * queries + 0] = local_max0; -+ den[warp * queries + 1] = local_max1; -+ den[warp * queries + 2] = local_max2; -+ den[warp * queries + 3] = local_max3; -+ } -+ __syncthreads(); -+ -+ float maximum0 = den[0]; -+ float maximum1 = den[1]; -+ float maximum2 = den[2]; -+ float maximum3 = den[3]; -+#pragma unroll -+ for (int w = 1; w < warps; ++w) { -+ maximum0 = fmaxf(maximum0, den[w * queries + 0]); -+ maximum1 = fmaxf(maximum1, den[w * queries + 1]); -+ maximum2 = fmaxf(maximum2, den[w * queries + 2]); -+ maximum3 = fmaxf(maximum3, den[w * queries + 3]); -+ } -+ __syncthreads(); -+ -+ float local_den0 = 0.0f; -+ float local_den1 = 0.0f; -+ float local_den2 = 0.0f; -+ float local_den3 = 0.0f; -+#pragma unroll -+ for (int c = 0; c < keys_per_warp; ++c) { -+ scores0[c] = __expf(scores0[c] - maximum0); -+ scores1[c] = __expf(scores1[c] - maximum1); -+ scores2[c] = __expf(scores2[c] - maximum2); -+ scores3[c] = __expf(scores3[c] - maximum3); -+ local_den0 += scores0[c]; -+ local_den1 += scores1[c]; -+ local_den2 += scores2[c]; -+ local_den3 += scores3[c]; -+ } -+ if (lane == 0) { -+ den[warp * queries + 0] = local_den0; -+ den[warp * queries + 1] = local_den1; -+ den[warp * queries + 2] = local_den2; -+ den[warp * queries + 3] = local_den3; -+ } -+ -+ float4 acc0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); -+ float4 acc1 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); -+ float4 acc2 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); -+ float4 acc3 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); -+#pragma unroll -+ for (int c = 0; c < keys_per_warp; ++c) { -+ const int j = j0 + c; -+ const float4 v4 = relpos_load_kv4( -+ Vh, Vch, j, cache_len, ring_head, v_sj, cache_sj, d4); -+#define RELPOS_ACCUM(acc, weight, value) \ -+ do { \ -+ (acc).x = fmaf((weight), (value).x, (acc).x); \ -+ (acc).y = fmaf((weight), (value).y, (acc).y); \ -+ (acc).z = fmaf((weight), (value).z, (acc).z); \ -+ (acc).w = fmaf((weight), (value).w, (acc).w); \ -+ } while (0) -+ RELPOS_ACCUM(acc0, scores0[c], v4); -+ RELPOS_ACCUM(acc1, scores1[c], v4); -+ RELPOS_ACCUM(acc2, scores2[c], v4); -+ RELPOS_ACCUM(acc3, scores3[c], v4); -+#undef RELPOS_ACCUM -+ } -+ *reinterpret_cast( -+ part + ((warp * queries + 0) * RELPOS_ATTN_DK_128) + d4) = acc0; -+ *reinterpret_cast( -+ part + ((warp * queries + 1) * RELPOS_ATTN_DK_128) + d4) = acc1; -+ *reinterpret_cast( -+ part + ((warp * queries + 2) * RELPOS_ATTN_DK_128) + d4) = acc2; -+ *reinterpret_cast( -+ part + ((warp * queries + 3) * RELPOS_ATTN_DK_128) + d4) = acc3; -+ __syncthreads(); -+ -+ for (int index = tid; index < queries * RELPOS_ATTN_DK_128; index += blockDim.x) { -+ const int query = index / RELPOS_ATTN_DK_128; -+ const int d = index % RELPOS_ATTN_DK_128; -+ float denominator = den[query]; -+ float context = part[query * RELPOS_ATTN_DK_128 + d]; -+#pragma unroll -+ for (int w = 1; w < warps; ++w) { -+ denominator += den[w * queries + query]; -+ context += part[(w * queries + query) * RELPOS_ATTN_DK_128 + d]; -+ } -+ ctx[(size_t) b * o_sb + (size_t) query * o_si + (size_t) h * o_sh + d] = -+ context / denominator; -+ } -+} -+ -+// 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. -+template -+__global__ __launch_bounds__(256, 4) void fused_relpos_attn_common_q2_register_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 q_len, int kv_len, 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 cache_sj, long cache_ss, -+ long p_sr, long p_sh, -+ long o_si, long o_sh, long o_sb, long m_sb) { -+ constexpr int warps = 8; -+ constexpr int queries = 2; -+ constexpr int keys_per_warp = 11; -+ -+ extern __shared__ float sh[]; -+ float * den = sh; -+ float * part = den + warps * queries; -+ -+ const int tid = threadIdx.x; -+ const int lane = tid & 31; -+ const int warp = tid >> 5; -+ const int h = blockIdx.x; -+ const int b = blockIdx.y; -+ const int query0 = blockIdx.z * queries; -+ const bool has_query1 = query0 + 1 < q_len; -+ const int d4 = lane * 4; -+ const int j0 = warp * keys_per_warp; -+ -+ const float * Qh0 = -+ Q + (size_t) b * q_sb + (size_t) query0 * q_sq + (size_t) h * q_sh; -+ const float * Qh1 = has_query1 ? Qh0 + q_sq : Qh0; -+ 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 = slot_ids[b]; -+ const int ring_head = ring_heads[b]; -+ const float * Kch = -+ Kcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128; -+ const float * Vch = -+ Vcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128; -+ 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); -+#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); -+ const float4 qc1 = RELPOS_ADD_BIAS(q14, bu4); -+ const float4 qp0 = RELPOS_ADD_BIAS(q04, bv4); -+ const float4 qp1 = RELPOS_ADD_BIAS(q14, bv4); -+#undef RELPOS_ADD_BIAS -+ -+ float scores0[keys_per_warp]; -+ float scores1[keys_per_warp]; -+ float local_max0 = -INFINITY; -+ float local_max1 = -INFINITY; -+#pragma unroll -+ for (int c = 0; c < keys_per_warp; ++c) { -+ const int j = j0 + c; -+ float score0 = -INFINITY; -+ float score1 = -INFINITY; -+ if (j < kv_len) { -+ const float4 k4 = relpos_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); -+#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)); -+ 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)); -+ } -+#undef RELPOS_SCORE -+ const float additive_mask = Mb ? Mb[j] : 0.0f; -+ score0 = score0 * scale + additive_mask; -+ if (has_query1) -+ score1 = score1 * scale + additive_mask; -+ } -+ scores0[c] = score0; -+ scores1[c] = score1; -+ local_max0 = fmaxf(local_max0, score0); -+ local_max1 = fmaxf(local_max1, score1); -+ } -+ if (lane == 0) { -+ den[warp * queries] = local_max0; -+ den[warp * queries + 1] = local_max1; -+ } -+ __syncthreads(); -+ -+ float maximum0 = den[0]; -+ float maximum1 = den[1]; -+#pragma unroll -+ for (int w = 1; w < warps; ++w) { -+ maximum0 = fmaxf(maximum0, den[w * queries]); -+ maximum1 = fmaxf(maximum1, den[w * queries + 1]); -+ } -+ __syncthreads(); -+ -+ float local_den0 = 0.0f; -+ float local_den1 = 0.0f; -+#pragma unroll -+ for (int c = 0; c < keys_per_warp; ++c) { -+ const int j = j0 + c; -+ if (j < kv_len) { -+ scores0[c] = __expf(scores0[c] - maximum0); -+ scores1[c] = has_query1 ? __expf(scores1[c] - maximum1) : 0.0f; -+ local_den0 += scores0[c]; -+ local_den1 += scores1[c]; -+ } else { -+ scores0[c] = 0.0f; -+ scores1[c] = 0.0f; -+ } -+ } -+ if (lane == 0) { -+ den[warp * queries] = local_den0; -+ den[warp * queries + 1] = local_den1; -+ } -+ -+ float4 acc0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); -+ float4 acc1 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); -+#pragma unroll -+ for (int c = 0; c < keys_per_warp; ++c) { -+ const int j = j0 + c; -+ if (j >= kv_len) -+ continue; -+ const float4 v4 = relpos_load_kv4( -+ Vh, Vch, j, cache_len, ring_head, v_sj, cache_sj, d4); -+#define RELPOS_ACCUM(acc, weight, value) \ -+ do { \ -+ (acc).x = fmaf((weight), (value).x, (acc).x); \ -+ (acc).y = fmaf((weight), (value).y, (acc).y); \ -+ (acc).z = fmaf((weight), (value).z, (acc).z); \ -+ (acc).w = fmaf((weight), (value).w, (acc).w); \ -+ } while (0) -+ RELPOS_ACCUM(acc0, scores0[c], v4); -+ RELPOS_ACCUM(acc1, scores1[c], v4); -+#undef RELPOS_ACCUM -+ } -+ *reinterpret_cast( -+ part + ((warp * queries) * RELPOS_ATTN_DK_128) + d4) = acc0; -+ *reinterpret_cast( -+ part + ((warp * queries + 1) * RELPOS_ATTN_DK_128) + d4) = acc1; -+ __syncthreads(); -+ -+ const int query = tid / RELPOS_ATTN_DK_128; -+ if (query == 0 || has_query1) { -+ const int d = tid % RELPOS_ATTN_DK_128; -+ float denominator = den[query]; -+ float context = part[query * RELPOS_ATTN_DK_128 + d]; -+#pragma unroll -+ for (int w = 1; w < warps; ++w) { -+ denominator += den[w * queries + query]; -+ context += part[(w * queries + query) * RELPOS_ATTN_DK_128 + d]; -+ } -+ ctx[(size_t) b * o_sb + (size_t) (query0 + query) * o_si + -+ (size_t) h * o_sh + d] = context / denominator; -+ } -+} -+ -+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,31 @@ 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)); -+ &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 -+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, -+ 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, - // 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 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,8 +1033,16 @@ 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 int slot = Cached ? slot_ids[b] : 0; -+ const int ring_head = Cached ? ring_heads[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 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( - 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]; -- } -+ 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]; - #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( - } - } 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]; -+ ac += relpos_load_kv( -+ Kh, Kch, j, cache_len, ring_head, k_sj, cache_sj, dd) * -+ Qu[dd]; - 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( - - // 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++) { -+ 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 - 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 bool cached = kv_cache != 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 - 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]; -+ 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[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[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 (cached) { -+ GGML_ASSERT(slot_ids != nullptr); -+ GGML_ASSERT(ring_heads != 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); -+ } else { -+ GGML_ASSERT(slot_ids == nullptr); -+ GGML_ASSERT(ring_heads == 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 (cached) { -+ // Streaming: shared across the batch or one key-mask column per -+ // stream. -+ GGML_ASSERT(mask->ne[1] == 1 || mask->ne[1] == batch); -+ GGML_ASSERT(mask->ne[2] == 1 && mask->ne[3] == 1); -+ } else { -+ // Offline: shared/per-batch key vector or one column per query. -+ 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; -+ (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,6 +1217,9 @@ 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 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 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 - GGML_ASSERT(shmem <= max_shmem && "fused_relpos_attn: 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; -+ -+ auto update_cache = [&]() { -+ if (!cached) { -+ return; -+ } -+ const dim3 update_grid((d_k * n_head + 255) / 256, batch, 2); -+ if (k->type == GGML_TYPE_F16) { -+ fused_relpos_attn_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<<>>( -+ (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 -@@ -523,6 +1267,114 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor - 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; -+ d_k == RELPOS_ATTN_DK_128 && q_len == 2 && kv_len == 72 && m_sq == 0; -+ const bool use_register_q2 = -+ 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 = -+ 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 = -+ 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(); -+ if (use_register_q4) { -+ constexpr int register_warps = 6; -+ constexpr int register_queries = 4; -+ const size_t register_shmem = -+ ((size_t) register_warps * register_queries + -+ (size_t) register_warps * register_queries * RELPOS_ATTN_DK_128) * sizeof(float); -+ GGML_ASSERT(register_shmem <= max_shmem); -+ const dim3 register_grid(n_head, batch, 1); -+ if (k->type == GGML_TYPE_F16) { -+ fused_relpos_attn_q4_register_kernel<<< -+ register_grid, 192, register_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, cache_k, cache_v, active_slots, active_ring_heads, 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_q4_register_kernel<<< -+ register_grid, 192, register_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, cache_k, cache_v, active_slots, active_ring_heads, 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); -+ } -+ update_cache(); -+ return; -+ } -+ if (use_register_q2) { -+ constexpr int register_warps = 8; -+ const size_t register_shmem = -+ ((size_t) register_warps * 2 + -+ (size_t) register_warps * 2 * RELPOS_ATTN_DK_128) * sizeof(float); -+ GGML_ASSERT(register_shmem <= max_shmem); -+ const dim3 register_grid(n_head, batch, 1); -+ if (k->type == GGML_TYPE_F16) { -+ fused_relpos_attn_q2_register_kernel<<< -+ register_grid, 256, register_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, cache_k, cache_v, active_slots, active_ring_heads, 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_q2_register_kernel<<< -+ register_grid, 256, register_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, cache_k, cache_v, active_slots, active_ring_heads, 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); -+ } -+ update_cache(); -+ return; -+ } -+ if (use_register_common) { -+ constexpr int register_warps = 8; -+ constexpr int register_queries = 2; -+ const size_t register_shmem = -+ ((size_t) register_warps * register_queries + -+ (size_t) register_warps * register_queries * RELPOS_ATTN_DK_128) * sizeof(float); -+ GGML_ASSERT(register_shmem <= max_shmem); -+ const dim3 register_grid(n_head, batch, (q_len + register_queries - 1) / register_queries); -+ if (k->type == GGML_TYPE_F16) { -+ fused_relpos_attn_common_q2_register_kernel<<< -+ register_grid, 256, register_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, cache_k, cache_v, active_slots, active_ring_heads, -+ q_len, 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_common_q2_register_kernel<<< -+ register_grid, 256, register_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, cache_k, cache_v, active_slots, active_ring_heads, -+ q_len, 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); -+ } -+ 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 +1383,86 @@ 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); -+ int max_single_blocks_per_sm; -+ if (k->type == GGML_TYPE_F16) { -+ max_single_blocks_per_sm = cached -+ ? relpos_attn_warp_128_max_blocks_per_sm(device, warp_shmem) -+ : relpos_attn_warp_128_max_blocks_per_sm(device, warp_shmem); -+ } else { -+ max_single_blocks_per_sm = cached -+ ? 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) { -+ 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<<< -+ 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); -+ (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<<< -+ 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); -+ (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); -+ } -+ }; -+ if (k->type == GGML_TYPE_F16) { -+ if (cached) { -+ 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); -+ 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); -+ 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, -+ auto launch_generic = [&](auto scalar, auto cache_tag) { -+ using T = decltype(scalar); -+ 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, - 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, -+ 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); -+ }; -+ if (k->type == GGML_TYPE_F16) { -+ if (cached) { -+ launch_generic(half{}, std::true_type{}); -+ } else { -+ launch_generic(half{}, std::false_type{}); -+ } - } 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, -- 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); -+ if (cached) { -+ launch_generic(float{}, std::true_type{}); -+ } else { -+ launch_generic(float{}, std::false_type{}); -+ } - } -+ update_cache(); - } -diff --git a/src/ggml.c b/src/ggml.c -index 80b5802f..ae1b8db3 100644 ---- a/src/ggml.c -+++ b/src/ggml.c -@@ -5377,7 +5377,7 @@ struct ggml_tensor * ggml_flash_attn_ext( - - // ggml_fused_relpos_attn - --struct ggml_tensor * ggml_fused_relpos_attn( -+static struct ggml_tensor * ggml_fused_relpos_attn_impl( - struct ggml_context * ctx, - struct ggml_tensor * q, - struct ggml_tensor * k, -@@ -5386,6 +5386,10 @@ struct ggml_tensor * ggml_fused_relpos_attn( - struct ggml_tensor * bias_u, - struct ggml_tensor * bias_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) { - // 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)); - -- 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]; -+ 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( - // 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(v->ne[1] == kv_len); -+ 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)); -+ if (kv_cache) { -+ GGML_ASSERT(cache_len > 0); -+ GGML_ASSERT(kv_cache->type == GGML_TYPE_F32 && ggml_is_contiguous(kv_cache)); -+ GGML_ASSERT(kv_cache->ne[0] == 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] == 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]); -+ } else { -+ GGML_ASSERT(cache_len == 0); -+ } - if (mask) { - GGML_ASSERT(ggml_is_contiguous(mask)); - GGML_ASSERT(mask->ne[0] == kv_len); -- // One shared key mask, or one column per batch item (the cache-aware -- // streaming layout, where each stream's history has its own validity). -- GGML_ASSERT(mask->ne[1] == 1 || mask->ne[1] == q->ne[3]); -- GGML_ASSERT(mask->ne[2] == 1 && mask->ne[3] == 1); -+ if (kv_cache) { -+ // Cache-aware streaming: one shared key mask, or one column per -+ // batch item whose history has its own validity. -+ GGML_ASSERT(mask->ne[1] == 1 || mask->ne[1] == q->ne[3]); -+ GGML_ASSERT(mask->ne[2] == 1 && mask->ne[3] == 1); -+ } else { -+ // Offline: a shared/per-batch key vector or a [key,query] -+ // local-attention mask, broadcast over heads. -+ GGML_ASSERT(mask->ne[1] == 1 || mask->ne[1] == q_len); -+ GGML_ASSERT(mask->ne[2] == 1 && (mask->ne[3] == 1 || mask->ne[3] == q->ne[3])); -+ } - } - - // Output mirrors q logically: [d_k, q_len, n_head, batch]. With -@@ -5434,6 +5461,7 @@ 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->src[0] = q; -@@ -5443,10 +5471,48 @@ 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; - - return result; - } - -+struct ggml_tensor * ggml_fused_relpos_attn( -+ struct ggml_context * ctx, -+ struct ggml_tensor * q, -+ struct ggml_tensor * k, -+ struct ggml_tensor * v, -+ struct ggml_tensor * p, -+ struct ggml_tensor * bias_u, -+ struct ggml_tensor * bias_v, -+ struct ggml_tensor * mask, -+ float scale, -+ bool merge_heads) { -+ return ggml_fused_relpos_attn_impl( -+ ctx, q, k, v, p, bias_u, bias_v, mask, NULL, NULL, NULL, 0, scale, merge_heads); -+} -+ -+struct ggml_tensor * ggml_fused_relpos_attn_cached( -+ struct ggml_context * ctx, -+ struct ggml_tensor * q, -+ struct ggml_tensor * k, -+ struct ggml_tensor * v, -+ struct ggml_tensor * p, -+ struct ggml_tensor * bias_u, -+ struct ggml_tensor * bias_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, p, bias_u, bias_v, mask, kv_cache, slot_ids, ring_heads, cache_len, scale, -+ merge_heads); -+} -+ - void ggml_flash_attn_ext_set_prec( - struct ggml_tensor * a, - enum ggml_prec prec) { 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/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 f3b7dca..409a75e 100644 --- a/ggml-patches/README.md +++ b/ggml-patches/README.md @@ -54,162 +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`. Turing and older GPUs retain stock block-Q8 - matmul; wide planar Q8 fails explicitly because its tensor-wide layout has - no stock fallback. 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; - 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 the cache-aware and offline FastConformer paths. 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 - 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. + 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 @@ -252,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/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/server/http/http_server.cpp b/server/http/http_server.cpp index a962c22..fee4fc2 100644 --- a/server/http/http_server.cpp +++ b/server/http/http_server.cpp @@ -10,9 +10,12 @@ #include #include #include +#include #include #include #include +#include +#include #include #include #include @@ -317,14 +320,70 @@ 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_.notify_all(); + 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)) { + 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) @@ -638,12 +697,58 @@ 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]" + << " 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..94ea243 100644 --- a/server/http/http_server.h +++ b/server/http/http_server.h @@ -19,6 +19,10 @@ 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; + // 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/src/runtime/ggml/backend.cpp b/src/runtime/ggml/backend.cpp index 886648a..ac60330 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,56 @@ 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 +145,9 @@ 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 4864071..aa68037 100644 --- a/src/runtime/ggml/runtime.h +++ b/src/runtime/ggml/runtime.h @@ -170,6 +170,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(); @@ -193,6 +197,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; @@ -233,6 +238,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( @@ -284,6 +293,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 f5b12dc..1a2d8df 100644 --- a/src/runtime/ggml/session.cpp +++ b/src/runtime/ggml/session.cpp @@ -321,6 +321,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/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/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/config.cpp b/src/tts/magpietts/config.cpp index 637ea02..eb2472d 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 @@ -188,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) { @@ -281,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/decoder.cpp b/src/tts/magpietts/decoder.cpp index 7e71a07..fa6881d 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,389 @@ 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, + int stacked_position_budget) + : model_(model), cross_kv_(&cross_kv), text_len_(text_len), + 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) { + 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, 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_; } + + 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 stacked_position_budget_ = 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, @@ -443,10 +830,54 @@ 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, - 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 { + 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 { + 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, 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_->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()); + persistent_runtime_.reset(); + cond_kv.clear(); + uncond_kv.clear(); + } + } 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, @@ -502,11 +933,39 @@ 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 +978,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 +998,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); @@ -588,10 +1048,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); @@ -601,7 +1065,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); @@ -623,15 +1089,15 @@ 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, - 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; @@ -642,12 +1108,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)); @@ -666,17 +1136,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 +1156,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); @@ -731,10 +1202,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); @@ -742,15 +1218,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); @@ -786,7 +1268,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 +1279,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; @@ -808,20 +1290,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)); @@ -842,11 +1332,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 +1369,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 +1386,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); } @@ -950,11 +1441,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); @@ -983,15 +1478,15 @@ 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, - 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) { @@ -1006,12 +1501,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)); @@ -1033,11 +1532,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 +1590,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 +1602,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 +1618,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; @@ -1164,10 +1664,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); @@ -1175,15 +1680,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); @@ -1219,7 +1730,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 +1741,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) { @@ -1247,20 +1758,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)); @@ -1374,14 +1893,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 +1914,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/decoder.h b/src/tts/magpietts/decoder.h index 996bc2e..84569f2 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, @@ -115,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, @@ -123,8 +131,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 292768b..dd4f068 100644 --- a/src/tts/magpietts/lt.cpp +++ b/src/tts/magpietts/lt.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include #include @@ -13,8 +14,37 @@ #include "graph.h" #include "nvtx_utils.h" +#if defined(MAGPIETTS_CUDA_SAMPLING) +#include "ggml-cuda.h" +#endif + 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; @@ -35,6 +65,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; @@ -64,6 +95,7 @@ class LocalTransformerGraphBank { DecoderKvCache single_cache; DecoderKvCache cond_cache; DecoderKvCache uncond_cache; + LocalTransformerCudaAttentionCache pair_cuda_attention_cache; }; using local_transformer_graph = LocalTransformerGraph; @@ -253,8 +285,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) || - !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)) { @@ -309,6 +343,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(); } @@ -328,6 +478,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; @@ -343,6 +494,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; @@ -367,6 +519,7 @@ LocalTransformerGraph::reset() { dec_uncond = nullptr; pos_emb = nullptr; prev_token = nullptr; + cache_state = nullptr; logits_cond = nullptr; logits_uncond = nullptr; codebook_idx = -1; @@ -393,6 +546,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; } @@ -404,12 +558,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( @@ -427,6 +585,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, @@ -452,6 +637,109 @@ 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, + 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 = 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; + 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, @@ -462,7 +750,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; } @@ -485,6 +773,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) { @@ -501,20 +795,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 = linear(graph.ctx, model.lt_in_w, graph.dec_cond, model.lt_in_b); + 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 = linear(graph.ctx, model.lt_in_w, graph.dec_uncond, model.lt_in_b); + input_uncond = graph.dec_uncond; } } else { const std::string name = "magpietts_local_transformer_prev_code"; @@ -523,10 +817,9 @@ 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); - if (pair) { - cur_uncond = cur_cond; - } + input_cond = emb; + if (pair) + input_uncond = emb; } graph.pos_emb = ggml_view_2d( @@ -534,30 +827,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)); + 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, + 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); + graph.logits_uncond = ggml_view_2d( + graph.ctx, logits_pair, h.audio_vocab_size, 1, logits_pair->nb[1], + 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); 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); } @@ -704,13 +1012,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) { @@ -737,7 +1073,20 @@ 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( @@ -754,18 +1103,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", @@ -827,6 +1173,67 @@ 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, @@ -837,13 +1244,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.audio_codebooks; ++c) { + for (int c = 0; c < stacked_codebooks; ++c) { std::vector logits; if (use_cfg) { std::vector uncond; @@ -864,9 +1282,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 - ? (*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); @@ -887,33 +1303,101 @@ 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.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); } } - codes.assign((size_t)h.audio_codebooks, 0); - argmax_codes.assign((size_t)h.audio_codebooks, 0); - char error[256] = {}; + 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); + error[0] = '\0'; 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/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 00448ba..5ff1528 100644 --- a/src/tts/magpietts/magpietts.cpp +++ b/src/tts/magpietts/magpietts.cpp @@ -352,6 +352,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 = {}; @@ -918,11 +929,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; } @@ -1097,6 +1109,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, @@ -1142,7 +1158,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; @@ -1257,7 +1273,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); @@ -1309,8 +1325,14 @@ 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"); + 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; @@ -1318,12 +1340,14 @@ 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; + 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; std::vector alignment_scores; @@ -1346,9 +1370,9 @@ 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, - &uncond_hidden_device, &cond_cross_kv, - decoder_attention_arg) + max_decoder_positions, nullptr, &text_cond_device, + &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, @@ -1399,8 +1423,9 @@ 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, @@ -1422,8 +1447,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(); @@ -1434,9 +1459,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, @@ -1472,13 +1497,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 { @@ -1525,8 +1560,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 && eos_lane < frames_remaining; + bool terminate_after_frame = false; if (has_eos) { ggml_nvtx::mark("magpietts_stream_eos"); if (params.verbose) { @@ -1535,7 +1577,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; @@ -1551,31 +1593,31 @@ 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 = 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]; + 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; @@ -1600,6 +1642,9 @@ stream_magpie_to_audio( } break; } + if (terminate_after_frame) { + break; + } } if (longform_active && !final_text_chunk) { diff --git a/src/tts/magpietts/magpietts.h b/src/tts/magpietts/magpietts.h index 40e4074..228dc75 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; @@ -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/magpietts_cuda_sampling.cu b/src/tts/magpietts/magpietts_cuda_sampling.cu index ad0ecf3..5fd454b 100644 --- a/src/tts/magpietts/magpietts_cuda_sampling.cu +++ b/src/tts/magpietts/magpietts_cuda_sampling.cu @@ -5,19 +5,45 @@ #include #include #include +#include #include "magpietts_cuda_sampling.h" 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 +124,25 @@ 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; + __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 +151,46 @@ 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; - } + // 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]; } - - 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]; - } - __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 +202,12 @@ 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 +230,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 +254,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 +264,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 +300,267 @@ 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 +588,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 +621,56 @@ 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) + : 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; + 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 +696,18 @@ 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 +716,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..5e7fa75 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 de1ee04..baf64c6 100644 --- a/src/tts/magpietts/model.cpp +++ b/src/tts/magpietts/model.cpp @@ -610,6 +610,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; @@ -672,6 +674,8 @@ MagpieModel::reset() { ctx = nullptr; } hparams = {}; + tokenizer_profile.clear(); + nemo_version.clear(); cuda_unified_memory = false; text_embedding = nullptr; audio_embeddings.clear(); @@ -715,6 +719,28 @@ 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", 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); @@ -793,10 +819,22 @@ magpietts_model_load_impl( return false; } - if (h.frame_stacking_factor != 1) { + 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 frame_stacking_factor=%d; this example currently supports 1\n", - h.frame_stacking_factor); + 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 (stored_stacked_codebooks != expected_stacked_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; } @@ -879,11 +917,23 @@ 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"); } @@ -908,28 +958,27 @@ 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( model, "local_transformer_out_projections." + std::to_string(i) + ".bias"); } - if (verbose) { - fprintf( - stderr, - "loaded MagpieTTS GGUF: text_vocab=%d audio_codebooks=%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.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(), - format_i32_list(h.apply_prior_to_layers).c_str()); - } + fprintf( + stderr, + "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.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(), + format_i32_list(h.apply_prior_to_layers).c_str()); return true; } @@ -1125,11 +1174,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; @@ -1149,8 +1198,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); @@ -1810,12 +1864,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)); @@ -1825,6 +1884,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"); @@ -1850,6 +1912,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, @@ -1904,7 +1970,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; @@ -1941,7 +2007,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; @@ -1957,14 +2023,20 @@ 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(); + 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, 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; @@ -1979,6 +2051,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 = @@ -1986,9 +2060,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, - &uncond_hidden_device, &cond_cross_kv, + 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( text_cond, (int)tokens.size(), audio_codes, params.speaker, @@ -2038,8 +2112,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, @@ -2061,8 +2135,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; } @@ -2072,8 +2146,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, @@ -2114,32 +2189,55 @@ 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)) { - ggml_nvtx::mark("magpietts_eos"); - fprintf(stderr, "%s EOS detected at frame %d\n", label, step); - break; + std::vector> codec_frames; + if (!magpietts_unstack_codes(next_codes, h, codec_frames)) { + fprintf(stderr, "sampled an invalid stacked MagpieTTS frame\n"); + return false; } - - generated_frames.push_back(next_codes); + 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) { - 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]); + } + } + 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); } - 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 e667f7b..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 @@ -50,6 +51,11 @@ 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 +93,67 @@ 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, + 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; +} + +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; @@ -203,6 +270,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 0209687..3436136 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, config_.verbose)) { 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 38e53b0..aafea86 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; @@ -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 349ec37..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())); + } } } @@ -304,7 +313,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/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..262a10f 100644 --- a/src/tts/tokenizer/mandarin_tokenizer.h +++ b/src/tts/tokenizer/mandarin_tokenizer.h @@ -9,7 +9,8 @@ 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 05d5623..9e8d763 100644 --- a/src/tts/tokenizer/tokenizer.cpp +++ b/src/tts/tokenizer/tokenizer.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -16,6 +17,508 @@ 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", + profile.id == "v2602" ? "upper" : "lower"); + 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) { @@ -49,6 +552,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 +652,7 @@ class MagpieNativeTokenizer::Impl { if (model_dir_.empty()) { throw std::invalid_argument("tokenizer model directory is required"); } + profile_ = load_tokenizer_profile(model_dir_); } MagpieTokenizationResult tokenize( @@ -158,18 +668,33 @@ class MagpieNativeTokenizer::Impl { p.language = MagpieNativeTokenizer::normalize_language_code(language_code); 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); + tokenizer_result native = tokenize_native(p, *selected); MagpieTokenizationResult out; out.language = native.language; out.tokenizer_name = native.tokenizer_name; @@ -188,17 +713,76 @@ class MagpieNativeTokenizer::Impl { return out; } + 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"}, + }; + 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); + } + } + 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 @@ -210,7 +794,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_; } @@ -221,13 +807,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; @@ -252,6 +845,7 @@ class MagpieNativeTokenizer::Impl { std::string model_dir_; MagpieTokenizerConfig config_; + tokenizer_profile profile_; mutable std::mutex cache_mutex_; mutable std::map> ipa_cache_; #ifdef NEMO_SPEECH_TTS_WITH_ZH @@ -299,27 +893,30 @@ 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(); +} + +const std::string& +MagpieNativeTokenizer::profile_id() const { + return impl_->profile_id(); +} + +int +MagpieNativeTokenizer::text_vocab_size() const { + return impl_->text_vocab_size(); } std::string @@ -337,7 +934,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; @@ -354,7 +952,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/src/tts/tokenizer/tokenizer.h b/src/tts/tokenizer/tokenizer.h index 6dedc50..9853752 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,9 @@ 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; + 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 7948fd6..13009cc 100644 --- a/src/tts/tokenizer/tokenizer_impl.cpp +++ b/src/tts/tokenizer/tokenizer_impl.cpp @@ -53,6 +53,18 @@ struct params { fs::path model; std::string text; std::string language = "en"; + 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; }; @@ -65,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; }; @@ -102,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'); @@ -331,21 +334,29 @@ hindi_tokens() { #ifdef NEMO_SPEECH_TTS_WITH_JA static std::vector -japanese_tokens() { - return {" ", "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", "!", "\"", "(", - ")", ",", "-", ".", "/", ":", ";", "?", "[", "]", "{", "}", "«", - "»", "•", "‥", "…", "‹", "›", "※", "◦", "、", "。", "〃", "〈", "〉", - "《", "》", "「", "」", "『", "』", "【", "】", "〒", "〓", "〔", "〕", "〖", - "〗", "〘", "〙", "〚", "〛", "〜", "〽", "・", "・・・", "ー", "﹅", "﹆", "！", - "＊", "？", "｟", "｠", "", ""}; +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", "!", "\"", "(", ")", ",", "-", ".", "/", ":", ";", "?", "[", + "]", "{", "}", "«", "»", "•", "‥", "…", "‹", "›", "※", "◦", "、", "。", + "〃", "〈", "〉", "《", "》", "「", "」", "『", "』", "【", "】", "〒", "〓", "〔", + "〕", "〖", "〗", "〘", "〙", "〚", "〛", "〜", "〽", "・", "・・・", "ー", "﹅", "﹆", + "！", "＊", "？", "｟", "｠", "", ""}; + 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 @@ -371,6 +382,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()); } @@ -850,50 +864,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"; -#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; @@ -907,34 +877,14 @@ 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}; - } - 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 @@ -976,18 +926,44 @@ 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", "{", "}", + "ð", "õ", "ĩ", "ŋ", "ũ", "ɐ", "ɑ", "ɒ", "ɔ", "ə", "ɛ", "ɜ", "ɡ", + "ɪ", "ɲ", "ɹ", "ɾ", "ʁ", "ʃ", "ʊ", "ʌ", "ʎ", "ʒ", "ʲ", "ˈ", "ˌ", + "ː", "̃", "θ", "ẽ", " ", "", ""}; + } + 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") { @@ -1054,17 +1030,17 @@ 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()) { + 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 containing '" + cfg_.dict_hint + "'"); + "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)) { @@ -1140,12 +1116,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; @@ -1334,42 +1310,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(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(p.offset + 1); 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; @@ -1377,17 +1331,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; @@ -1411,9 +1368,64 @@ 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); + pad_short_text_chunk_before_eos(ch, result.eos_id, pad_id); + result.chunks.push_back(std::move(ch)); + } + 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) { + 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 = 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; + 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(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)); @@ -1426,7 +1438,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; @@ -1441,12 +1454,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; } } @@ -1484,13 +1499,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; } @@ -1499,7 +1514,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); @@ -1536,7 +1551,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( @@ -1544,22 +1559,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); @@ -1577,7 +1597,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); @@ -1587,59 +1607,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; - } -#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); - } -#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/cli/cli_contract_test.py b/tests/cli/cli_contract_test.py index fa4bb53..8c8f1de 100644 --- a/tests/cli/cli_contract_test.py +++ b/tests/cli/cli_contract_test.py @@ -74,7 +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: + 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() @@ -142,7 +145,14 @@ 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"] if has_tts_options else []), + *missing_arguments, + "--no-warmup", + ), 3, "missing_model", ) 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() diff --git a/tests/conversion/tts_tokenizer_profiles_test.py b/tests/conversion/tts_tokenizer_profiles_test.py new file mode 100644 index 0000000..d6c908e --- /dev/null +++ b/tests/conversion/tts_tokenizer_profiles_test.py @@ -0,0 +1,76 @@ +#!/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/CMakeLists.txt b/tests/cpp/CMakeLists.txt index 9ff5030..4824102 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 @@ -72,6 +79,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}) @@ -87,6 +104,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/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 8f62b6b..0d558fe 100644 --- a/tests/cpp/tts/CMakeLists.txt +++ b/tests/cpp/tts/CMakeLists.txt @@ -8,6 +8,15 @@ 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_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) + 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 @@ -38,6 +47,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) @@ -45,6 +64,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..d6cf3ff 100644 --- a/tests/cpp/tts/test_grpc_tts_config.cpp +++ b/tests/cpp/tts/test_grpc_tts_config.cpp @@ -111,6 +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(); nemo_speech::GrpcTtsService service(std::move(synthesizer)); nr_tts::RivaSynthesisConfigRequest req; @@ -120,7 +121,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_asr.cpp b/tests/cpp/tts/test_magpietts_asr.cpp index dc8fce1..c780949 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 new file mode 100644 index 0000000..704c1b6 --- /dev/null +++ b/tests/cpp/tts/test_magpietts_cached_attention.cpp @@ -0,0 +1,210 @@ +// 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" +#include "tts/magpietts/model.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; + } + + 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(), + /*.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; +} diff --git a/tests/cpp/tts/test_magpietts_file.cpp b/tests/cpp/tts/test_magpietts_file.cpp index 23d63cf..2a4efa6 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" @@ -42,7 +46,8 @@ 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" " --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,52 @@ 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( @@ -258,17 +316,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; @@ -288,6 +380,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; } 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..7290e33 --- /dev/null +++ b/tests/cpp/tts/test_magpietts_frame_stacking.cpp @@ -0,0 +1,101 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include +#include + +#include "tts/magpietts/lt.h" +#include "tts/magpietts/model.h" + +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; + 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; + } + + 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) { + 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; + } + + 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; +} 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 a04a03b..1ed87e2 100644 --- a/tests/cpp/tts/test_tokenizer_single_chars.cpp +++ b/tests/cpp/tts/test_tokenizer_single_chars.cpp @@ -72,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; @@ -130,7 +145,76 @@ main(int argc, char** argv) { } tts::MagpieNativeTokenizer tokenizer(argv[1]); + 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( + 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, "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}); 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( 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(