Add files using upload-large-folder tool

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+

.gitignore

@@ -0,0 +1,6 @@
+.cache/
+.venv/
+__pycache__/
+*.py[cod]
+output.wav
+

LICENSE

@@ -0,0 +1,23 @@
+MIT License
+
+Copyright (c) 2026 Resemble AI
+Copyright (c) 2026 Chatterbox-Flash NVFP4 contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+

README.md

@@ -0,0 +1,288 @@
+---
+license: mit
+language:
+- en
+library_name: chatterbox-flash
+pipeline_tag: text-to-speech
+base_model: ResembleAI/chatterbox-flash
+base_model_relation: quantized
+inference: false
+tags:
+- text-to-speech
+- zero-shot-tts
+- voice-cloning
+- chatterbox-flash
+- nvfp4
+- w4a16
+- quantized
+- nvidia
+- blackwell
+- flashinfer
+- safetensors
+---
+
+# Chatterbox-Flash T3 W4A16 NVFP4
+
+This is an unofficial community NVFP4 derivative of
+[ResembleAI/chatterbox-flash](https://huggingface.co/ResembleAI/chatterbox-flash).
+It is not affiliated with or endorsed by Resemble AI or NVIDIA.
+
+The package replaces all 120 Chatterbox-Flash T3 transformer projection
+operators with source-free NVFP4 weights. Weights use native FP4 (E2M1) with
+E4M3 block scales; activations remain BF16, so the precise description is
+**T3 all-projection W4A16 NVFP4**. Embeddings, norms, conditioning layers,
+speech head, S3Gen, and the voice encoder remain at their upstream precision.
+This is not a whole-pipeline 4-bit model.
+
+The canonical artifact does not contain the original BF16 T3 projection
+weights. Its custom loader constructs those parameters on the meta device,
+prepares the canonical FP4 codes/scales directly for FlashInfer, and never
+calls `nvfp4_quantize` at load time. Engine rebuilds reuse the model-attached
+projection provider.
+
+## What this release actually improves
+
+The primary benefit is **much lower storage and VRAM with a modest speed gain**,
+not a claim that NVFP4 makes the complete TTS pipeline dramatically faster.
+S3Gen and several small T3 operations remain high precision, and some
+irregular prefix projections are slower in FP4 than BF16.
+
+On the tested RTX 5060 Ti, the recommended D16 quality profile delivered:
+
+- **83.97% smaller T3 checkpoint**: 341 MB instead of 2.13 GB.
+- **55.88% less complete model-asset storage/download** after adding the
+  unchanged S3Gen, voice encoder, and tokenizer: 1.412 GB instead of 3.200 GB.
+- **33.2% lower allocated VRAM immediately after full pipeline load**:
+  1.444 GB instead of 2.162 GB.
+- **18.82% lower maximum allocated VRAM in the natural eight-text suite**:
+  5.215 GB instead of 6.424 GB. Fixed-work peak allocation was 34.09% lower.
+- **1.0811x aggregate-RTF speedup** (7.50% lower duration-normalized latency)
+  and **1.1131x valid-token T3 throughput**. The fixed-work T3 speedup was
+  1.1260x; fixed whole-pipeline speedup was 1.0631x.
+
+The D24 fast profile reaches 25.13% lower RTF than the D16 BF16 baseline, but
+that is a combined block-size-plus-NVFP4 result. Against matched D24 BF16,
+NVFP4 alone provides only a 1.0309x RTF speedup. This distinction matters when
+deciding whether the model is useful for your deployment.
+
+## How quality was evaluated
+
+A Gemma Audio evaluator listened to eight English test utterances generated
+from one zero-shot reference voice. It transcribed each output, scored
+intelligibility, pronunciation, voice similarity, naturalness, and artifacts,
+and flagged hard failures. BF16 and NVFP4 were also compared blindly in both
+A/B orders to expose position bias.
+
+- D16 BF16/NVFP4 means were **94.125/94.625**. Both produced exact transcripts
+  for 8/8 texts with zero hard failures. After reversing A/B order, the
+  per-case result resolved to seven ties and one consistent BF16 preference.
+- D24 BF16/NVFP4 means were **95.0/93.125**. Both produced exact transcripts
+  for 8/8 texts with zero hard failures. Order control resolved to six ties,
+  one BF16 preference, and one NVFP4 preference.
+
+This is encouraging automated evidence, not human MOS, a standard Seed-TTS
+benchmark, proof of perceptual identity, or a multi-speaker evaluation. The
+audio was intentionally not published because consent for public
+redistribution of the reference voice was not established.
+
+## Hardware and software
+
+The accepted release configuration is intentionally narrow:
+
+- Linux x86-64 and Python 3.12
+- NVIDIA RTX 50-series SM120/SM121 (tested on RTX 5060 Ti)
+- PyTorch 2.7.1+cu128 and CUDA 12.8 wheels
+- FlashInfer 0.6.14 with `nvidia-cutlass-dsl` 4.5.2
+- BF16 activations, FlashInfer CUDA graphs, user batch size 1
+
+There is no CPU, MLX, Torch-SDPA, older-GPU, or non-NVFP4 fallback because the
+BF16 projection sources are deliberately absent. RTX 5080 and SM121 are
+allowed by the loader but have not yet been benchmarked. This is custom code,
+not a Transformers `AutoModel` checkpoint or a standard Hub quantization
+format, so the hosted inference widget is disabled.
+
+`nvidia-cutlass-dsl` is installed under NVIDIA's separate software license.
+Review [THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md) before installation.
+
+## Install
+
+Clone the model repository, then use the included lock. Plain `pip install`
+does not honor the upstream Torch override or the official cu128 wheel source.
+
+```bash
+git clone https://huggingface.co/ajh-code/chatterbox-flash-t3-w4a16-nvfp4
+cd chatterbox-flash-t3-w4a16-nvfp4
+
+uv sync \
+  --project runtime/chatterbox-flash \
+  --locked \
+  --extra flashinfer \
+  --python 3.12
+```
+
+The lock resolves the exact validated critical stack. A dry verification is:
+
+```bash
+runtime/chatterbox-flash/.venv/bin/python - <<'PY'
+import importlib.metadata as md
+import flashinfer
+import torch
+
+assert torch.__version__ == "2.7.1+cu128"
+assert torch.version.cuda == "12.8"
+assert md.version("flashinfer-python") == "0.6.14"
+assert md.version("nvidia-cutlass-dsl") == "4.5.2"
+assert hasattr(flashinfer, "prepare_bf16_fp4_weights")
+assert hasattr(flashinfer, "mm_bf16_fp4")
+assert torch.cuda.get_device_capability() in {(12, 0), (12, 1)}
+print(torch.cuda.get_device_name(), "is ready")
+PY
+```
+
+## Zero-shot voice cloning
+
+Use only reference audio that you have permission to clone.
+
+```bash
+runtime/chatterbox-flash/.venv/bin/python example.py \
+  --prompt-wav your-reference.wav \
+  --text "This is Chatterbox Flash running native NVFP4 projections." \
+  --output output.wav \
+  --block-size 16
+```
+
+The example downloads only the unchanged S3Gen, voice-encoder, and tokenizer
+assets from the immutable base revision
+`4385507288b8197e6dab8b4e6b1603328d549d9d`. It verifies their sizes and
+SHA-256 fingerprints from `nvfp4_config.json`. It never downloads the original
+2.13 GB `t3_flash.safetensors`.
+
+The first decode for a new GPU/exact block shape invokes FlashInfer setup and
+tuning. On the RTX 5060 Ti, a fresh process with the selected D16 tactic already
+cached took 10.57 seconds for the first 1.28-second utterance versus 0.154
+seconds for the same-process warm repeat. A completely missing tactic has taken
+roughly 25 seconds in development. Autotune results are GPU-specific, so this
+repository does not ship a universal cache.
+
+## Profiles
+
+`D=16` is the quality-default profile. `D=24` is an explicitly optional fast
+profile: it is materially faster overall but had a slightly lower automated
+mean score and less consistent pairwise preference. `D=32` was tested and
+rejected after token-cap failures and an OOM in the natural suite; its cached
+prior remains serialized for research provenance, but the release loader
+refuses it.
+
+| Profile | Intended use | NVFP4 aggregate RTF | NVFP4 T3 tok/s | Automated mean |
+|---|---|---:|---:|---:|
+| D16 | Quality default | 0.08725 | 365.4 | 94.625 |
+| D24 | Optional fast | 0.07062 | 483.3 | 93.125 |
+
+## RTX 5060 Ti results
+
+The natural test is an eight-text, one-reference suite. Pipeline and T3 times
+are sums of per-case medians; aggregate RTF divides summed pipeline time by
+summed generated-audio duration. Memory is the maximum per-case median
+PyTorch CUDA peak, not total `nvidia-smi` process memory.
+
+| Profile | Comparison | RTF speedup | T3 throughput | Peak allocated reduction |
+|---|---|---:|---:|---:|
+| D16 | W4A16 vs matched BF16 | 1.0811x | 1.1131x | 18.82% / 1.209 GB |
+| D24 | W4A16 vs matched BF16 | 1.0309x | 1.0450x | 21.57% / 1.410 GB |
+
+A fixed capped-work control removes sampled-length attribution. D16 T3 is
+`1.1260x` faster and D24 T3 is `1.0656x` faster; each reduces fixed-work peak
+allocated memory by about 34%. D16 fixed whole-pipeline speedup is `1.0631x`.
+
+Do not quote the larger raw natural-suite wall-time ratios as pure NVFP4
+speedups: the BF16 and W4A16 models sampled different token counts. The
+D24-versus-D16 speed difference is primarily the larger block size, not FP4
+alone. Complete numeric records and caveats are in
+[benchmark_results.json](benchmark_results.json).
+
+Automated quality was close but is not proof of perceptual identity:
+
+- D16 BF16/NVFP4 means were 94.125/94.625; both transcribed 8/8 exactly with
+  zero hard failures. Order-controlled pairwise scoring resolved to seven ties
+  and one BF16 win.
+- D24 BF16/NVFP4 means were 95.0/93.125; both transcribed 8/8 exactly with zero
+  hard failures. Order control resolved to six ties, one BF16 win, and one
+  NVFP4 win.
+
+This was not human MOS, a standard Seed-TTS evaluation, or a multi-speaker
+release study.
+
+## Artifact and load behavior
+
+- `t3_nvfp4.safetensors`: 341,439,040 bytes, 447 tensors, SHA-256
+  `c5ea2d1d1441f657b383d8dad54bbae3778447d675822fb15918cb99b3587f1e`
+- 360 canonical NVFP4 tensors cover 30 layers × 4 fused projection kinds ×
+  FP4 codes/E4M3 scales/FP32 alpha.
+- 82 non-projection T3 tensors remain BF16; three fixed priors and two rotary
+  buffers support source-free construction.
+- Original T3 checkpoint: 2,129,657,904 bytes. The derivative T3 file is
+  83.97% smaller.
+- Derivative T3 plus the three unchanged downloaded base assets totals
+  1,412,035,402 bytes versus 3,200,254,266 bytes upstream: 55.88% less model
+  asset storage/download, excluding small runtime files and cache.
+
+In the final source-free release smoke, fresh-process pipeline load with base
+hash verification and warm JIT state took 2.675 seconds and settled at 1.444 GB
+PyTorch allocated / 1.478 GB reserved before the first generation. It prepared
+120 projections into 283,116,000 bytes, without materializing BF16 projection
+weights or calling the quantizer. The short generation peaked at 2.548 GB
+allocated. These are one-machine diagnostics; first-generation setup is
+reported separately above.
+
+The canonical D16 loader regression reproduced the accepted stripped runtime
+bit-for-bit: token history, waveform tensor SHA-256, and saved WAV SHA-256 all
+matched. A longer request also rebuilt the engine successfully while retaining
+the same source-free projection provider.
+
+## Limitations and responsible use
+
+- Only RTX 5060 Ti / SM120, batch one, and the documented settings received the
+  complete technical and small automated-audio gate.
+- D16 is the only quality-default profile. D24 is optional; D32 is rejected.
+- S3Gen and voice-encoder execution are not quantized by this artifact.
+- Prefix projections can remain slower than BF16 at irregular shapes; the
+  decode block projections provide the measured NVFP4 gain.
+- Streaming, multi-request throughput, RTX 5080 tuning, and human multi-voice
+  listening remain future work.
+- Generated length can change after quantization, affecting latency, memory,
+  pauses, and termination even when transcripts are correct.
+- Voice cloning can enable impersonation or deception. Obtain the speaker's
+  consent, disclose synthetic audio where appropriate, and comply with local
+  law and platform policies.
+
+No evaluation audio is included because the reference clip's redistribution
+and voice-consent status was not established for a public model repository.
+
+## Attribution and citation
+
+Chatterbox-Flash model weights and vendored runtime code are Copyright (c) 2026
+Resemble AI and licensed under MIT. This runtime is based on source commit
+`74e05baa8ce574bf2cc571702391a21f1b0d48c5`; the derivative modifies the
+FlashInfer projection path and adds the canonical NVFP4 loader. See
+[THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md).
+
+```bibtex
+@misc{seo2026chatterboxflashpriorcalibratedblockdiffusion,
+  title={Chatterbox-Flash: Prior-Calibrated Block Diffusion for Streaming Zero-Shot TTS},
+  author={Deokjin Seo and Gangin Park and Kihyun Nam},
+  year={2026},
+  eprint={2605.30748},
+  archivePrefix={arXiv},
+  primaryClass={cs.SD},
+  url={https://arxiv.org/abs/2605.30748}
+}
+```
+
+## Repository files
+
+- `t3_nvfp4.safetensors`: canonical source-free checkpoint
+- `nvfp4_config.json`: complete schema, fingerprints, revisions, and profiles
+- `runtime/chatterbox-flash/`: pinned installable derivative runtime and lock
+- `example.py`: tested local-bundle zero-shot generation entry point
+- `benchmark_results.json`: machine-readable release measurements
+- `SHA256SUMS`: release integrity manifest

SHA256SUMS

@@ -0,0 +1,31 @@
+d5fadc36773f0a6e03ae542aac79a3747abd47f77028ab304cb878dc2beb3ec1  ./.gitattributes
+672ba8b147a82d1ecc3f9f1dc899d69986efb2f7e4fd9f939ce9f38948922b41  ./.gitignore
+2eba8e98035070234c35a91bb6e3c4185505cf073ff2879c8ce5fc1a462f5534  ./LICENSE
+c3c8b178f6fca0debb7e7072c6fce42376f922caee993d50c9fb36285ed60aea  ./README.md
+6a8b7141a4bb5ac41462c6e4c0547bd6ca701aef28fa5c4453d191658193a485  ./THIRD_PARTY_NOTICES.md
+aaf7a6043f9de80a4a9ad99082274310d2dbfa9d6024753dc106755327bc0c75  ./benchmark_results.json
+ff1be650be410c64a7e7d9bc2288c40638aed4af7383cbc560e77c4e87de38be  ./example.py
+7521540282c2b716c3135a768b810b8f39c01ab5f773b0f8c82c9f9d4f8c44b5  ./nvfp4_config.json
+d6baef68ab47dd53741421f0440d62fc5632e7955c33d9c08b7b33a6c164b8c2  ./runtime/chatterbox-flash/.gitignore
+a80ea1b21580d334f010acf2f1c055d57d0959b1e6f578cc2b0a3de2eb96a8bb  ./runtime/chatterbox-flash/LICENSE
+d41a6f6cdedb1284e64119630d61aba97bbe325f9a0a6e0c6cbeb8ba9057bcb4  ./runtime/chatterbox-flash/README.md
+d1047075186b8f28646977a5e70e70f37a00abbd6d7ff5d77aa56a1cfd8a0697  ./runtime/chatterbox-flash/chatterbox_flash/__init__.py
+d3676e229c668a2d12e9a08d8abf9990203f6f6664f2fb97d3bfa75494549a75  ./runtime/chatterbox-flash/chatterbox_flash/attention_flashinfer.py
+6c7769bb30d846985751513a97be0d962ed378a82dbef124beb2d7b422a38991  ./runtime/chatterbox-flash/chatterbox_flash/calibration.py
+e8703a7bf949f6e8d119962fcf84775a3b62169cd0d77c4e99c0a2cc85089858  ./runtime/chatterbox-flash/chatterbox_flash/cfg_guidance.py
+f77672b56d510ef9d78f03d64ff1e192ef341b57f114f94d907bf8d49aad5cfe  ./runtime/chatterbox-flash/chatterbox_flash/engines/__init__.py
+0db90fd941bde83a834e5a3600806cfa8cab1fdb36c372de0dbfe95fe1f8dbd2  ./runtime/chatterbox-flash/chatterbox_flash/engines/base.py
+177d1d750b56e4922503c4b20b8ffb7bb04965effafccdcf52a0af56cc4a79f4  ./runtime/chatterbox-flash/chatterbox_flash/engines/flashinfer.py
+7ae970bfb9342a521bdf83576c736f6c8da5fd6324f5d54a13e053c20e4ddcd6  ./runtime/chatterbox-flash/chatterbox_flash/engines/mlx.py
+b03799ffd9aad68eb33700fcfdeea490e999c884092687b9a16d877a334ec3ca  ./runtime/chatterbox-flash/chatterbox_flash/engines/torch_sdpa.py
+a6e51fb3237d6b5569749fdce66bea1056bdccdc4245ccc28bbeb5646ff314c3  ./runtime/chatterbox-flash/chatterbox_flash/model.py
+276d1bc84f5c6c911fbe2b7b52097c015ae4e838bee2759acfc9ddebf4d68560  ./runtime/chatterbox-flash/chatterbox_flash/nvfp4.py
+561e11514455dbab23bfcff7b7e9d8252c8aeedd0b36e1e906f585e5981220fa  ./runtime/chatterbox-flash/chatterbox_flash/text_norm/__init__.py
+4b221a4829756695224bc62b2b4f0c1993f40311f21f299ca5d98e14feb6dee3  ./runtime/chatterbox-flash/chatterbox_flash/text_norm/cleaners.py
+1ffc6ea6d4649ee7498d1e3e184442b660811c2765cdd250d7e6c7021fa595ae  ./runtime/chatterbox-flash/chatterbox_flash/text_norm/number_norm.py
+26e26739d80282692aeecc9c44c7c0c83a116e1dcfedbddb370e066094e56d24  ./runtime/chatterbox-flash/chatterbox_flash/text_norm/phone_norm.py
+a4254d1b82781a79c18902d5151ee35e3c871926a72e87cb5b1562843b5bb055  ./runtime/chatterbox-flash/chatterbox_flash/text_norm/time_norm.py
+0c0665d0c1d9593422eb2c1bcb02f1ed2ce6bf6aed45729e10311525f985b57d  ./runtime/chatterbox-flash/chatterbox_flash/tts.py
+3ed7cf63c7fb04f46b453b2b60a2a890d1c64eac7945dabeea19aef6b40585d1  ./runtime/chatterbox-flash/pyproject.toml
+763867108090620985de8aa8f3019f0eca356ffe37519b1d45ae56e1fdcc78b4  ./runtime/chatterbox-flash/uv.lock
+c5ea2d1d1441f657b383d8dad54bbae3778447d675822fb15918cb99b3587f1e  ./t3_nvfp4.safetensors

THIRD_PARTY_NOTICES.md

@@ -0,0 +1,39 @@
+# Third-party notices
+
+This bundle is an unofficial community derivative. It is not affiliated with
+or endorsed by Resemble AI or NVIDIA.
+
+## Chatterbox-Flash
+
+The model weights and vendored runtime are derived from
+[ResembleAI/chatterbox-flash](https://github.com/resemble-ai/chatterbox-flash),
+source commit `74e05baa8ce574bf2cc571702391a21f1b0d48c5`.
+
+Copyright (c) 2026 Resemble AI. Licensed under the MIT License. The complete
+license text is reproduced in [LICENSE](LICENSE) and in
+`runtime/chatterbox-flash/LICENSE`.
+
+The derivative changes the FlashInfer engine/provider integration and adds a
+canonical, source-free NVFP4 loader. The unchanged S3Gen, voice-encoder, and
+tokenizer assets are not redistributed; the loader fetches them from the
+immutable Chatterbox-Flash Hub revision recorded in `nvfp4_config.json`.
+
+## Chatterbox-TTS
+
+Chatterbox-Flash uses architecture and runtime components from
+[ResembleAI/chatterbox](https://github.com/resemble-ai/chatterbox).
+
+Copyright (c) 2025 Resemble AI. Licensed under the MIT License.
+
+## Runtime dependencies
+
+[FlashInfer](https://github.com/flashinfer-ai/flashinfer) is an Apache-2.0
+runtime dependency and is not vendored.
+
+`nvidia-cutlass-dsl` and `nvidia-cutlass-dsl-libs-base` are installed from
+their official Python packages and are not vendored. They are distributed
+under the NVIDIA Software License Agreement included in those packages. Users
+must review and accept that agreement when installing or using the dependency.
+
+NVIDIA, CUDA, and Blackwell are trademarks and do not imply endorsement.
+

benchmark_results.json

@@ -0,0 +1,142 @@
+{
+  "schema_version": 1,
+  "date": "2026-07-12",
+  "hardware": {
+    "gpu": "NVIDIA GeForce RTX 5060 Ti",
+    "compute_capability": "SM120",
+    "user_batch_size": 1
+  },
+  "software": {
+    "python": "3.12",
+    "torch": "2.7.1+cu128",
+    "torch_cuda": "12.8",
+    "flashinfer": "0.6.14",
+    "nvidia_cutlass_dsl": "4.5.2"
+  },
+  "natural_eight_text_suite": {
+    "definition": "Pipeline and T3 seconds are sums of per-case medians; aggregate RTF is summed pipeline time divided by summed generated-audio duration; T3 throughput uses valid speech tokens; memory is maximum per-case median PyTorch CUDA peak.",
+    "d16_quality": {
+      "bf16": {
+        "audio_seconds": 98.64,
+        "pipeline_seconds": 9.304738,
+        "aggregate_rtf": 0.0943303,
+        "t3_seconds": 7.512483,
+        "valid_tokens": 2466,
+        "t3_tokens_per_second": 328.254,
+        "peak_allocated_bytes": 6423756800,
+        "peak_reserved_bytes": 7107248128
+      },
+      "nvfp4_w4a16": {
+        "audio_seconds": 92.52,
+        "pipeline_seconds": 8.072466,
+        "aggregate_rtf": 0.087251,
+        "t3_seconds": 6.330483,
+        "valid_tokens": 2313,
+        "t3_tokens_per_second": 365.375,
+        "peak_allocated_bytes": 5214659584,
+        "peak_reserved_bytes": 6322913280
+      },
+      "nvfp4_vs_bf16": {
+        "aggregate_rtf_speedup": 1.081136,
+        "t3_throughput_speedup": 1.113087,
+        "peak_allocated_reduction_bytes": 1209097216,
+        "peak_allocated_reduction_percent": 18.8223,
+        "peak_reserved_reduction_bytes": 784334848,
+        "peak_reserved_reduction_percent": 11.0357
+      }
+    },
+    "d24_fast": {
+      "bf16": {
+        "audio_seconds": 93.36,
+        "pipeline_seconds": 6.797084,
+        "aggregate_rtf": 0.0728051,
+        "t3_seconds": 5.04703,
+        "valid_tokens": 2334,
+        "t3_tokens_per_second": 462.45,
+        "peak_allocated_bytes": 6539530752,
+        "peak_reserved_bytes": 7411335168
+      },
+      "nvfp4_w4a16": {
+        "audio_seconds": 91.8,
+        "pipeline_seconds": 6.483264,
+        "aggregate_rtf": 0.0706238,
+        "t3_seconds": 4.748782,
+        "valid_tokens": 2295,
+        "t3_tokens_per_second": 483.282,
+        "peak_allocated_bytes": 5129031168,
+        "peak_reserved_bytes": 6064963584
+      },
+      "nvfp4_vs_bf16": {
+        "aggregate_rtf_speedup": 1.030886,
+        "t3_throughput_speedup": 1.045046,
+        "peak_allocated_reduction_bytes": 1410499584,
+        "peak_allocated_reduction_percent": 21.5688,
+        "peak_reserved_reduction_bytes": 1346371584,
+        "peak_reserved_reduction_percent": 18.1664
+      }
+    }
+  },
+  "fixed_capped_work": {
+    "definition": "One fixed-seed case, five-repeat medians, 1024-token reservation. D16 caps at 32 speech tokens; D24 caps at 48, so compare quantized and BF16 only within each block size.",
+    "d16": {
+      "bf16_t3_ms": 88.363,
+      "nvfp4_t3_ms": 78.475,
+      "t3_speedup": 1.125998,
+      "bf16_pipeline_ms": 162.897,
+      "nvfp4_pipeline_ms": 153.235,
+      "pipeline_speedup": 1.063052,
+      "bf16_peak_allocated_bytes": 4165647360,
+      "nvfp4_peak_allocated_bytes": 2745680896,
+      "peak_allocated_reduction_percent": 34.0875,
+      "peak_reserved_reduction_percent": 30.0444
+    },
+    "d24": {
+      "bf16_t3_ms": 91.192,
+      "nvfp4_t3_ms": 85.574,
+      "t3_speedup": 1.065645,
+      "bf16_pipeline_ms": 180.938,
+      "nvfp4_pipeline_ms": 176.006,
+      "pipeline_speedup": 1.028023,
+      "bf16_peak_allocated_bytes": 4173167104,
+      "nvfp4_peak_allocated_bytes": 2754244096,
+      "peak_allocated_reduction_percent": 34.0011,
+      "peak_reserved_reduction_percent": 30.0689
+    }
+  },
+  "automated_quality": {
+    "warning": "Small eight-text, one-reference automated evaluation; not human MOS or a multi-speaker benchmark.",
+    "d16": {
+      "bf16_mean_score": 94.125,
+      "nvfp4_mean_score": 94.625,
+      "exact_transcripts_each": "8/8",
+      "hard_failures_each": "0/8",
+      "order_controlled_pairwise": "7 ties, 1 BF16 win"
+    },
+    "d24": {
+      "bf16_mean_score": 95.0,
+      "nvfp4_mean_score": 93.125,
+      "exact_transcripts_each": "8/8",
+      "hard_failures_each": "0/8",
+      "order_controlled_pairwise": "6 ties, 1 BF16 win, 1 NVFP4 win"
+    }
+  },
+  "source_free_load_diagnostic": {
+    "fresh_process_load_seconds_with_base_hash_verification_and_jit_cache_hit": 2.675073,
+    "allocated_bytes_after_load": 1444340224,
+    "reserved_bytes_after_load": 1478492160,
+    "prepared_projection_bytes": 283116000,
+    "prepared_projection_count": 120,
+    "short_generation_peak_allocated_bytes": 2548431360,
+    "first_same_process_generation_seconds_with_exact_shape_config_cache": 10.567298,
+    "second_same_process_generation_seconds": 0.153511,
+    "first_to_warm_ratio": 68.8376,
+    "notes": "Canonical weights were prepared without BF16 projection materialization or nvfp4_quantize. The first-generation measurement includes engine creation, CUDA graph capture, and per-process FlashInfer setup even though the exact shape's selected config was cached. A completely missing exact-shape tuning entry can take longer. Caches are GPU-specific and are intentionally not distributed."
+  },
+  "interpretation_warnings": [
+    "Do not present raw natural-suite wall-time ratios as pure NVFP4 speedups because BF16 and NVFP4 sampled different token counts.",
+    "PyTorch allocated/reserved peaks are not total board VRAM or nvidia-smi process usage.",
+    "The D24 versus D16 speed difference is primarily the block-size change, not quantization alone.",
+    "The D16 BF16 natural control had one token-cap hit and one duration warning; NVFP4 did not.",
+    "D32 was rejected for release use after cap hits and an OOM in the natural suite."
+  ]
+}

example.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""Generate one zero-shot voice-cloned utterance with the local NVFP4 bundle."""
+
+from __future__ import annotations
+
+import argparse
+import re
+from pathlib import Path
+
+import soundfile as sf
+import torch
+from huggingface_hub import snapshot_download
+
+from chatterbox_flash.nvfp4 import (
+    BASE_ASSET_FILENAMES,
+    BASE_CHECKPOINT_REVISION,
+    BASE_REPO_ID,
+    load_prepacked_tts_local,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--prompt-wav", type=Path, required=True)
+    parser.add_argument("--text", required=True)
+    parser.add_argument("--output", type=Path, default=Path("output.wav"))
+    parser.add_argument(
+        "--block-size",
+        type=int,
+        choices=(16, 24),
+        default=16,
+        help="16 is the quality profile; 24 is the optional faster profile",
+    )
+    parser.add_argument("--seed", type=int, default=4242)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if not args.prompt_wav.is_file():
+        raise SystemExit(f"Missing reference audio: {args.prompt_wav}")
+
+    bundle = Path(__file__).resolve().parent
+    base_dir = Path(snapshot_download(
+        repo_id=BASE_REPO_ID,
+        revision=BASE_CHECKPOINT_REVISION,
+        allow_patterns=list(BASE_ASSET_FILENAMES),
+    ))
+    capability = torch.cuda.get_device_capability()
+    gpu_slug = re.sub(r"[^a-zA-Z0-9_.-]+", "-", torch.cuda.get_device_name())
+    autotune_cache = (
+        Path.home()
+        / ".cache"
+        / "chatterbox-flash-nvfp4"
+        / f"flashinfer-{gpu_slug}-sm{capability[0]}{capability[1]}-d{args.block_size}.json"
+    )
+
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed_all(args.seed)
+    tts = load_prepacked_tts_local(
+        bundle,
+        base_dir,
+        device="cuda",
+        drf_block_size=args.block_size,
+        autotune_cache=autotune_cache,
+    )
+    waveform = tts.generate(
+        args.text,
+        audio_prompt_path=args.prompt_wav,
+        exaggeration=0.5,
+        num_steps=10,
+        temperature=0.2,
+        time_shift_tau=0.5,
+        omnivoice_schedule_t_shift=0.5,
+        cfg_scale=1.0,
+        position_temperature=5.0,
+        n_cfm_timesteps=2,
+        backend="flashinfer",
+        use_cuda_graph=True,
+    )
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    sf.write(args.output, waveform.float().numpy(), tts.sr)
+    print(f"Wrote {args.output} at {tts.sr} Hz")
+
+
+if __name__ == "__main__":
+    main()
+

runtime/chatterbox-flash/.gitignore

@@ -0,0 +1,30 @@
+# Python
+__pycache__/
+*.py[cod]
+*.egg-info/
+build/
+dist/
+.eggs/
+
+# Virtual environments (uv venv, etc.)
+.venv/
+.venv*/
+venv/
+env/
+
+# Model checkpoints — released via the Hugging Face Hub, too large for git
+checkpoints/
+
+# Generated audio / test outputs
+/test_out/
+/syn_out/
+/out.wav
+/*.wav
+eval_out_smoke/
+
+# Internal-only scripts (not part of the released package)
+smoke_test.py
+upload_hf.py
+scripts/convert_checkpoint.py
+# Demo-page sites are separate git repos (own GitHub remotes), not part of this package
+*_demopage/

runtime/chatterbox-flash/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Resemble AI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

runtime/chatterbox-flash/README.md

@@ -0,0 +1,11 @@
+# Chatterbox-Flash NVFP4 runtime
+
+This is the installable runtime shipped with the unofficial
+Chatterbox-Flash T3 W4A16 NVFP4 model bundle. It vendors the MIT-licensed
+Chatterbox-Flash source at commit
+`74e05baa8ce574bf2cc571702391a21f1b0d48c5` and adds an engine-local,
+source-free NVFP4 projection provider plus canonical checkpoint loader.
+
+Install only through the bundle's locked `uv sync --extra flashinfer` command.
+See the model card at the bundle root for hardware requirements, usage,
+benchmarks, attribution, and limitations.

runtime/chatterbox-flash/chatterbox_flash/__init__.py

@@ -0,0 +1,36 @@
+"""Chatterbox-Flash — prior-calibrated block-diffusion zero-shot TTS.
+
+Public API:
+    ChatterboxFlashTTS     — end-to-end TTS pipeline (T3 + S3Gen + VE + tokenizer)
+    ChatterboxFlashT3      — block-diffusion T3 decoder (low-level)
+    Conditionals           — cacheable speaker / prompt conditioning bundle
+    en_us_cleaner          — English text normalization used at training / eval time
+"""
+
+from .engines import FLASHINFER_AVAILABLE, MLX_AVAILABLE, build_engine
+from .model import ChatterboxFlashT3
+from .nvfp4 import (
+    CanonicalNVFP4ProjectionProvider,
+    load_prepacked_t3,
+    load_prepacked_tts_from_hub,
+    load_prepacked_tts_local,
+)
+from .text_norm import en_us_cleaner
+from .tts import ChatterboxFlashTTS, Conditionals
+
+__version__ = "0.1.0+nvfp4.1"
+
+__all__ = [
+    "ChatterboxFlashT3",
+    "ChatterboxFlashTTS",
+    "CanonicalNVFP4ProjectionProvider",
+    "Conditionals",
+    "FLASHINFER_AVAILABLE",
+    "MLX_AVAILABLE",
+    "build_engine",
+    "en_us_cleaner",
+    "load_prepacked_t3",
+    "load_prepacked_tts_from_hub",
+    "load_prepacked_tts_local",
+    "__version__",
+]

runtime/chatterbox-flash/chatterbox_flash/attention_flashinfer.py

@@ -0,0 +1,20 @@
+"""Backwards-compatibility shim — :class:`FlashInferEngine` now lives in
+:mod:`chatterbox_flash.engines.flashinfer`.
+
+The original module path is preserved so existing scripts that did
+
+    from chatterbox_flash.attention_flashinfer import FlashInferEngine
+
+continue to work.  New code should import from
+``chatterbox_flash.engines`` (or call :func:`chatterbox_flash.engines.build_engine`)
+instead.
+"""
+
+from __future__ import annotations
+
+from .engines.flashinfer import (  # noqa: F401
+    FlashInferEngine,
+    flashinfer_available,
+)
+
+__all__ = ["FlashInferEngine", "flashinfer_available"]

runtime/chatterbox-flash/chatterbox_flash/calibration.py

@@ -0,0 +1,82 @@
+"""Prior-calibrated PMI scoring for Chatterbox-Flash block decoding.
+
+The decoder ranks masked positions by a self-calibrated pointwise mutual
+information score that suppresses the dominant-token bias of discrete speech
+codecs (no per-token silence IDs needed):
+
+    s_i = log p_i(x_i) - lambda * log p_bar(x_i)
+
+where ``p_bar`` is a reference marginal — either the precomputed unconditional
+block prior (see :meth:`ChatterboxFlashT3._compute_uncond_block_prior`) or the
+step-0 block-mean ``mean_j p_j``.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor
+
+_LOG_EPS = 1e-8
+
+
+def _marginal_for_sampled(
+    probs_bl: Tensor,
+    sampled: Tensor,
+    marginal_prior: Tensor | None,
+) -> Tensor:
+    """Return per-position marginal probability for each sampled token.
+
+    If ``marginal_prior`` is given (shape ``(V,)``), index into it directly.
+    Otherwise compute the block-mean marginal on the fly.
+    """
+    if marginal_prior is not None:
+        return marginal_prior[sampled]
+    bl = probs_bl.size(0)
+    sampled_cols = sampled.unsqueeze(0).expand(bl, -1)
+    return probs_bl.gather(1, sampled_cols).mean(dim=0)
+
+
+def self_calibrated_pmi(
+    probs_bl: Tensor,
+    sampled: Tensor,
+    eps: float = _LOG_EPS,
+    marginal_prior: Tensor | None = None,
+    prior_lambda: float = 1.0,
+) -> Tensor:
+    """
+    Per-position PMI-style score vs a marginal prior.
+
+    s_i = log p_i(x_i) - lambda * log p_bar(x_i),
+    p_bar(v) = mean_j p_j(v).
+
+    When ``marginal_prior`` (shape ``(V,)``) is provided, it is used as
+    ``p_bar`` instead of the current block-mean — typically the precomputed
+    unconditional block prior for a consistent, uninformed reference.
+    """
+    pi = probs_bl.gather(1, sampled.unsqueeze(1)).squeeze(1)
+    p_bar_sampled = _marginal_for_sampled(probs_bl, sampled, marginal_prior)
+    return torch.log(pi.clamp_min(eps)) - (
+        float(prior_lambda) * torch.log(p_bar_sampled.clamp_min(eps))
+    )
+
+
+def compute_pmi(
+    probs_bl: Tensor,
+    sampled: Tensor,
+    *,
+    marginal_prior: Tensor | None = None,
+    prior_lambda: float = 1.0,
+    eps: float = _LOG_EPS,
+) -> Tensor:
+    """Compute self-calibrated PMI.
+
+    When ``marginal_prior`` (shape ``(V,)``) is given, it is used as the
+    denominator instead of computing ``probs_bl.mean(dim=0)`` on the fly.
+    """
+    return self_calibrated_pmi(
+        probs_bl,
+        sampled,
+        eps=eps,
+        marginal_prior=marginal_prior,
+        prior_lambda=prior_lambda,
+    )

runtime/chatterbox-flash/chatterbox_flash/cfg_guidance.py

@@ -0,0 +1,57 @@
+"""Classifier-free guidance helpers for Chatterbox-Flash.
+
+We only keep the ``zero_text_batch`` recipe used in the paper: at each block
+step the forward batch is doubled to ``2 * B``; rows ``0..B-1`` carry the full
+text + conditioning, rows ``B..2B-1`` carry a null branch where the text and
+the entire conditioning prefix are zeroed (``cfg_null_cond_mode="zero_all"``).
+The two halves of the logits are then combined via
+``log p' = log p_c + w * (log p_c - log p_u)``.
+
+The same combination is applied on the **PMI scale** as well (see
+:func:`pmi_cfg_combine`) — this is the only PMI/CFG combination supported;
+no alternative ``cfg_prior_mode`` is exposed.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+
+def omnivoice_cfg_log_probs(
+    c_log_probs: Tensor,
+    u_log_probs: Tensor,
+    guidance_scale: float,
+) -> Tensor:
+    """``log p' = log p_c + w * (log p_c - log p_u)`` (elementwise)."""
+    if guidance_scale == 0.0:
+        return c_log_probs
+    return c_log_probs + guidance_scale * (c_log_probs - u_log_probs)
+
+
+def apply_zero_text_cfg_from_logits(
+    logits_cond: Tensor,
+    logits_u: Tensor,
+    guidance_scale: float,
+) -> Tensor:
+    """CFG combination from two logit tensors of identical shape.
+
+    Operates on log-softmax so the result is a valid scoring/sampling logit
+    (softmax recovers the guided distribution).
+    """
+    if guidance_scale == 0.0:
+        return logits_cond
+    c_log = F.log_softmax(logits_cond, dim=-1)
+    u_log = F.log_softmax(logits_u, dim=-1)
+    return omnivoice_cfg_log_probs(c_log, u_log, guidance_scale)
+
+
+def pmi_cfg_combine(
+    pmi_c: Tensor,
+    pmi_u: Tensor,
+    guidance_scale: float,
+) -> Tensor:
+    """Guidance-combine PMI scores: ``pmi = (1+w) * pmi_c - w * pmi_u``."""
+    w = float(guidance_scale)
+    return (1.0 + w) * pmi_c - w * pmi_u

runtime/chatterbox-flash/chatterbox_flash/engines/__init__.py

@@ -0,0 +1,168 @@
+"""Inference engines for Chatterbox-Flash.
+
+Three backends are provided, all implementing the :class:`InferenceEngine`
+protocol consumed by :class:`~chatterbox_flash.model.ChatterboxFlashT3`:
+
+* :class:`FlashInferEngine` — paged KV cache + custom attention kernels
+  via the ``flashinfer-python`` package; supports CUDA-graph capture of
+  the block forward.  This is what the paper numbers were produced with
+  and what should be used for any throughput-sensitive workload on
+  CUDA.
+
+* :class:`TorchSDPAEngine` — pure PyTorch fallback using the HuggingFace
+  transformers SDPA backend with a :class:`DynamicCache`.  No extra
+  dependencies, runs on any CUDA / CPU / MPS device the underlying T3
+  supports.  No CUDA graphs, slower per step, but produces numerically
+  matching output (modulo small fp32/bf16 reduction order differences).
+
+* :class:`MLXEngine` *(experimental)* — Apple-Silicon native Metal
+  backend.  Wraps the mlx-lm LLaMA backbone and bypasses its
+  ``embed_tokens`` layer so the chatterbox-flash custom embeddings can
+  be used unchanged.  Requires ``mlx`` and ``mlx-lm`` installed.
+
+Use :func:`build_engine` to instantiate; it auto-falls-back to SDPA when
+FlashInfer is unavailable.  MLX must be selected explicitly via
+``backend="mlx"`` (or the ``CHATTERBOX_FLASH_ENGINE=mlx`` env var).
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import logging
+import os
+from typing import Literal
+
+import torch
+
+from .base import InferenceEngine
+from .torch_sdpa import TorchSDPAEngine
+
+logger = logging.getLogger(__name__)
+
+
+FLASHINFER_AVAILABLE: bool = importlib.util.find_spec("flashinfer") is not None
+MLX_AVAILABLE: bool = (
+    importlib.util.find_spec("mlx") is not None
+    and importlib.util.find_spec("mlx_lm") is not None
+)
+
+
+_Backend = Literal["auto", "flashinfer", "torch", "mlx"]
+
+
+def _resolve_backend(requested: _Backend) -> Literal["flashinfer", "torch", "mlx"]:
+    """Pick a backend honoring ``CHATTERBOX_FLASH_ENGINE`` env override."""
+    env = os.environ.get("CHATTERBOX_FLASH_ENGINE", "").strip().lower()
+    if env in ("flashinfer", "torch", "sdpa", "mlx"):
+        requested = "torch" if env == "sdpa" else env  # type: ignore[assignment]
+
+    if requested == "auto":
+        return "flashinfer" if FLASHINFER_AVAILABLE else "torch"
+    if requested == "flashinfer":
+        if not FLASHINFER_AVAILABLE:
+            raise ImportError(
+                "flashinfer-python is not installed. Install with "
+                "`pip install flashinfer-python` or pass backend='torch'.",
+            )
+        return "flashinfer"
+    if requested == "mlx":
+        if not MLX_AVAILABLE:
+            raise ImportError(
+                "MLX backend requires the 'mlx' and 'mlx-lm' packages. "
+                "Install with `pip install mlx mlx-lm` "
+                "(Apple Silicon Mac only).",
+            )
+        return "mlx"
+    if requested in ("torch", "sdpa"):
+        return "torch"
+    raise ValueError(
+        f"backend must be one of 'auto', 'flashinfer', 'torch', 'mlx'; "
+        f"got {requested!r}",
+    )
+
+
+def build_engine(
+    model: torch.nn.Module,
+    max_seq_len: int,
+    dtype: torch.dtype,
+    *,
+    backend: _Backend = "auto",
+    batch_size: int = 1,
+    page_size: int = 16,
+) -> InferenceEngine:
+    """Construct an inference engine.
+
+    Parameters
+    ----------
+    model: the :class:`ChatterboxFlashT3` instance.
+    max_seq_len: upper bound on ``prefix_len + total_speech_len`` for the
+        cache reservation; needed by the FlashInfer paged buffer.  The
+        torch SDPA / MLX engines use this only as a hint.
+    dtype: model parameter dtype (bf16 / fp16 / fp32).
+    backend: ``"auto"`` (default, prefer FlashInfer) / ``"flashinfer"`` /
+        ``"torch"`` / ``"mlx"``.  Environment variable
+        ``CHATTERBOX_FLASH_ENGINE`` overrides this.
+    batch_size: number of forward-batch rows (with zero-text-batch CFG
+        this is ``2 × B_usr``).
+    page_size: FlashInfer paged-cache page size (ignored elsewhere).
+    """
+    chosen = _resolve_backend(backend)
+    projection_provider = getattr(model, "_flashinfer_projection_provider", None)
+    if (
+        projection_provider is not None
+        and bool(getattr(projection_provider, "source_free", False))
+        and chosen != "flashinfer"
+    ):
+        raise RuntimeError(
+            "This model has no torch/MLX transformer projection weights. "
+            "Its source-free NVFP4 projection provider requires "
+            "backend='flashinfer'."
+        )
+    if chosen == "flashinfer":
+        from .flashinfer import FlashInferEngine
+        logger.debug(
+            "build_engine: backend=flashinfer max_seq=%d dtype=%s batch=%d page=%d",
+            max_seq_len, dtype, batch_size, page_size,
+        )
+        return FlashInferEngine(
+            model, max_seq_len, dtype,
+            batch_size=batch_size, page_size=page_size,
+        )
+    if chosen == "mlx":
+        from .mlx import MLXEngine
+        logger.debug(
+            "build_engine: backend=mlx max_seq=%d dtype=%s batch=%d",
+            max_seq_len, dtype, batch_size,
+        )
+        return MLXEngine(
+            model, max_seq_len, dtype, batch_size=batch_size,
+        )
+    logger.debug(
+        "build_engine: backend=torch (SDPA fallback) max_seq=%d dtype=%s batch=%d",
+        max_seq_len, dtype, batch_size,
+    )
+    return TorchSDPAEngine(
+        model, max_seq_len, dtype, batch_size=batch_size,
+    )
+
+
+def __getattr__(name: str):  # noqa: D401
+    """Lazy attribute access — only import optional backends if asked for."""
+    if name == "FlashInferEngine":
+        from .flashinfer import FlashInferEngine
+        return FlashInferEngine
+    if name == "MLXEngine":
+        from .mlx import MLXEngine
+        return MLXEngine
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "InferenceEngine",
+    "TorchSDPAEngine",
+    "FlashInferEngine",
+    "MLXEngine",
+    "FLASHINFER_AVAILABLE",
+    "MLX_AVAILABLE",
+    "build_engine",
+]

runtime/chatterbox-flash/chatterbox_flash/engines/base.py

@@ -0,0 +1,141 @@
+"""Abstract inference-engine protocol shared by the FlashInfer and the
+torch-SDPA backends.
+
+The block-diffusion ``generate()`` loop is engine-agnostic: it only ever
+talks to the engine through these six methods.  Concrete backends differ
+in how they store the KV cache (paged buffer vs. HuggingFace
+:class:`~transformers.DynamicCache`), how they batch the prefix
+forward, and whether they can capture the per-block forward as a CUDA
+graph.
+
+Contract
+--------
+The engine manages a single KV cache shared by the conditioning prefix
+and the speech-block forwards.  Lifecycle for one ``generate()`` call::
+
+    engine.reset()
+    shift_ctx = engine.prefix_forward(prefix_emb_or_list)   # (B, 1, dim)
+    engine.capture_cuda_graph(BS, speech_head=speech_head)   # optional
+
+    for b in range(num_blocks):
+        for k in range(K):
+            engine.set_shift_ctx(shift_ctx)                  # graph engines only
+            hidden, logits = engine.block_forward_graph(block_emb, cs)
+            #   or: hidden = engine.block_forward(block_emb, cs)
+
+        engine.advance_cache(block_len)
+        # caller commits the block's clean tokens via block_forward(...)
+        # before advance_cache; shift_ctx is updated from that hidden.
+"""
+
+from __future__ import annotations
+
+from typing import List, Optional, Protocol, Tuple, Union, runtime_checkable
+
+import torch
+from torch import Tensor
+
+
+CacheStarts = Union[int, List[int]]
+
+
+@runtime_checkable
+class InferenceEngine(Protocol):
+    """Minimal protocol implemented by all Chatterbox-Flash engines."""
+
+    has_cuda_graph: bool
+    """``True`` iff :meth:`capture_cuda_graph` was called and succeeded."""
+
+    @property
+    def cache_len(self) -> int:
+        """Current per-row KV cache occupancy (uniform across the batch)."""
+        ...
+
+    def reset(self) -> None:
+        """Clear cache contents and free CUDA-graph state for the next call."""
+        ...
+
+    def can_reuse(
+        self,
+        max_seq_len: int,
+        dtype: torch.dtype,
+        *,
+        batch_size: int,
+        page_size: int = 16,
+    ) -> bool:
+        """Whether this engine can be re-used for a request with these shape/dtype."""
+        ...
+
+    # --- prefix / block forwards -------------------------------------- #
+
+    def prefix_forward(
+        self,
+        prefix_emb_or_list: Union[Tensor, List[Tensor]],
+    ) -> Tensor:
+        """Encode the conditioning prefix once.
+
+        ``prefix_emb_or_list`` is either a stacked tensor ``(B_fwd, lp, dim)``
+        (same-text path) or a list of length ``B_fwd`` of per-row prefixes
+        with different lengths (cross-text path).  Writes the prefix KV
+        into the engine's cache and returns ``shift_ctx`` — the hidden
+        state of the SOS row, shape ``(B_fwd, 1, dim)``.
+        """
+        ...
+
+    def block_forward(
+        self,
+        block_emb: Tensor,
+        cache_starts: CacheStarts,
+    ) -> Tensor:
+        """Eager (non-CUDA-graph) block forward.
+
+        ``block_emb`` shape: ``(B_fwd, bl, dim)``.
+        ``cache_starts`` is either an int (uniform start position) or a
+        per-row list (cross-text / compact-null-prefix path).
+        Returns hidden state ``(B_fwd, bl, dim)``; logits are computed by
+        the caller via ``speech_head``.
+        """
+        ...
+
+    def block_forward_graph(
+        self,
+        block_emb: Tensor,
+        cache_starts: CacheStarts,
+    ) -> Tuple[Tensor, Optional[Tensor]]:
+        """CUDA-graph replay of the block forward.
+
+        Engines that don't support CUDA graphs (e.g. torch SDPA) should
+        return ``(self.block_forward(block_emb, cache_starts), None)`` so
+        the caller can fall back uniformly.  When the engine pre-computes
+        logits inside the graph (FlashInfer), the second element of the
+        return tuple is the logits tensor ``(B_fwd, bl, V)``.
+        """
+        ...
+
+    def advance_cache(self, block_len: int) -> None:
+        """Commit ``block_len`` newly-written entries to the cache."""
+        ...
+
+    def set_shift_ctx(self, ctx: Tensor) -> None:
+        """Stage the SOS / previous-block-end hidden state for the graph.
+
+        FlashInfer copies this into a graph-captured buffer at the start
+        of each block; the SDPA engine ignores it (it computes the shift
+        in-line during :meth:`block_forward`).
+        """
+        ...
+
+    def capture_cuda_graph(
+        self,
+        block_size: int,
+        *,
+        speech_head: Optional[torch.nn.Linear] = None,
+    ) -> None:
+        """Capture the per-block forward as a CUDA graph.
+
+        ``block_size`` is the fixed ``D`` the graph is captured for.
+        Engines that don't support CUDA graphs should leave
+        ``has_cuda_graph = False``.  Optionally folds the speech head
+        into the graph if supplied.
+        """
+        ...

runtime/chatterbox-flash/chatterbox_flash/engines/flashinfer.py

@@ -0,0 +1,1656 @@
+"""
+FlashInfer attention engine with **paged KV cache** for batched
+block-diffusion inference.
+
+Supports ``batch_size >= 1`` via ``BatchPrefillWithPagedKVCacheWrapper``.
+Each sample gets a static slice of pages in the pool; prefix pages are
+computed once and replicated for same-text multi-generation, or computed
+independently per sample for cross-text batching.
+
+KV cache layout (NHD, per layer):
+    ``(total_pages, 2, page_size, num_kv_heads, head_dim)``
+
+Static page allocation: sample *i* → pages
+``[i * pages_per_sample : (i+1) * pages_per_sample]``.
+
+Fused projections (same as before):
+  - QKV: 1 GEMM per layer instead of 3.
+  - Gate-up MLP: 1 GEMM instead of 2 (Llama SwiGLU).
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+from typing import List, Protocol, runtime_checkable
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+logger = logging.getLogger(__name__)
+
+try:
+    import flashinfer
+    _FLASHINFER_AVAILABLE = True
+    # Module-level alias avoids a name collision with the engine's
+    # ``self._fi_fused_add_rmsnorm`` helper, which forwards through to the
+    # real flashinfer kernel.
+    _flashinfer_fused_add_rmsnorm = flashinfer.norm.fused_add_rmsnorm
+    try:
+        from flashinfer.rope import apply_rope_inplace as _fi_apply_rope
+        _FI_ROPE_AVAILABLE = True
+    except (ImportError, AttributeError):
+        _FI_ROPE_AVAILABLE = False
+except ImportError:
+    _FLASHINFER_AVAILABLE = False
+    _FI_ROPE_AVAILABLE = False
+    _flashinfer_fused_add_rmsnorm = None  # type: ignore[assignment]
+
+
+def flashinfer_available() -> bool:
+    return _FLASHINFER_AVAILABLE
+
+
+@runtime_checkable
+class FlashInferProjectionProvider(Protocol):
+    """Model-attached source of engine-local transformer projections.
+
+    A provider may own backend-prepared quantized weights instead of PyTorch
+    ``Parameter`` objects.  Dispatch is deliberately confined to the four
+    transformer projection boundaries in :class:`FlashInferEngine`; global
+    ``torch.nn.functional.linear`` behavior is never modified.
+    """
+
+    source_free: bool
+
+    def validate_engine(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        num_layers: int,
+    ) -> None:
+        """Reject an incompatible engine before any forward executes."""
+
+    def prepare_for_engine(
+        self,
+        *,
+        batch_size: int,
+        block_size: int,
+    ) -> None:
+        """Optionally select/cache kernels for an engine shape."""
+
+    def project(
+        self,
+        kind: str,
+        layer_idx: int,
+        activation: Tensor,
+    ) -> Tensor:
+        """Apply one of ``qkv``, ``o``, ``gate_up``, or ``down``."""
+
+
+class FlashInferEngine:
+    """FlashInfer engine with paged KV cache for batched block-diffusion."""
+
+    # ------------------------------------------------------------------ #
+    #  Construction
+    # ------------------------------------------------------------------ #
+
+    def __init__(
+        self,
+        model,
+        max_seq_len: int,
+        dtype: torch.dtype,
+        batch_size: int = 1,
+        page_size: int = 16,
+    ):
+        if not _FLASHINFER_AVAILABLE:
+            raise ImportError(
+                "FlashInfer not installed.  Install with: "
+                "pip install flashinfer-python"
+            )
+
+        device = model.device
+        config = model.tfmr.config
+
+        # ---- Hopper PDL (programmatic dependent launch) ----
+        # PDL lets back-to-back kernels overlap their epilogue/prologue on
+        # Hopper, which helps the memory-bound norm / activation kernels in the
+        # decoder hot loop. Auto-detected from the device capability; can be
+        # forced via ``RESEMBLETRON_FLASHINFER_PDL`` (``0`` to disable).
+        import os
+        _pdl_env = os.environ.get("RESEMBLETRON_FLASHINFER_PDL")
+        if _pdl_env is not None:
+            self._enable_pdl: bool = _pdl_env not in ("0", "false", "False")
+        elif device.type == "cuda":
+            cap = torch.cuda.get_device_capability(device)
+            self._enable_pdl = cap[0] >= 9
+        else:
+            self._enable_pdl = False
+        # ---- fp16 qk reduction for bf16 attention ----
+        # Costs one extra cast on QK but lets the tensor cores reduce in fp16,
+        # which is materially faster on Hopper. Numerically equivalent in
+        # practice for inference-time bf16 models (vLLM/SGLang default).
+        self._use_fp16_qk_red: bool = dtype in (torch.float16, torch.bfloat16)
+
+        num_layers = config.num_hidden_layers
+        num_heads = config.num_attention_heads
+        num_kv_heads = config.num_key_value_heads
+        head_dim = getattr(config, "head_dim", config.hidden_size // num_heads)
+
+        self._projection_provider: FlashInferProjectionProvider | None = getattr(
+            model, "_flashinfer_projection_provider", None,
+        )
+        if self._projection_provider is not None:
+            missing_provider_methods = [
+                name for name in ("validate_engine", "prepare_for_engine", "project")
+                if not callable(getattr(self._projection_provider, name, None))
+            ]
+            if missing_provider_methods:
+                raise TypeError(
+                    "model._flashinfer_projection_provider is missing callable "
+                    f"methods: {missing_provider_methods}"
+                )
+            self._projection_provider.validate_engine(
+                device=device,
+                dtype=dtype,
+                num_layers=num_layers,
+            )
+
+        B = max(1, batch_size)
+        ps = page_size
+        P = math.ceil(max_seq_len / ps)
+        total_pages = B * P
+
+        # ---- Paged KV caches ----
+        self._paged_kv: List[Tensor] = [
+            torch.zeros(
+                total_pages, 2, ps, num_kv_heads, head_dim,
+                device=device, dtype=dtype,
+            )
+            for _ in range(num_layers)
+        ]
+
+        self._max_seq_len = max_seq_len
+        self._num_layers = num_layers
+        self._num_heads = num_heads
+        self._num_kv_heads = num_kv_heads
+        self._head_dim = head_dim
+        self._device = device
+        self._dtype = dtype
+        self._batch_size = B
+        self._page_size = ps
+        self._pages_per_sample = P
+
+        q_dim = num_heads * head_dim
+        kv_dim = num_kv_heads * head_dim
+        self._q_dim = q_dim
+        self._kv_dim = kv_dim
+        hidden_size = num_heads * head_dim
+        self._hidden_size = hidden_size
+
+        # ---- Fused projection weights ----
+        self._fused_qkv_w: List[Tensor] = []
+        self._fused_gate_up_w: List[Tensor] = []
+        self._o_proj_w: List[Tensor] = []
+        self._down_proj_w: List[Tensor] = []
+        self._input_ln: list = []
+        self._post_ln: list = []
+        self._mlp_modules: list = []
+        # ---- Cached RMSNorm weight/eps for FlashInfer fused kernels ----
+        # ``flashinfer.norm.{rmsnorm, fused_add_rmsnorm}`` needs the raw weight
+        # tensor + scalar eps. Caching avoids re-fetching attrs in the hot loop.
+        self._input_ln_w: List[Tensor] = []
+        self._input_ln_eps: List[float] = []
+        self._post_ln_w: List[Tensor] = []
+        self._post_ln_eps: List[float] = []
+
+        for layer in model.tfmr.layers:
+            attn = layer.self_attn
+            mlp = layer.mlp
+            if self._projection_provider is None:
+                self._fused_qkv_w.append(
+                    torch.cat(
+                        [attn.q_proj.weight.data, attn.k_proj.weight.data,
+                         attn.v_proj.weight.data], dim=0,
+                    )
+                )
+                self._o_proj_w.append(attn.o_proj.weight)
+                if hasattr(mlp, "gate_proj") and hasattr(mlp, "up_proj"):
+                    self._fused_gate_up_w.append(
+                        torch.cat(
+                            [mlp.gate_proj.weight.data, mlp.up_proj.weight.data],
+                            dim=0,
+                        )
+                    )
+                    self._down_proj_w.append(mlp.down_proj.weight)
+            self._input_ln.append(layer.input_layernorm)
+            self._post_ln.append(layer.post_attention_layernorm)
+            self._mlp_modules.append(mlp)
+            self._input_ln_w.append(layer.input_layernorm.weight.data)
+            self._input_ln_eps.append(
+                float(layer.input_layernorm.variance_epsilon)
+            )
+            self._post_ln_w.append(
+                layer.post_attention_layernorm.weight.data
+            )
+            self._post_ln_eps.append(
+                float(layer.post_attention_layernorm.variance_epsilon)
+            )
+
+        self._mlp_act_fn = model.tfmr.layers[0].mlp.act_fn
+        self._final_norm = model.tfmr.norm
+        self._final_norm_w = model.tfmr.norm.weight.data
+        self._final_norm_eps = float(model.tfmr.norm.variance_epsilon)
+        # FlashInfer rmsnorm kernels support fp16 / bf16. fp32 falls back to the
+        # legacy LlamaRMSNorm module path.
+        self._use_fused_norm: bool = dtype in (torch.float16, torch.bfloat16)
+        self._use_fused_silu: bool = (
+            self._use_fused_norm
+            and (
+                self._projection_provider is not None
+                or bool(self._fused_gate_up_w)
+            )
+        )
+
+        if self._projection_provider is not None:
+            self._projection_provider.prepare_for_engine(
+                batch_size=B,
+                block_size=int(model.drf_block_size),
+            )
+
+        # ---- RoPE ----
+        # transformers ≥5.9 consolidates rope settings into a single
+        # ``config.rope_parameters`` dict and DROPS the top-level
+        # ``config.rope_theta`` / ``config.rope_scaling`` attributes; older
+        # versions expose them at top level. Read both so we get the real
+        # values regardless of layout (the T3 backbone uses theta=500000 with
+        # llama3 scaling, both of which now live only in rope_parameters).
+        _rope_params = getattr(config, "rope_parameters", None) or {}
+        self._rope_theta: float = float(
+            getattr(config, "rope_theta", None)
+            or _rope_params.get("rope_theta")
+            or 10000.0
+        )
+        # ``flashinfer.rope.apply_rope_inplace`` only implements plain-theta
+        # RoPE (``rope_scale=1.0``); it does NOT apply llama3-style frequency
+        # scaling. With llama3 scaling the inplace kernel encodes wrong
+        # positions and the model degenerates into repetitive output. When any
+        # non-plain rope_type is configured, fall back to the precomputed HF
+        # cos/sin table below (built from ``model.tfmr.rotary_emb``, which bakes
+        # in the scaling) — exactly matching the torch_sdpa path.
+        _rope_scaling = getattr(config, "rope_scaling", None) or _rope_params or None
+        # transformers has used both "rope_type" (≥4.40) and the legacy "type"
+        # key; accept either, and treat anything other than default/linear
+        # (e.g. "llama3", "dynamic", "yarn") as needing the table fallback.
+        _rope_type = (
+            (_rope_scaling.get("rope_type") or _rope_scaling.get("type"))
+            if _rope_scaling
+            else None
+        )
+        _plain_rope = _rope_type in (None, "default", "linear")
+        self._use_rope_inplace: bool = (
+            _FI_ROPE_AVAILABLE and device.type == "cuda" and _plain_rope
+        )
+
+        if self._use_rope_inplace:
+            self._rope_indptr = torch.zeros(
+                B + 1, dtype=torch.int32, device=device,
+            )
+            self._rope_offset = torch.zeros(
+                B, dtype=torch.int32, device=device,
+            )
+            self._rope_indptr_1 = torch.zeros(
+                2, dtype=torch.int32, device=device,
+            )
+            self._rope_offset_1 = torch.zeros(
+                1, dtype=torch.int32, device=device,
+            )
+
+        dim = config.hidden_size
+        _dummy = torch.zeros(1, 1, dim, device=device, dtype=dtype)
+        _full_pos = torch.arange(max_seq_len, device=device).unsqueeze(0)
+        self._rope_cos, self._rope_sin = model.tfmr.rotary_emb(
+            _dummy, _full_pos,
+        )
+        self._rope_cos = self._rope_cos.detach()
+        self._rope_sin = self._rope_sin.detach()
+
+        from transformers.models.llama.modeling_llama import (
+            apply_rotary_pos_emb,
+        )
+        self._apply_rotary_pos_emb = apply_rotary_pos_emb
+
+        # ---- Per-sample cache state ----
+        self._cache_lens: List[int] = [0] * B
+        self._prefix_lens: List[int] = [0] * B
+
+        # ---- Eager paged wrapper (block forward without CUDA graph) ----
+        self._fi_workspace = torch.empty(
+            256 * 1024 * 1024, dtype=torch.uint8, device=device,
+        )
+        self._fi_eager_wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
+            self._fi_workspace, kv_layout="NHD",
+        )
+        # ---- Ragged prefill wrapper (cross-text prefix forward, eager only) ----
+        # Used to fuse per-sample ``single_prefill_with_kv_cache`` calls into one
+        # ragged-batched prefill per layer. Has its own workspace so its plan
+        # state cannot clobber the paged wrapper.
+        self._fi_ragged_workspace = torch.empty(
+            64 * 1024 * 1024, dtype=torch.uint8, device=device,
+        )
+        self._fi_ragged_wrapper = (
+            flashinfer.BatchPrefillWithRaggedKVCacheWrapper(
+                self._fi_ragged_workspace, kv_layout="NHD",
+            )
+        )
+
+        # ---- CUDA graph state ----
+        self._use_cuda_graph: bool = False
+        self._cuda_graph = None
+        self._graph_pool = None
+        self._g_block_size: int = 0
+        self._g_inputs_embeds: Tensor | None = None
+        self._g_output_hidden: Tensor | None = None
+        self._g_cos: Tensor | None = None
+        self._g_sin: Tensor | None = None
+        self._fi_graph_wrapper = None
+        self._g_page_write_idx: Tensor | None = None
+        self._g_page_write_off: Tensor | None = None
+        self._last_plan_cache_start: int | tuple = -1
+        self._plan_qo_indptr: Tensor | None = None
+        self._plan_paged_kv_indptr: Tensor | None = None
+        self._plan_paged_kv_indices: Tensor | None = None
+        self._plan_paged_kv_last_page_len: Tensor | None = None
+
+        # ---- Speech head in graph (optional) ----
+        self._has_speech_head: bool = False
+        self._speech_head_w: Tensor | None = None
+        self._speech_head_b: Tensor | None = None
+        self._g_shift_ctx: Tensor | None = None
+        self._g_logits: Tensor | None = None
+
+        logger.info(
+            "FlashInferEngine: max_seq=%d layers=%d heads=%d/%d dim=%d "
+            "batch=%d page_size=%d pages/sample=%d total_pages=%d "
+            "rope_inplace=%s dtype=%s",
+            max_seq_len, num_layers, num_heads, num_kv_heads, head_dim,
+            B, ps, P, total_pages, self._use_rope_inplace, dtype,
+        )
+
+    # ------------------------------------------------------------------ #
+    #  Page-index helpers
+    # ------------------------------------------------------------------ #
+
+    def _compute_page_write_indices(
+        self, cache_starts: int | List[int], bl: int,
+    ) -> tuple[Tensor, Tensor]:
+        """Flat page indices + offsets for writing ``B * bl`` KV entries."""
+        B = self._batch_size
+        P = self._pages_per_sample
+        ps = self._page_size
+        dev = self._device
+
+        if isinstance(cache_starts, int):
+            pos = torch.arange(cache_starts, cache_starts + bl, device=dev)
+            p_in_s = pos // ps
+            off = pos % ps
+            base = torch.arange(B, device=dev, dtype=torch.long) * P
+            return (
+                (base[:, None] + p_in_s[None, :]).reshape(-1),
+                off.repeat(B),
+            )
+
+        parts_idx: list[Tensor] = []
+        parts_off: list[Tensor] = []
+        for bi in range(B):
+            cs = cache_starts[bi]
+            pos = torch.arange(cs, cs + bl, device=dev)
+            parts_idx.append(bi * P + pos // ps)
+            parts_off.append(pos % ps)
+        return torch.cat(parts_idx), torch.cat(parts_off)
+
+    def _plan_paged_attention(
+        self,
+        wrapper,
+        kv_lens: int | List[int],
+        seq_len: int,
+        causal: bool = False,
+        *,
+        qo_indptr_buf: Tensor | None = None,
+        paged_kv_indptr_buf: Tensor | None = None,
+        paged_kv_indices_buf: Tensor | None = None,
+        paged_kv_last_page_len_buf: Tensor | None = None,
+    ):
+        """Call ``wrapper.plan()`` with page-table metadata."""
+        B = self._batch_size
+        P = self._pages_per_sample
+        ps = self._page_size
+        dev = self._device
+
+        kv_list = [kv_lens] * B if isinstance(kv_lens, int) else list(kv_lens)
+        pages_used = [math.ceil(kl / ps) for kl in kv_list]
+        total_pu = sum(pages_used)
+
+        use_buf = qo_indptr_buf is not None
+
+        qo = qo_indptr_buf if use_buf else torch.zeros(
+            B + 1, dtype=torch.int32, device=dev,
+        )
+        kvi = paged_kv_indptr_buf if use_buf else torch.zeros(
+            B + 1, dtype=torch.int32, device=dev,
+        )
+        kvidx = paged_kv_indices_buf if use_buf else torch.empty(
+            total_pu, dtype=torch.int32, device=dev,
+        )
+        kvlp = paged_kv_last_page_len_buf if use_buf else torch.empty(
+            B, dtype=torch.int32, device=dev,
+        )
+
+        if use_buf:
+            qo[0] = 0
+            kvi[0] = 0
+
+        idx = 0
+        for bi in range(B):
+            pu = pages_used[bi]
+            qo[bi + 1] = qo[bi] + seq_len
+            kvi[bi + 1] = kvi[bi] + pu
+            kvidx[idx: idx + pu] = torch.arange(
+                bi * P, bi * P + pu, dtype=torch.int32, device=dev,
+            )
+            idx += pu
+            last = kv_list[bi] % ps
+            kvlp[bi] = last if last > 0 else ps
+
+        wrapper.plan(
+            qo_indptr=qo,
+            paged_kv_indptr=kvi,
+            paged_kv_indices=kvidx[:total_pu],
+            paged_kv_last_page_len=kvlp,
+            num_qo_heads=self._num_heads,
+            num_kv_heads=self._num_kv_heads,
+            head_dim_qk=self._head_dim,
+            page_size=ps,
+            causal=causal,
+            q_data_type=self._dtype,
+            use_fp16_qk_reduction=self._use_fp16_qk_red,
+        )
+
+    # ------------------------------------------------------------------ #
+    #  MLP helper (shared by prefix + block + graph paths)
+    # ------------------------------------------------------------------ #
+
+    def _project(
+        self,
+        kind: str,
+        layer_idx: int,
+        activation: Tensor,
+    ) -> Tensor:
+        """Apply an engine transformer projection at its explicit boundary."""
+        if self._projection_provider is not None:
+            return self._projection_provider.project(
+                kind, layer_idx, activation,
+            )
+        weights = {
+            "qkv": self._fused_qkv_w,
+            "o": self._o_proj_w,
+            "gate_up": self._fused_gate_up_w,
+            "down": self._down_proj_w,
+        }
+        try:
+            weight = weights[kind][layer_idx]
+        except (KeyError, IndexError) as exc:
+            raise RuntimeError(
+                f"Missing {kind} projection for transformer layer {layer_idx}"
+            ) from exc
+        return F.linear(activation, weight)
+
+    # ------------------------------------------------------------------ #
+    #  FlashInfer kernel wrappers (forward PDL flag once)
+    # ------------------------------------------------------------------ #
+
+    def _fi_rmsnorm(self, x: Tensor, w: Tensor, eps: float) -> Tensor:
+        return flashinfer.norm.rmsnorm(
+            x, w, eps, enable_pdl=self._enable_pdl,
+        )
+
+    def _fi_fused_add_rmsnorm(
+        self, x: Tensor, residual: Tensor, w: Tensor, eps: float,
+    ) -> None:
+        _flashinfer_fused_add_rmsnorm(
+            x, residual, w, eps, enable_pdl=self._enable_pdl,
+        )
+
+    def _mlp_forward(self, layer_idx: int, hidden: Tensor) -> Tensor:
+        """Llama SwiGLU MLP. Accepts any leading dims ``(..., dim)``.
+
+        When fused gate-up weight + bf16/fp16 are available, the
+        ``silu(gate) * up`` elementwise op is replaced by FlashInfer's
+        ``silu_and_mul`` (single fused kernel, 1 read + 1 write of the
+        intermediate tensor instead of 2 reads + 1 write).
+        """
+        if self._projection_provider is None and not self._fused_gate_up_w:
+            return self._mlp_modules[layer_idx](hidden)
+
+        gate_up = self._project("gate_up", layer_idx, hidden)
+        if self._use_fused_silu:
+            orig_shape = gate_up.shape
+            gate_up_2d = gate_up.reshape(-1, orig_shape[-1])
+            act_2d = flashinfer.activation.silu_and_mul(
+                gate_up_2d, enable_pdl=self._enable_pdl,
+            )
+            act = act_2d.view(*orig_shape[:-1], act_2d.size(-1))
+        else:
+            mid = gate_up.size(-1) // 2
+            act = self._mlp_act_fn(gate_up[..., :mid]) * gate_up[..., mid:]
+        return self._project("down", layer_idx, act)
+
+    # ------------------------------------------------------------------ #
+    #  Prefix forward  (single-request per sample, paged KV write)
+    # ------------------------------------------------------------------ #
+
+    def prefix_forward(
+        self,
+        prefix_emb_or_list: Tensor | List[Tensor],
+        prefix_lens: List[int] | None = None,
+    ) -> Tensor:
+        """Compute prefix KV and store in paged cache.
+
+        Same-text mode
+            ``prefix_emb_or_list`` is ``(1, lp, dim)``.
+            Computed once for sample 0, pages copied to samples 1..B-1.
+
+        Cross-text mode
+            ``prefix_emb_or_list`` is a list of ``(1, lp_i, dim)``.
+            Each sample processed independently.
+
+        Returns:
+            ``(B, 1, dim)`` shift context (SOS hidden per sample).
+        """
+        B = self._batch_size
+
+        if isinstance(prefix_emb_or_list, list):
+            # Cross-text: pack all samples into one ragged batched prefill so
+            # we issue O(L) kernel launches instead of O(L * B).
+            return self._prefix_forward_batched(prefix_emb_or_list)
+
+        emb = prefix_emb_or_list[:1] if prefix_emb_or_list.size(0) > 1 else prefix_emb_or_list
+        ctx = self._prefix_forward_one(emb, sample_idx=0)
+        lp = prefix_emb_or_list.size(1)
+        pfx_pages = math.ceil(lp / self._page_size)
+        P = self._pages_per_sample
+        for li in range(self._num_layers):
+            src = self._paged_kv[li][:pfx_pages]
+            for bi in range(1, B):
+                s = bi * P
+                self._paged_kv[li][s: s + pfx_pages] = src
+        for bi in range(1, B):
+            self._cache_lens[bi] = lp
+            self._prefix_lens[bi] = lp
+        return ctx.expand(B, -1, -1).contiguous()
+
+    def _prefix_forward_one(
+        self, prefix_emb: Tensor, sample_idx: int,
+    ) -> Tensor:
+        """Single-sample prefix prefill, writes into ``sample_idx`` page slice.
+
+        Decoder body uses the vLLM-style ``(hidden, residual)`` carry pattern
+        with :func:`flashinfer.norm.fused_add_rmsnorm` (in-place residual add
+        + RMSNorm in a single kernel). Attention is the optimized
+        :func:`flashinfer.single_prefill_with_kv_cache` (causal).
+        """
+        lp = prefix_emb.size(1)
+        page_off = sample_idx * self._pages_per_sample
+        ps = self._page_size
+        dev = self._device
+
+        pos = torch.arange(lp, device=dev)
+        pw_idx = page_off + pos // ps
+        pw_off = pos % ps
+
+        if self._use_rope_inplace:
+            self._rope_indptr_1[1] = lp
+            self._rope_offset_1[0] = 0
+
+        nh = self._num_heads
+        nkv = self._num_kv_heads
+        hd = self._head_dim
+        qd = self._q_dim
+        kvd = self._kv_dim
+
+        if self._use_fused_norm:
+            return self._prefix_forward_one_fused(
+                prefix_emb, sample_idx, lp, pw_idx, pw_off,
+                nh, nkv, hd, qd, kvd,
+            )
+
+        hidden = prefix_emb
+
+        for li in range(self._num_layers):
+            residual = hidden
+            hidden = self._input_ln[li](hidden)
+
+            qkv = self._project("qkv", li, hidden)
+            q = qkv[:, :, :qd].view(1, lp, nh, hd)
+            k = qkv[:, :, qd: qd + kvd].view(1, lp, nkv, hd)
+            v = qkv[:, :, qd + kvd:].view(1, lp, nkv, hd)
+
+            qs = q.squeeze(0)
+            ks = k.squeeze(0)
+            vs = v.squeeze(0)
+
+            if self._use_rope_inplace:
+                _fi_apply_rope(
+                    qs, ks, self._rope_indptr_1, self._rope_offset_1,
+                    interleave=False, rope_scale=1.0,
+                    rope_theta=self._rope_theta,
+                )
+            else:
+                cos = self._rope_cos[:, :lp, :]
+                sin = self._rope_sin[:, :lp, :]
+                q, k = self._apply_rotary_pos_emb(
+                    q, k, cos, sin, unsqueeze_dim=2,
+                )
+                qs = q.squeeze(0)
+                ks = k.squeeze(0)
+
+            attn_out = flashinfer.single_prefill_with_kv_cache(
+                qs, ks, vs, causal=True, kv_layout="NHD",
+                use_fp16_qk_reduction=self._use_fp16_qk_red,
+            )
+
+            self._paged_kv[li][pw_idx, 0, pw_off] = ks
+            self._paged_kv[li][pw_idx, 1, pw_off] = vs
+
+            attn_out = attn_out.unsqueeze(0).reshape(1, lp, -1)
+            attn_out = self._project("o", li, attn_out)
+
+            hidden = residual + attn_out
+            residual = hidden
+            hidden = self._post_ln[li](hidden)
+            hidden = self._mlp_forward(li, hidden)
+            hidden = residual + hidden
+
+        hidden = self._final_norm(hidden)
+        self._cache_lens[sample_idx] = lp
+        self._prefix_lens[sample_idx] = lp
+        return hidden[:, lp - 1: lp, :]
+
+    def _prefix_forward_one_fused(
+        self,
+        prefix_emb: Tensor,
+        sample_idx: int,
+        lp: int,
+        pw_idx: Tensor,
+        pw_off: Tensor,
+        nh: int,
+        nkv: int,
+        hd: int,
+        qd: int,
+        kvd: int,
+    ) -> Tensor:
+        """fused_add_rmsnorm + silu_and_mul variant of :meth:`_prefix_forward_one`."""
+        L = self._num_layers
+        # Work in 2D ``(lp, dim)`` so fused_add_rmsnorm (which wants 2D contig)
+        # and silu_and_mul (any leading dims) both apply cleanly.
+        # Clone so we don't mutate the caller's ``prefix_emb``.
+        residual = prefix_emb.reshape(lp, -1).contiguous().clone()
+        hidden = self._fi_rmsnorm(
+            residual, self._input_ln_w[0], self._input_ln_eps[0],
+        )
+
+        for li in range(L):
+            qkv = self._project("qkv", li, hidden)
+            qs = qkv[:, :qd].view(lp, nh, hd)
+            ks = qkv[:, qd: qd + kvd].view(lp, nkv, hd)
+            vs = qkv[:, qd + kvd:].view(lp, nkv, hd)
+
+            if self._use_rope_inplace:
+                _fi_apply_rope(
+                    qs, ks, self._rope_indptr_1, self._rope_offset_1,
+                    interleave=False, rope_scale=1.0,
+                    rope_theta=self._rope_theta,
+                )
+            else:
+                # Apply RoPE via HF helper (expects ``(B, S, h, hd)``).
+                q3 = qs.unsqueeze(0)
+                k3 = ks.unsqueeze(0)
+                cos = self._rope_cos[:, :lp, :]
+                sin = self._rope_sin[:, :lp, :]
+                q3, k3 = self._apply_rotary_pos_emb(
+                    q3, k3, cos, sin, unsqueeze_dim=2,
+                )
+                qs = q3.squeeze(0)
+                ks = k3.squeeze(0)
+
+            attn_out = flashinfer.single_prefill_with_kv_cache(
+                qs, ks, vs, causal=True, kv_layout="NHD",
+                use_fp16_qk_reduction=self._use_fp16_qk_red,
+            )
+
+            self._paged_kv[li][pw_idx, 0, pw_off] = ks
+            self._paged_kv[li][pw_idx, 1, pw_off] = vs
+
+            attn_out_2d = attn_out.reshape(lp, -1)
+            attn_out_2d = self._project("o", li, attn_out_2d).contiguous()
+
+            # residual += attn_out; attn_out = rmsnorm(residual, post_ln_w)
+            self._fi_fused_add_rmsnorm(
+                attn_out_2d, residual,
+                self._post_ln_w[li], self._post_ln_eps[li],
+            )
+
+            mlp_out = self._mlp_forward(li, attn_out_2d)
+            if not mlp_out.is_contiguous():
+                mlp_out = mlp_out.contiguous()
+
+            if li + 1 < L:
+                self._fi_fused_add_rmsnorm(
+                    mlp_out, residual,
+                    self._input_ln_w[li + 1], self._input_ln_eps[li + 1],
+                )
+                hidden = mlp_out
+            else:
+                residual.add_(mlp_out)
+                hidden = residual
+
+        # Final norm and per-sample SOS hidden.
+        out = self._fi_rmsnorm(
+            hidden, self._final_norm_w, self._final_norm_eps,
+        )
+        self._cache_lens[sample_idx] = lp
+        self._prefix_lens[sample_idx] = lp
+        return out[lp - 1: lp, :].view(1, 1, -1)
+
+    # ------------------------------------------------------------------ #
+    #  Cross-text ragged batched prefix forward
+    # ------------------------------------------------------------------ #
+
+    def _prefix_forward_batched(
+        self, emb_list: List[Tensor],
+    ) -> Tensor:
+        """Cross-text prefix prefill via ragged batched attention.
+
+        Replaces the per-sample loop of
+        :meth:`_prefix_forward_one` with a single
+        :class:`flashinfer.BatchPrefillWithRaggedKVCacheWrapper.run`
+        per layer. Sample ``bi`` may have its own prefix length ``lp_i``;
+        embeddings are packed into a flat ``(sum_i lp_i, dim)`` tensor.
+
+        Each sample's KV is scattered into its own paged-cache slice
+        ``[bi*P : bi*P + ceil(lp_i/ps)]``.
+
+        Returns:
+            ``(B_actual, 1, dim)`` per-sample SOS hidden.
+        """
+        B_actual = len(emb_list)
+        if B_actual == 1:
+            # Single sample: paged ragged setup overhead is wasteful. Reuse the
+            # well-tuned single-sample path which calls single_prefill directly.
+            return self._prefix_forward_one(emb_list[0], 0)
+
+        dev = self._device
+        ps = self._page_size
+        P = self._pages_per_sample
+        nh = self._num_heads
+        nkv = self._num_kv_heads
+        hd = self._head_dim
+        qd = self._q_dim
+        kvd = self._kv_dim
+        L = self._num_layers
+        dtype = self._dtype
+
+        lp_list = [int(e.size(1)) for e in emb_list]
+        total_lp = sum(lp_list)
+
+        # Pack embeddings to (total_lp, dim).
+        packed = torch.cat(
+            [e.reshape(-1, e.size(-1)) for e in emb_list], dim=0,
+        ).to(dtype=dtype).contiguous()
+
+        # qo_indptr == kv_indptr for self-attention prefill.
+        indptr_host = [0]
+        for lp_i in lp_list:
+            indptr_host.append(indptr_host[-1] + lp_i)
+        qo_indptr = torch.tensor(
+            indptr_host, dtype=torch.int32, device=dev,
+        )
+
+        # Paged-write indices: sample bi's positions [0, lp_i) → pages
+        # [bi*P, bi*P + ceil(lp_i/ps)).
+        pw_idx_parts: list[Tensor] = []
+        pw_off_parts: list[Tensor] = []
+        for bi, lp_i in enumerate(lp_list):
+            pos = torch.arange(lp_i, device=dev)
+            pw_idx_parts.append(bi * P + pos // ps)
+            pw_off_parts.append(pos % ps)
+        pw_idx = torch.cat(pw_idx_parts)
+        pw_off = torch.cat(pw_off_parts)
+
+        # RoPE: each sample starts from position 0 (prefix). Reuse indptr as
+        # rope_indptr; offsets are all zero. ``apply_rope_inplace`` wants int32.
+        if self._use_rope_inplace:
+            rope_indptr = qo_indptr.to(torch.int32)
+            rope_offset = torch.zeros(
+                B_actual, dtype=torch.int32, device=dev,
+            )
+        else:
+            # Pre-build packed cos/sin once; shape (total_lp, head_dim).
+            cos_parts = [self._rope_cos[0, :lp_i, :] for lp_i in lp_list]
+            sin_parts = [self._rope_sin[0, :lp_i, :] for lp_i in lp_list]
+            packed_cos = torch.cat(cos_parts, dim=0)
+            packed_sin = torch.cat(sin_parts, dim=0)
+
+        # Plan ragged wrapper once for the whole layer stack (causal).
+        self._fi_ragged_wrapper.plan(
+            qo_indptr=qo_indptr,
+            kv_indptr=qo_indptr,
+            num_qo_heads=nh,
+            num_kv_heads=nkv,
+            head_dim_qk=hd,
+            causal=True,
+            q_data_type=dtype,
+            use_fp16_qk_reduction=self._use_fp16_qk_red,
+        )
+
+        use_fused = self._use_fused_norm
+        if use_fused:
+            residual = packed.clone()
+            hidden = self._fi_rmsnorm(
+                residual, self._input_ln_w[0], self._input_ln_eps[0],
+            )
+        else:
+            hidden = packed  # (total_lp, dim)
+
+        for li in range(L):
+            if use_fused:
+                qkv = self._project("qkv", li, hidden)
+            else:
+                # Module-style norm requires 3D ``(1, total_lp, dim)``.
+                normed = self._input_ln[li](hidden.unsqueeze(0)).squeeze(0)
+                qkv = self._project("qkv", li, normed)
+
+            qs = qkv[:, :qd].view(total_lp, nh, hd)
+            ks = qkv[:, qd: qd + kvd].view(total_lp, nkv, hd)
+            vs = qkv[:, qd + kvd:].view(total_lp, nkv, hd)
+
+            if self._use_rope_inplace:
+                _fi_apply_rope(
+                    qs, ks, rope_indptr, rope_offset,
+                    interleave=False, rope_scale=1.0,
+                    rope_theta=self._rope_theta,
+                )
+            else:
+                # ``apply_rotary_pos_emb`` expects ``(B, S, h, hd)``; mimic with
+                # a synthetic batch dim over the packed sequence.
+                cos = packed_cos.unsqueeze(0)
+                sin = packed_sin.unsqueeze(0)
+                q3 = qs.unsqueeze(0)
+                k3 = ks.unsqueeze(0)
+                q3, k3 = self._apply_rotary_pos_emb(
+                    q3, k3, cos, sin, unsqueeze_dim=2,
+                )
+                qs = q3.squeeze(0)
+                ks = k3.squeeze(0)
+
+            self._paged_kv[li][pw_idx, 0, pw_off] = ks
+            self._paged_kv[li][pw_idx, 1, pw_off] = vs
+
+            attn_out = self._fi_ragged_wrapper.run(qs, ks, vs)
+            attn_out_2d = attn_out.reshape(total_lp, -1)
+            attn_out_2d = self._project(
+                "o", li, attn_out_2d,
+            ).contiguous()
+
+            if use_fused:
+                self._fi_fused_add_rmsnorm(
+                    attn_out_2d, residual,
+                    self._post_ln_w[li], self._post_ln_eps[li],
+                )
+                mlp_out = self._mlp_forward(li, attn_out_2d)
+                if not mlp_out.is_contiguous():
+                    mlp_out = mlp_out.contiguous()
+                if li + 1 < L:
+                    self._fi_fused_add_rmsnorm(
+                        mlp_out, residual,
+                        self._input_ln_w[li + 1], self._input_ln_eps[li + 1],
+                    )
+                    hidden = mlp_out
+                else:
+                    residual.add_(mlp_out)
+                    hidden = residual
+            else:
+                # Non-fused fallback: keep classic residual carry.
+                hidden = hidden + attn_out_2d  # residual after attn
+                residual_after_attn = hidden
+                hidden_n = self._post_ln[li](hidden.unsqueeze(0)).squeeze(0)
+                mlp_out = self._mlp_forward(li, hidden_n)
+                hidden = residual_after_attn + mlp_out
+
+        if use_fused:
+            out = self._fi_rmsnorm(
+                hidden, self._final_norm_w, self._final_norm_eps,
+            )
+        else:
+            out = self._final_norm(hidden.unsqueeze(0)).squeeze(0)
+
+        # Per-sample SOS hidden = out[lp_i - 1] within each sample's slice.
+        shift_idx = (qo_indptr[1:] - 1).to(torch.long)
+        shift_ctx = out.index_select(0, shift_idx).unsqueeze(1)  # (B_actual, 1, dim)
+
+        for bi, lp_i in enumerate(lp_list):
+            self._cache_lens[bi] = lp_i
+            self._prefix_lens[bi] = lp_i
+
+        return shift_ctx.contiguous()
+
+    # ------------------------------------------------------------------ #
+    #  Batched block forward  (non-causal, paged attention)
+    # ------------------------------------------------------------------ #
+
+    def block_forward(
+        self,
+        block_emb: Tensor,
+        cache_starts: int | List[int],
+    ) -> Tensor:
+        """Batched block forward with non-causal paged attention.
+
+        Args:
+            block_emb:    ``(B, bl, dim)``
+            cache_starts: int (uniform) or list[int] (per-sample).
+
+        Returns:
+            ``(B, bl, dim)`` after final RMSNorm.
+        """
+        B = self._batch_size
+        bl = block_emb.size(1)
+
+        pw_idx, pw_off = self._compute_page_write_indices(cache_starts, bl)
+
+        if isinstance(cache_starts, int):
+            kv_lens: int | List[int] = cache_starts + bl
+        else:
+            kv_lens = [cs + bl for cs in cache_starts]
+
+        self._plan_paged_attention(
+            self._fi_eager_wrapper, kv_lens, seq_len=bl, causal=False,
+        )
+
+        if self._use_rope_inplace:
+            for bi in range(B):
+                self._rope_indptr[bi + 1] = (bi + 1) * bl
+            if isinstance(cache_starts, int):
+                self._rope_offset[:] = cache_starts
+            else:
+                for bi in range(B):
+                    self._rope_offset[bi] = cache_starts[bi]
+
+        nh = self._num_heads
+        nkv = self._num_kv_heads
+        hd = self._head_dim
+        qd = self._q_dim
+        kvd = self._kv_dim
+
+        if self._use_fused_norm:
+            return self._block_forward_fused(
+                block_emb, cache_starts, B, bl, pw_idx, pw_off,
+                nh, nkv, hd, qd, kvd,
+            )
+
+        hidden = block_emb
+
+        for li in range(self._num_layers):
+            residual = hidden
+            hidden = self._input_ln[li](hidden)
+
+            qkv = self._project("qkv", li, hidden)
+            q = qkv[:, :, :qd].view(B, bl, nh, hd)
+            k = qkv[:, :, qd: qd + kvd].view(B, bl, nkv, hd)
+            v = qkv[:, :, qd + kvd:].view(B, bl, nkv, hd)
+
+            qs = q.reshape(B * bl, nh, hd)
+            ks = k.reshape(B * bl, nkv, hd)
+            vs = v.reshape(B * bl, nkv, hd)
+
+            if self._use_rope_inplace:
+                _fi_apply_rope(
+                    qs, ks, self._rope_indptr, self._rope_offset,
+                    interleave=False, rope_scale=1.0,
+                    rope_theta=self._rope_theta,
+                )
+            else:
+                if isinstance(cache_starts, int):
+                    cos = self._rope_cos[
+                        :, cache_starts: cache_starts + bl, :
+                    ]
+                    sin = self._rope_sin[
+                        :, cache_starts: cache_starts + bl, :
+                    ]
+                else:
+                    cos = torch.cat([
+                        self._rope_cos[:, cs: cs + bl, :]
+                        for cs in cache_starts
+                    ], dim=0)
+                    sin = torch.cat([
+                        self._rope_sin[:, cs: cs + bl, :]
+                        for cs in cache_starts
+                    ], dim=0)
+                q, k = self._apply_rotary_pos_emb(
+                    q, k, cos, sin, unsqueeze_dim=2,
+                )
+                qs = q.reshape(B * bl, nh, hd)
+                ks = k.reshape(B * bl, nkv, hd)
+
+            self._paged_kv[li][pw_idx, 0, pw_off] = ks
+            self._paged_kv[li][pw_idx, 1, pw_off] = vs
+
+            attn_out = self._fi_eager_wrapper.run(
+                qs, self._paged_kv[li],
+            )
+            attn_out = attn_out.reshape(B, bl, -1)
+            attn_out = self._project("o", li, attn_out)
+
+            hidden = residual + attn_out
+            residual = hidden
+            hidden = self._post_ln[li](hidden)
+            hidden = self._mlp_forward(li, hidden)
+            hidden = residual + hidden
+
+        return self._final_norm(hidden)
+
+    def _block_forward_fused(
+        self,
+        block_emb: Tensor,
+        cache_starts: int | List[int],
+        B: int,
+        bl: int,
+        pw_idx: Tensor,
+        pw_off: Tensor,
+        nh: int,
+        nkv: int,
+        hd: int,
+        qd: int,
+        kvd: int,
+    ) -> Tensor:
+        """``block_forward`` body with fused_add_rmsnorm + silu_and_mul.
+
+        Carries ``(hidden, residual)`` as 2D ``(B*bl, dim)``. ``residual`` is
+        a private clone of ``block_emb`` so the caller's tensor is not mutated.
+        """
+        L = self._num_layers
+        N = B * bl
+
+        residual = block_emb.reshape(N, -1).contiguous().clone()
+        hidden = self._fi_rmsnorm(
+            residual, self._input_ln_w[0], self._input_ln_eps[0],
+        )
+
+        for li in range(L):
+            qkv = self._project("qkv", li, hidden)
+            qs = qkv[:, :qd].view(N, nh, hd)
+            ks = qkv[:, qd: qd + kvd].view(N, nkv, hd)
+            vs = qkv[:, qd + kvd:].view(N, nkv, hd)
+
+            if self._use_rope_inplace:
+                _fi_apply_rope(
+                    qs, ks, self._rope_indptr, self._rope_offset,
+                    interleave=False, rope_scale=1.0,
+                    rope_theta=self._rope_theta,
+                )
+            else:
+                if isinstance(cache_starts, int):
+                    cos = self._rope_cos[
+                        :, cache_starts: cache_starts + bl, :
+                    ]
+                    sin = self._rope_sin[
+                        :, cache_starts: cache_starts + bl, :
+                    ]
+                else:
+                    cos = torch.cat([
+                        self._rope_cos[:, cs: cs + bl, :]
+                        for cs in cache_starts
+                    ], dim=0)
+                    sin = torch.cat([
+                        self._rope_sin[:, cs: cs + bl, :]
+                        for cs in cache_starts
+                    ], dim=0)
+                q3 = qs.view(B, bl, nh, hd)
+                k3 = ks.view(B, bl, nkv, hd)
+                q3, k3 = self._apply_rotary_pos_emb(
+                    q3, k3, cos, sin, unsqueeze_dim=2,
+                )
+                qs = q3.reshape(N, nh, hd)
+                ks = k3.reshape(N, nkv, hd)
+
+            self._paged_kv[li][pw_idx, 0, pw_off] = ks
+            self._paged_kv[li][pw_idx, 1, pw_off] = vs
+
+            attn_out = self._fi_eager_wrapper.run(
+                qs, self._paged_kv[li],
+            )
+            attn_out_2d = attn_out.reshape(N, -1)
+            attn_out_2d = self._project("o", li, attn_out_2d).contiguous()
+
+            self._fi_fused_add_rmsnorm(
+                attn_out_2d, residual,
+                self._post_ln_w[li], self._post_ln_eps[li],
+            )
+
+            mlp_out = self._mlp_forward(li, attn_out_2d)
+            if not mlp_out.is_contiguous():
+                mlp_out = mlp_out.contiguous()
+
+            if li + 1 < L:
+                self._fi_fused_add_rmsnorm(
+                    mlp_out, residual,
+                    self._input_ln_w[li + 1], self._input_ln_eps[li + 1],
+                )
+                hidden = mlp_out
+            else:
+                residual.add_(mlp_out)
+                hidden = residual
+
+        out = self._fi_rmsnorm(
+            hidden, self._final_norm_w, self._final_norm_eps,
+        )
+        return out.view(B, bl, -1)
+
+    # ------------------------------------------------------------------ #
+    #  Cache management
+    # ------------------------------------------------------------------ #
+
+    @property
+    def batch_size(self) -> int:
+        return self._batch_size
+
+    @property
+    def cache_len(self) -> int:
+        """Backward-compat: cache length of sample 0."""
+        return self._cache_lens[0]
+
+    @cache_len.setter
+    def cache_len(self, value: int):
+        self._cache_lens = [value] * self._batch_size
+
+    @property
+    def cache_lens(self) -> List[int]:
+        return self._cache_lens
+
+    @property
+    def prefix_lens(self) -> List[int]:
+        return self._prefix_lens
+
+    def advance_cache(self, block_len: int):
+        for bi in range(self._batch_size):
+            self._cache_lens[bi] += block_len
+
+    def reset(self):
+        B = self._batch_size
+        self._cache_lens = [0] * B
+        self._prefix_lens = [0] * B
+        self._last_plan_cache_start = -1
+        # CUDA graph / pool are kept across generate() calls. Per-call prefix and
+        # paged-KV state are rebuilt in prefix_forward; block_forward_graph replans
+        # when cache_starts change. capture_cuda_graph early-returns when the same
+        # block_size is already captured (avoids allocator issues + recapture cost).
+
+    def snapshot_kv(self):
+        """Deep copy of paged KV + lengths for :meth:`restore_kv` (CFG uncond pass)."""
+        kv_copy = [t.clone() for t in self._paged_kv]
+        return (
+            kv_copy,
+            list(self._cache_lens),
+            list(self._prefix_lens),
+            int(getattr(self, "_last_plan_cache_start", -1)),
+        )
+
+    def restore_kv(self, snap) -> None:
+        """Restore state from :meth:`snapshot_kv`."""
+        kv_copy, cache_lens, prefix_lens, last_plan = snap
+        for li, t in enumerate(kv_copy):
+            self._paged_kv[li].copy_(t)
+        self._cache_lens = cache_lens
+        self._prefix_lens = prefix_lens
+        self._last_plan_cache_start = last_plan
+
+    # ------------------------------------------------------------------ #
+    #  CUDA Graph
+    # ------------------------------------------------------------------ #
+
+    @property
+    def has_cuda_graph(self) -> bool:
+        return self._use_cuda_graph
+
+    def can_reuse(
+        self,
+        max_seq_len: int,
+        dtype: torch.dtype,
+        batch_size: int = 1,
+        page_size: int = 16,
+    ) -> bool:
+        return (
+            self._max_seq_len >= max_seq_len
+            and self._dtype == dtype
+            and self._batch_size == max(1, batch_size)
+            and self._page_size == page_size
+        )
+
+    def capture_cuda_graph(
+        self,
+        block_size: int,
+        speech_head: "torch.nn.Linear | None" = None,
+    ):
+        if self._use_cuda_graph and self._g_block_size == block_size:
+            return
+
+        import gc
+        import os
+        import shutil
+        import glob as glob_mod
+
+        device = self._device
+        dtype = self._dtype
+        bl = block_size
+        B = self._batch_size
+        P = self._pages_per_sample
+        ps = self._page_size
+        hs = self._hidden_size
+
+        cache_root = os.path.expanduser("~/.cache/flashinfer")
+        for partial in glob_mod.glob(
+            os.path.join(cache_root, "**/cached_ops/batch_prefill_*"),
+            recursive=True,
+        ):
+            if os.path.isdir(partial):
+                so_files = glob_mod.glob(os.path.join(partial, "*.so"))
+                if not so_files:
+                    logger.info(
+                        "FlashInferEngine: cleaning partial JIT: %s", partial,
+                    )
+                    shutil.rmtree(partial, ignore_errors=True)
+
+        # ---- Plan buffers (pre-allocated for CUDA-graph wrapper) ----
+        max_plan_pages = B * P
+        self._plan_qo_indptr = torch.zeros(
+            B + 1, dtype=torch.int32, device=device,
+        )
+        self._plan_paged_kv_indptr = torch.zeros(
+            B + 1, dtype=torch.int32, device=device,
+        )
+        self._plan_paged_kv_indices = torch.zeros(
+            max_plan_pages, dtype=torch.int32, device=device,
+        )
+        self._plan_paged_kv_last_page_len = torch.zeros(
+            B, dtype=torch.int32, device=device,
+        )
+
+        fi_ws = torch.empty(
+            256 * 1024 * 1024, dtype=torch.uint8, device=device,
+        )
+        self._fi_graph_wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
+            fi_ws,
+            kv_layout="NHD",
+            use_cuda_graph=True,
+            qo_indptr_buf=self._plan_qo_indptr,
+            paged_kv_indptr_buf=self._plan_paged_kv_indptr,
+            paged_kv_indices_buf=self._plan_paged_kv_indices,
+            paged_kv_last_page_len_buf=self._plan_paged_kv_last_page_len,
+        )
+
+        # ---- Graph IO ----
+        self._g_inputs_embeds = torch.zeros(
+            B, bl, hs, device=device, dtype=dtype,
+        )
+        self._g_output_hidden = torch.zeros(
+            B, bl, hs, device=device, dtype=dtype,
+        )
+        self._g_block_size = bl
+
+        uniform_cache = len(set(self._cache_lens)) <= 1
+        if uniform_cache:
+            cs0 = self._cache_lens[0]
+            pw_idx, pw_off = self._compute_page_write_indices(cs0, bl)
+            kv_plan: int | List[int] = cs0 + bl
+        else:
+            starts = list(self._cache_lens)
+            pw_idx, pw_off = self._compute_page_write_indices(starts, bl)
+            kv_plan = [c + bl for c in self._cache_lens]
+
+        self._g_page_write_idx = pw_idx.clone()
+        self._g_page_write_off = pw_off.clone()
+
+        if self._use_rope_inplace:
+            for bi in range(B):
+                self._rope_indptr[bi + 1] = (bi + 1) * bl
+            if uniform_cache:
+                self._rope_offset[:] = self._cache_lens[0]
+            else:
+                for bi in range(B):
+                    self._rope_offset[bi] = self._cache_lens[bi]
+        else:
+            self._g_cos = torch.zeros(
+                B, bl, self._head_dim, device=device, dtype=dtype,
+            )
+            self._g_sin = torch.zeros(
+                B, bl, self._head_dim, device=device, dtype=dtype,
+            )
+            for bi in range(B):
+                cs_bi = self._cache_lens[bi]
+                self._g_cos[bi].copy_(self._rope_cos[0, cs_bi: cs_bi + bl, :])
+                self._g_sin[bi].copy_(self._rope_sin[0, cs_bi: cs_bi + bl, :])
+
+        # ---- Speech head in graph ----
+        if speech_head is not None:
+            self._speech_head_w = speech_head.weight.detach()
+            self._speech_head_b = (
+                speech_head.bias.detach()
+                if speech_head.bias is not None
+                else None
+            )
+            vocab_size = speech_head.weight.size(0)
+            self._g_shift_ctx = torch.zeros(
+                B, 1, hs, device=device, dtype=dtype,
+            )
+            self._g_logits = torch.zeros(
+                B, bl, vocab_size, device=device, dtype=dtype,
+            )
+            self._has_speech_head = True
+            logger.info(
+                "FlashInferEngine: speech_head in graph (vocab=%d)",
+                vocab_size,
+            )
+
+        # ---- Initial plan ----
+        logger.info(
+            "FlashInferEngine: CUDA graph JIT compiling batch_prefill "
+            "(MAX_JOBS=%s) ...",
+            os.environ.get("MAX_JOBS", "unset"),
+        )
+        self._plan_paged_attention(
+            self._fi_graph_wrapper, kv_plan, seq_len=bl, causal=False,
+            qo_indptr_buf=self._plan_qo_indptr,
+            paged_kv_indptr_buf=self._plan_paged_kv_indptr,
+            paged_kv_indices_buf=self._plan_paged_kv_indices,
+            paged_kv_last_page_len_buf=self._plan_paged_kv_last_page_len,
+        )
+        # Tuple matches :meth:`block_forward_graph` when ``cache_starts`` is a list;
+        # same-text int ``cache_starts`` triggers one replan on first replay (harmless).
+        self._last_plan_cache_start = tuple(self._cache_lens)
+
+        # ---- Warmup ----
+        logger.info("FlashInferEngine: CUDA graph warmup forward ...")
+        self._graph_body()
+        torch.cuda.synchronize()
+        gc.collect()
+        torch.cuda.synchronize()
+
+        # ---- Capture ----
+        logger.info("FlashInferEngine: CUDA graph capturing ...")
+        self._cuda_graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(self._cuda_graph, pool=self._graph_pool):
+            self._graph_body()
+
+        if self._graph_pool is None:
+            self._graph_pool = self._cuda_graph.pool()
+
+        self._use_cuda_graph = True
+        logger.info(
+            "FlashInferEngine: CUDA graph captured "
+            "(block_size=%d, batch=%d, speech_head=%s)",
+            bl, B, self._has_speech_head,
+        )
+
+    def _graph_body(self):
+        """Graph-capturable batched forward for block denoising.
+
+        In fused-norm mode the layer body carries ``(hidden, residual)`` as 2D
+        ``(B*bl, dim)`` and uses :func:`flashinfer.norm.fused_add_rmsnorm` for
+        every "residual add + RMSNorm" pair. The persistent
+        ``self._g_inputs_embeds`` buffer doubles as the running residual; it is
+        overwritten via ``copy_(block_emb)`` before each replay, so in-place
+        mutation inside the graph is safe.
+        """
+        B = self._batch_size
+        bl = self._g_block_size
+        nh = self._num_heads
+        nkv = self._num_kv_heads
+        hd = self._head_dim
+        qd = self._q_dim
+        kvd = self._kv_dim
+        use_ri = self._use_rope_inplace
+        L = self._num_layers
+        N = B * bl
+
+        if self._use_fused_norm:
+            # 2D view of the persistent input buffer; in-place mutation OK
+            # because the buffer is reset every replay (see ``block_forward_graph``).
+            residual = self._g_inputs_embeds.view(N, -1)
+            hidden = self._fi_rmsnorm(
+                residual, self._input_ln_w[0], self._input_ln_eps[0],
+            )
+
+            for li in range(L):
+                qkv = self._project("qkv", li, hidden)
+                qs = qkv[:, :qd].view(N, nh, hd)
+                ks = qkv[:, qd: qd + kvd].view(N, nkv, hd)
+                vs = qkv[:, qd + kvd:].view(N, nkv, hd)
+
+                if use_ri:
+                    _fi_apply_rope(
+                        qs, ks, self._rope_indptr, self._rope_offset,
+                        interleave=False, rope_scale=1.0,
+                        rope_theta=self._rope_theta,
+                    )
+                else:
+                    q3 = qs.view(B, bl, nh, hd)
+                    k3 = ks.view(B, bl, nkv, hd)
+                    q3, k3 = self._apply_rotary_pos_emb(
+                        q3, k3, self._g_cos, self._g_sin, unsqueeze_dim=2,
+                    )
+                    qs = q3.reshape(N, nh, hd)
+                    ks = k3.reshape(N, nkv, hd)
+
+                self._paged_kv[li][
+                    self._g_page_write_idx, 0, self._g_page_write_off
+                ] = ks
+                self._paged_kv[li][
+                    self._g_page_write_idx, 1, self._g_page_write_off
+                ] = vs
+
+                attn_out = self._fi_graph_wrapper.run(
+                    qs, self._paged_kv[li],
+                )
+                attn_out_2d = attn_out.reshape(N, -1)
+                attn_out_2d = self._project(
+                    "o", li, attn_out_2d,
+                ).contiguous()
+
+                self._fi_fused_add_rmsnorm(
+                    attn_out_2d, residual,
+                    self._post_ln_w[li], self._post_ln_eps[li],
+                )
+
+                mlp_out = self._mlp_forward(li, attn_out_2d)
+                if not mlp_out.is_contiguous():
+                    mlp_out = mlp_out.contiguous()
+
+                if li + 1 < L:
+                    self._fi_fused_add_rmsnorm(
+                        mlp_out, residual,
+                        self._input_ln_w[li + 1], self._input_ln_eps[li + 1],
+                    )
+                    hidden = mlp_out
+                else:
+                    residual.add_(mlp_out)
+                    hidden = residual
+
+            out = self._fi_rmsnorm(
+                hidden, self._final_norm_w, self._final_norm_eps,
+            )
+            self._g_output_hidden[:] = out.view(B, bl, -1)
+        else:
+            hidden = self._g_inputs_embeds
+
+            for li in range(L):
+                residual = hidden
+                hidden = self._input_ln[li](hidden)
+
+                qkv = self._project("qkv", li, hidden)
+                q = qkv[:, :, :qd].view(B, bl, nh, hd)
+                k = qkv[:, :, qd: qd + kvd].view(B, bl, nkv, hd)
+                v = qkv[:, :, qd + kvd:].view(B, bl, nkv, hd)
+
+                qs = q.reshape(B * bl, nh, hd)
+                ks = k.reshape(B * bl, nkv, hd)
+                vs = v.reshape(B * bl, nkv, hd)
+
+                if use_ri:
+                    _fi_apply_rope(
+                        qs, ks, self._rope_indptr, self._rope_offset,
+                        interleave=False, rope_scale=1.0,
+                        rope_theta=self._rope_theta,
+                    )
+                else:
+                    cos = self._g_cos
+                    sin = self._g_sin
+                    q, k = self._apply_rotary_pos_emb(
+                        q, k, cos, sin, unsqueeze_dim=2,
+                    )
+                    qs = q.reshape(B * bl, nh, hd)
+                    ks = k.reshape(B * bl, nkv, hd)
+
+                self._paged_kv[li][
+                    self._g_page_write_idx, 0, self._g_page_write_off
+                ] = ks
+                self._paged_kv[li][
+                    self._g_page_write_idx, 1, self._g_page_write_off
+                ] = vs
+
+                attn_out = self._fi_graph_wrapper.run(
+                    qs, self._paged_kv[li],
+                )
+                attn_out = attn_out.reshape(B, bl, -1)
+                attn_out = self._project("o", li, attn_out)
+
+                hidden = residual + attn_out
+                residual = hidden
+                hidden = self._post_ln[li](hidden)
+                hidden = self._mlp_forward(li, hidden)
+                hidden = residual + hidden
+
+            self._g_output_hidden[:] = self._final_norm(hidden)
+
+        if self._has_speech_head:
+            shift_hidden = torch.cat(
+                [self._g_shift_ctx, self._g_output_hidden[:, : bl - 1]],
+                dim=1,
+            )
+            self._g_logits[:] = F.linear(
+                shift_hidden,
+                self._speech_head_w,
+                self._speech_head_b,
+            )
+
+    def set_shift_ctx(self, ctx: Tensor):
+        """Copy shift context into graph buffer (once per block).
+
+        Args:
+            ctx: ``(B, 1, dim)`` or ``(1, 1, dim)`` (broadcast).
+        """
+        if self._g_shift_ctx is not None:
+            if ctx.size(0) < self._batch_size:
+                self._g_shift_ctx[:] = ctx.expand(
+                    self._batch_size, -1, -1,
+                )
+            else:
+                self._g_shift_ctx.copy_(ctx)
+
+    def block_forward_graph(
+        self,
+        block_emb: Tensor,
+        cache_starts: int | List[int],
+    ) -> tuple[Tensor, Tensor | None]:
+        """Batched block forward via CUDA graph replay.
+
+        Falls back to eager :meth:`block_forward` for partial blocks.
+
+        Args:
+            block_emb:    ``(B, bl, dim)``
+            cache_starts: int (uniform) or list[int] per-sample.
+
+        Returns:
+            ``(hidden, logits)`` — ``(B, bl, dim)`` and
+            ``(B, bl, vocab)`` or ``None``.
+        """
+        bl = block_emb.size(1)
+        if bl != self._g_block_size:
+            return self.block_forward(block_emb, cache_starts), None
+
+        self._g_inputs_embeds.copy_(block_emb)
+
+        cs_key: int | tuple
+        if isinstance(cache_starts, int):
+            cs_key = cache_starts
+        else:
+            cs_key = tuple(cache_starts)
+
+        if cs_key != self._last_plan_cache_start:
+            B = self._batch_size
+            bl_ = bl
+
+            pw_idx, pw_off = self._compute_page_write_indices(
+                cache_starts, bl_,
+            )
+            self._g_page_write_idx.copy_(pw_idx)
+            self._g_page_write_off.copy_(pw_off)
+
+            if self._use_rope_inplace:
+                for bi in range(B):
+                    self._rope_indptr[bi + 1] = (bi + 1) * bl_
+                if isinstance(cache_starts, int):
+                    self._rope_offset[:] = cache_starts
+                else:
+                    for bi in range(B):
+                        self._rope_offset[bi] = cache_starts[bi]
+            else:
+                if isinstance(cache_starts, int):
+                    for bi in range(B):
+                        self._g_cos[bi].copy_(
+                            self._rope_cos[0, cache_starts: cache_starts + bl_, :],
+                        )
+                        self._g_sin[bi].copy_(
+                            self._rope_sin[0, cache_starts: cache_starts + bl_, :],
+                        )
+                else:
+                    for bi in range(B):
+                        csb = cache_starts[bi]
+                        self._g_cos[bi].copy_(
+                            self._rope_cos[0, csb: csb + bl_, :],
+                        )
+                        self._g_sin[bi].copy_(
+                            self._rope_sin[0, csb: csb + bl_, :],
+                        )
+
+            if isinstance(cache_starts, int):
+                kv_lens: int | List[int] = cache_starts + bl_
+            else:
+                kv_lens = [c + bl_ for c in cache_starts]
+
+            self._plan_paged_attention(
+                self._fi_graph_wrapper, kv_lens, seq_len=bl_,
+                causal=False,
+                qo_indptr_buf=self._plan_qo_indptr,
+                paged_kv_indptr_buf=self._plan_paged_kv_indptr,
+                paged_kv_indices_buf=self._plan_paged_kv_indices,
+                paged_kv_last_page_len_buf=self._plan_paged_kv_last_page_len,
+            )
+            self._last_plan_cache_start = cs_key
+
+        self._cuda_graph.replay()
+
+        logits = (
+            self._g_logits if self._has_speech_head else None
+        )
+        return self._g_output_hidden, logits

runtime/chatterbox-flash/chatterbox_flash/engines/mlx.py

@@ -0,0 +1,621 @@
+"""MLX (Apple Silicon) inference engine for Chatterbox-Flash.
+
+Runs the LLaMA backbone of the T3 decoder natively on Metal via the
+`mlx <https://github.com/ml-explore/mlx>`_ framework, while keeping the
+small auxiliary modules (speech / text / cond embeddings, speech head,
+conditioning encoder) on PyTorch.  The PyTorch ↔ MLX boundary is at
+the engine surface, so the block-diffusion ``generate()`` loop does not
+need to know about MLX at all.
+
+Memory layout
+-------------
+* T3 conditioning encoder + speech_emb + text_emb + speech_head live
+  on the PyTorch device (CPU on a Mac; the host model is loaded with
+  ``device='cpu'`` for the MLX backend).
+* The LLaMA stack (embed-free, 30 layers for Llama_520M) is rebuilt in
+  MLX, with weights converted once from the PyTorch ``state_dict``.
+* A per-layer KV cache lives entirely inside MLX.
+
+Tentative-then-commit pattern
+-----------------------------
+The block-diffusion loop calls :meth:`block_forward` multiple times for
+the same block (one per inner-loop step), each time overwriting the
+previous step's tentative KV.  We implement this by tracking a
+*committed* length per layer and cropping the cache back to that length
+at the start of each block forward.  :meth:`advance_cache` bumps the
+committed length so the next block sees the just-written entries as
+real prefix.
+
+Status
+------
+Experimental.  Tested with mlx ≥ 0.13 and mlx-lm ≥ 0.10.  This module
+imports MLX lazily; the rest of the package works without MLX installed.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import logging
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from torch import Tensor
+
+logger = logging.getLogger(__name__)
+
+
+MLX_AVAILABLE: bool = (
+    importlib.util.find_spec("mlx") is not None
+    and importlib.util.find_spec("mlx_lm") is not None
+)
+
+
+def _require_mlx() -> tuple["module", "module"]:  # type: ignore[name-defined]
+    if not MLX_AVAILABLE:
+        raise ImportError(
+            "MLX backend requires the 'mlx' and 'mlx-lm' packages.\n"
+            "  Install (Apple Silicon Mac only):\n"
+            "    pip install mlx mlx-lm\n"
+            "  or:\n"
+            "    pip install 'chatterbox-flash[mlx]'",
+        )
+    import mlx.core as mx
+    import mlx_lm  # noqa: F401  (force import, validate version)
+    return mx, mlx_lm
+
+
+# --------------------------------------------------------------------------- #
+#  Weight conversion: PyTorch HF LlamaModel state_dict  →  MLX state_dict
+# --------------------------------------------------------------------------- #
+
+
+def _torch_to_mx(t: torch.Tensor, mx_module):  # noqa: ANN001
+    """Convert a torch tensor to an mx.array via numpy (zero extra copy)."""
+    arr = t.detach().to(torch.float32).cpu().numpy()
+    return mx_module.array(arr)
+
+
+def _convert_llama_weights(
+    tfmr_state: dict[str, torch.Tensor],
+    mx_module,
+) -> dict[str, "mx_module.array"]:  # type: ignore[name-defined]
+    """Convert HF LlamaModel state_dict into an MLX-flat dict.
+
+    Both PyTorch and mlx-lm store nn.Linear as ``(out, in)`` so no
+    transposition is needed.  RoPE inv_freq / cos_cached / sin_cached
+    buffers are recomputed inside MLX and therefore skipped.
+    """
+    SKIP_SUBSTRINGS = (
+        "rotary_emb.inv_freq",
+        "rotary_emb.cos_cached",
+        "rotary_emb.sin_cached",
+    )
+    out: dict[str, object] = {}
+    for k, v in tfmr_state.items():
+        if any(s in k for s in SKIP_SUBSTRINGS):
+            continue
+        out[k] = _torch_to_mx(v, mx_module)
+    return out
+
+
+# --------------------------------------------------------------------------- #
+#  Embed-free MLX LlamaModel wrapper
+# --------------------------------------------------------------------------- #
+
+
+class _EmbedFreeLlama:
+    """Thin wrapper around mlx-lm's LlamaModel that accepts inputs_embeds.
+
+    Reuses the entire mlx-lm LLaMA layer stack (incl. RoPE LLaMA3 scaling
+    that ships with mlx-lm) but bypasses the input ``embed_tokens`` layer
+    — we already have the speech / text / cond embeddings on the
+    PyTorch side.
+    """
+
+    def __init__(
+        self,
+        tfmr_state: dict[str, torch.Tensor],
+        hf_config,  # noqa: ANN001
+        rotary_emb=None,  # noqa: ANN001
+    ) -> None:
+        mx, _ = _require_mlx()
+        from mlx_lm.models.llama import LlamaModel, ModelArgs
+
+        self._mx = mx
+        args = ModelArgs(
+            model_type="llama",
+            hidden_size=int(hf_config.hidden_size),
+            num_hidden_layers=int(hf_config.num_hidden_layers),
+            intermediate_size=int(hf_config.intermediate_size),
+            num_attention_heads=int(hf_config.num_attention_heads),
+            num_key_value_heads=int(hf_config.num_key_value_heads),
+            rms_norm_eps=float(hf_config.rms_norm_eps),
+            vocab_size=int(hf_config.vocab_size),
+            head_dim=int(getattr(hf_config, "head_dim", 0))
+            or int(hf_config.hidden_size // hf_config.num_attention_heads),
+            max_position_embeddings=int(
+                getattr(hf_config, "max_position_embeddings", 4096)
+            ),
+            rope_theta=float(getattr(hf_config, "rope_theta", 10000.0)),
+            rope_traditional=False,
+            rope_scaling=dict(hf_config.rope_scaling)
+            if getattr(hf_config, "rope_scaling", None) is not None
+            else None,
+            tie_word_embeddings=bool(hf_config.tie_word_embeddings),
+        )
+        self._args = args
+        self._llama = LlamaModel(args)
+
+        # Convert weights (skip rotary buffers; mlx-lm rebuilds them).
+        mlx_state = _convert_llama_weights(tfmr_state, mx)
+        from mlx.utils import tree_unflatten
+        self._llama.load_weights(list(mlx_state.items()), strict=False)
+
+        # Override RoPE with HF's *exact* inv_freq. Reconstructing it from the
+        # config (rope_theta / rope_scaling) is fragile: across transformers
+        # versions ``config.rope_theta`` may be absent, in which case the
+        # ModelArgs above silently falls back to base=10000 while HF actually
+        # used e.g. 500000 — giving a position-dependent RoPE mismatch that
+        # scrambles generation. HF's ``rotary_emb.inv_freq`` already bakes in
+        # the base *and* any (llama3 / yarn / …) scaling, so using it verbatim
+        # is correct regardless of config layout or scaling type.
+        self._override_rope_from_hf(rotary_emb)
+
+        # Optional weight-only quantization of the backbone (4- / 8-bit).
+        # Driven by an env var so the tts/model signatures stay untouched
+        # (same convention as CHATTERBOX_FLASH_ENGINE). Must run AFTER
+        # load_weights so the fp weights are present to quantize, and only
+        # touches this LLaMA stack — S3Gen / VoiceEncoder live on the torch
+        # side and are unaffected.
+        self._quant_bits = self._maybe_quantize()
+
+        # Force materialization to catch missing keys early.
+        mx.eval(self._llama.parameters())
+        self._cache_cls = None
+        self._init_cache_cls()
+
+    def _override_rope_from_hf(self, rotary_emb) -> None:  # noqa: ANN001
+        """Replace each layer's RoPE with one driven by HF's exact inv_freq."""
+        mx = self._mx
+        if rotary_emb is None or not hasattr(rotary_emb, "inv_freq"):
+            logger.warning(
+                "MLX backbone: HF rotary_emb/inv_freq unavailable; falling back "
+                "to config-reconstructed RoPE (may mismatch if config lacks "
+                "rope_theta).",
+            )
+            return
+
+        import numpy as _np
+        import mlx.nn as nn
+
+        inv = rotary_emb.inv_freq.detach().to(torch.float32).cpu().numpy()
+        # mx.fast.rope takes per-pair ``freqs`` = wavelengths = 1 / inv_freq.
+        freqs = mx.array((1.0 / inv).astype(_np.float32))
+        att_scale = float(getattr(rotary_emb, "attention_scaling", 1.0) or 1.0)
+        dims = int(self._args.head_dim)
+        traditional = bool(self._args.rope_traditional)
+
+        class _HFRoPE(nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.dims = dims
+                self._freqs = freqs
+                self.scale = att_scale
+                self.traditional = traditional
+
+            def __call__(self, x, offset=0):  # noqa: ANN001
+                y = mx.fast.rope(
+                    x, self.dims, traditional=self.traditional,
+                    base=None, scale=1.0, offset=offset, freqs=self._freqs,
+                )
+                return y if self.scale == 1.0 else y * self.scale
+
+        for layer in self._llama.layers:
+            layer.self_attn.rope = _HFRoPE()
+        logger.info(
+            "MLX backbone RoPE overridden from HF inv_freq "
+            "(dims=%d, attention_scaling=%.4f).", dims, att_scale,
+        )
+
+    def _maybe_quantize(self) -> Optional[int]:
+        import os
+
+        raw = os.environ.get("CHATTERBOX_FLASH_MLX_QUANT_BITS", "").strip()
+        if not raw:
+            return None
+        try:
+            bits = int(raw)
+        except ValueError:
+            logger.warning("Ignoring invalid CHATTERBOX_FLASH_MLX_QUANT_BITS=%r", raw)
+            return None
+        if bits not in (2, 3, 4, 6, 8):
+            logger.warning("Unsupported MLX quant bits=%d; skipping.", bits)
+            return None
+        group_size = int(os.environ.get("CHATTERBOX_FLASH_MLX_QUANT_GROUP", "64"))
+
+        import mlx.nn as nn
+
+        # Skip layers whose dims aren't divisible by group_size — mlx requires
+        # in_features % group_size == 0, and the speech/text heads live on the
+        # torch side anyway, so any odd-shaped Linear here is left in fp.
+        def _can_quant(_path: str, module) -> bool:  # noqa: ANN001
+            if not hasattr(module, "to_quantized"):
+                return False
+            w = getattr(module, "weight", None)
+            return w is not None and w.ndim == 2 and (w.shape[-1] % group_size == 0)
+
+        nn.quantize(
+            self._llama, group_size=group_size, bits=bits,
+            class_predicate=_can_quant,
+        )
+        logger.info(
+            "MLX backbone quantized to %d-bit (group_size=%d).", bits, group_size,
+        )
+        return bits
+
+    def _init_cache_cls(self) -> None:
+        try:
+            from mlx_lm.models.cache import KVCache
+            self._cache_cls = KVCache
+        except ImportError:  # very old mlx-lm
+            from mlx_lm.models.base import KVCache  # type: ignore
+            self._cache_cls = KVCache
+
+    def make_cache(self) -> list:
+        return [self._cache_cls() for _ in range(self._args.num_hidden_layers)]
+
+    def forward(
+        self,
+        inputs_embeds_mx,
+        cache: list,
+        *,
+        mask=None,
+    ):
+        """Run the LLaMA stack on continuous embeddings, return hidden_states."""
+        mx = self._mx
+        x = inputs_embeds_mx
+        # Manually mirror mlx-lm LlamaModel.__call__ but without embed_tokens.
+        for layer, c in zip(self._llama.layers, cache):
+            x = layer(x, mask=mask, cache=c)
+        x = self._llama.norm(x)
+        return x
+
+
+# --------------------------------------------------------------------------- #
+#  Engine
+# --------------------------------------------------------------------------- #
+
+
+CacheStarts = Union[int, List[int]]
+
+
+class MLXEngine:
+    """Apple-Silicon MLX engine. Implements the chatterbox_flash engine protocol."""
+
+    has_cuda_graph: bool = False  # MLX has no CUDA-graph concept
+
+    def __init__(
+        self,
+        model,  # ChatterboxFlashT3
+        max_seq_len: int,
+        dtype: torch.dtype,
+        *,
+        batch_size: int = 1,
+        page_size: int = 16,  # unused; protocol parity
+    ) -> None:
+        self._model = model
+        self._max_seq_len = int(max_seq_len)
+        self._dtype = dtype
+        self._batch_size = max(1, int(batch_size))
+        self._device = model.device
+
+        mx, _ = _require_mlx()
+        self._mx = mx
+        if dtype == torch.bfloat16:
+            self._mx_dtype = mx.bfloat16
+        elif dtype == torch.float16:
+            self._mx_dtype = mx.float16
+        else:
+            self._mx_dtype = mx.float32
+
+        # Build the MLX backbone exactly once per engine instance.
+        tfmr_state = {k: v for k, v in model.tfmr.state_dict().items()}
+        self._backbone = _EmbedFreeLlama(
+            tfmr_state, model.tfmr.config,
+            rotary_emb=getattr(model.tfmr, "rotary_emb", None),
+        )
+        self._cache = self._backbone.make_cache()
+
+        # Per-row batch-1 caches, populated only on the ragged (cross-text /
+        # per-row prefix length) path.  A single dense batch-B cache cannot
+        # represent rows with differing sequence lengths, so the ragged path
+        # serializes over these instead of slicing ``self._cache`` live.
+        self._row_caches: Optional[List[list]] = None
+
+        # Per-row bookkeeping.  MLX's KV cache is shared across the batch
+        # dim (same as the SDPA path) so this is uniform across rows.
+        self._cache_lens: List[int] = [0] * self._batch_size
+        self._prefix_lens: List[int] = [0] * self._batch_size
+
+    # ------------------------------------------------------------------ #
+    #  Introspection / lifecycle
+    # ------------------------------------------------------------------ #
+
+    @property
+    def batch_size(self) -> int:
+        return self._batch_size
+
+    @property
+    def cache_len(self) -> int:
+        return self._cache_lens[0]
+
+    @cache_len.setter
+    def cache_len(self, value: int) -> None:
+        self._cache_lens = [int(value)] * self._batch_size
+
+    @property
+    def cache_lens(self) -> List[int]:
+        return self._cache_lens
+
+    @property
+    def prefix_lens(self) -> List[int]:
+        return self._prefix_lens
+
+    def reset(self) -> None:
+        self._cache = self._backbone.make_cache()
+        self._row_caches = None
+        self._cache_lens = [0] * self._batch_size
+        self._prefix_lens = [0] * self._batch_size
+
+    def can_reuse(
+        self,
+        max_seq_len: int,
+        dtype: torch.dtype,
+        *,
+        batch_size: int = 1,
+        page_size: int = 16,
+    ) -> bool:
+        return (
+            self._max_seq_len >= max_seq_len
+            and self._dtype == dtype
+            and self._batch_size == max(1, batch_size)
+        )
+
+    # ------------------------------------------------------------------ #
+    #  Tensor boundary
+    # ------------------------------------------------------------------ #
+
+    def _torch_to_mx(self, t: Tensor):
+        """torch (cpu / mps / cuda) → mx.array (Metal). Routes via numpy on host."""
+        return self._mx.array(t.detach().to(torch.float32).cpu().numpy()).astype(
+            self._mx_dtype,
+        )
+
+    def _mx_to_torch(self, a) -> Tensor:
+        """mx.array → torch on the host device, restoring our dtype."""
+        # mx.array → numpy → torch (always lands on the model's device).
+        arr = np.array(a.astype(self._mx.float32))
+        return torch.from_numpy(arr).to(device=self._device, dtype=self._dtype)
+
+    # ------------------------------------------------------------------ #
+    #  KV-cache cropping (tentative-then-commit pattern)
+    # ------------------------------------------------------------------ #
+
+    def _crop_cache(self, cache: list, length: int) -> None:
+        """Truncate every layer of ``cache`` to ``length`` time steps."""
+        mx = self._mx
+        for c in cache:
+            # mlx-lm's KVCache exposes ``keys`` / ``values`` (shape
+            # ``(B, n_kv_h, T, head_dim)``) and an ``offset`` int.
+            if c.keys is None:
+                continue
+            c.keys = c.keys[..., :length, :]
+            c.values = c.values[..., :length, :]
+            c.offset = int(length)
+        mx.eval(*[(c.keys, c.values) for c in cache if c.keys is not None])
+
+    def _crop_cache_to(self, length: int) -> None:
+        """Truncate the shared (uniform-path) KV cache to ``length`` steps."""
+        self._crop_cache(self._cache, length)
+
+    # ------------------------------------------------------------------ #
+    #  Prefix forward
+    # ------------------------------------------------------------------ #
+
+    def prefix_forward(
+        self,
+        prefix_emb_or_list: Union[Tensor, List[Tensor]],
+    ) -> Tensor:
+        """Run the cond/text/SOS prefix once and seed the cache."""
+        if isinstance(prefix_emb_or_list, list):
+            return self._prefix_forward_ragged(prefix_emb_or_list)
+
+        prefix_emb = prefix_emb_or_list
+        B_pfx = prefix_emb.size(0)
+        if B_pfx not in (1, self._batch_size):
+            raise ValueError(
+                f"MLX prefix_forward: prefix batch {B_pfx} must be 1 or "
+                f"engine batch {self._batch_size}",
+            )
+        if B_pfx == 1 and self._batch_size > 1:
+            prefix_emb = prefix_emb.expand(self._batch_size, -1, -1).contiguous()
+
+        lp = int(prefix_emb.size(1))
+        # MLX expects a 3D input (B, T, D).
+        x_mx = self._torch_to_mx(prefix_emb)
+        # The conditioning prefix must be processed CAUSALLY (matching the HF
+        # SDPA path, which gets an implicit causal mask). _EmbedFreeLlama
+        # bypasses LlamaModel.__call__, so mask=None would leave the prefix
+        # fully bidirectional and corrupt the cached context. Use mlx's native
+        # "causal" string mask — exactly what mlx-lm passes for a causal
+        # forward — so it is dtype-safe regardless of the backbone weight dtype.
+        out_mx = self._backbone.forward(x_mx, self._cache, mask="causal")
+        self._mx.eval(out_mx)
+
+        for bi in range(self._batch_size):
+            self._cache_lens[bi] = lp
+            self._prefix_lens[bi] = lp
+
+        # SOS is at index lp - 1.
+        shift_ctx = self._mx_to_torch(out_mx[:, lp - 1 : lp, :])
+        return shift_ctx.contiguous()
+
+    def _prefix_forward_ragged(self, emb_list: List[Tensor]) -> Tensor:
+        """Pad-and-mask path for per-row prefix lengths (cross-text / null prefix)."""
+        if len(emb_list) == 1:
+            return self.prefix_forward(emb_list[0])
+
+        mx = self._mx
+        B = len(emb_list)
+        if B != self._batch_size:
+            raise ValueError(
+                f"MLX ragged prefix: list length {B} != engine batch "
+                f"{self._batch_size}",
+            )
+        lp_list = [int(e.size(1)) for e in emb_list]
+        lp_max = max(lp_list)
+        dim = emb_list[0].size(-1)
+
+        padded = torch.zeros(B, lp_max, dim, dtype=self._dtype, device=self._device)
+        for bi, e in enumerate(emb_list):
+            padded[bi, : e.size(1)] = e.squeeze(0).to(self._dtype)
+
+        # Additive padding mask, combined with causal triangular mask.
+        x_mx = self._torch_to_mx(padded)
+        neg_inf = mx.array(-1e9, dtype=self._mx_dtype)
+        pad_mask = mx.zeros((B, lp_max), dtype=self._mx_dtype)
+        for bi, lp_i in enumerate(lp_list):
+            if lp_i < lp_max:
+                pad_mask[bi, lp_i:] = neg_inf
+        causal = mx.triu(
+            mx.full((lp_max, lp_max), neg_inf, dtype=self._mx_dtype),
+            k=1,
+        )
+        attn_mask = pad_mask[:, None, None, :] + causal[None, None]
+        out_mx = self._backbone.forward(x_mx, self._cache, mask=attn_mask)
+        mx.eval(out_mx)
+
+        for bi, lp_i in enumerate(lp_list):
+            self._cache_lens[bi] = lp_i
+            self._prefix_lens[bi] = lp_i
+
+        # Split the padded batch-B prefix cache into per-row batch-1 caches,
+        # each trimmed to that row's *valid* prefix length (the padded tail
+        # carries masked garbage and must be dropped).  The ragged
+        # ``block_forward`` serializes over these so a batch-1 block input
+        # never gets concatenated against batch-B cache keys.
+        row_caches: List[list] = []
+        for bi, lp_i in enumerate(lp_list):
+            rc = self._backbone.make_cache()
+            for src, dst in zip(self._cache, rc):
+                if src.keys is None:
+                    continue
+                dst.keys = src.keys[bi : bi + 1, :, :lp_i, :]
+                dst.values = src.values[bi : bi + 1, :, :lp_i, :]
+                dst.offset = int(lp_i)
+            mx.eval(*[(c.keys, c.values) for c in rc if c.keys is not None])
+            row_caches.append(rc)
+        self._row_caches = row_caches
+
+        # SOS row index per sample (lp_i - 1).  Gather on the torch side.
+        hidden = self._mx_to_torch(out_mx)
+        idx = torch.tensor(
+            [lp_i - 1 for lp_i in lp_list], device=self._device, dtype=torch.long,
+        )
+        shift_ctx = hidden.gather(
+            1, idx[:, None, None].expand(B, 1, hidden.size(-1)),
+        )
+        return shift_ctx.contiguous()
+
+    # ------------------------------------------------------------------ #
+    #  Block forward
+    # ------------------------------------------------------------------ #
+
+    def block_forward(
+        self,
+        block_emb: Tensor,
+        cache_starts: CacheStarts,
+    ) -> Tensor:
+        """Tentative block forward; KV is written, advance_cache commits it."""
+        B = int(block_emb.size(0))
+        bl = int(block_emb.size(1))
+
+        if isinstance(cache_starts, int):
+            cs_uniform: Optional[int] = cache_starts
+        else:
+            cs_list = list(cache_starts)
+            if len(set(cs_list)) == 1:
+                cs_uniform = cs_list[0]
+            else:
+                cs_uniform = None
+
+        if cs_uniform is not None:
+            self._crop_cache_to(int(cs_uniform))
+            x_mx = self._torch_to_mx(block_emb)
+            # mlx-lm builds a causal mask internally when T > 1 and the
+            # cache has an offset; the block sees the cropped prefix +
+            # itself.  No explicit mask needed here.
+            out_mx = self._backbone.forward(x_mx, self._cache)
+            self._mx.eval(out_mx)
+            return self._mx_to_torch(out_mx)
+
+        # Ragged: per-row cache_starts differ (cross-text / null-prefix).
+        # mlx-lm exposes no per-row position offset, and a dense batch-B
+        # cache cannot hold rows of differing length, so we serialize over
+        # the per-row batch-1 caches seeded at prefix time.  Each row's
+        # cache is cropped to that row's committed length (cs_i), the block
+        # is forwarded tentatively, and the write persists in the row cache
+        # for the next step / commit — mirroring the uniform path per row.
+        # We deliberately serialize to keep correctness > throughput.
+        mx = self._mx
+        if self._row_caches is None:
+            raise RuntimeError(
+                "MLX ragged block_forward requires per-row caches; "
+                "prefix_forward must run with a per-row prefix list first.",
+            )
+        outs = []
+        for bi, cs_i in enumerate(cs_list):  # type: ignore[name-defined]
+            rc = self._row_caches[bi]
+            self._crop_cache(rc, int(cs_i))
+            x_mx_i = self._torch_to_mx(block_emb[bi : bi + 1])
+            out_mx_i = self._backbone.forward(x_mx_i, rc)
+            mx.eval(out_mx_i)
+            outs.append(self._mx_to_torch(out_mx_i))
+        return torch.cat(outs, dim=0)
+
+    def block_forward_graph(
+        self,
+        block_emb: Tensor,
+        cache_starts: CacheStarts,
+    ) -> Tuple[Tensor, Optional[Tensor]]:
+        """No graph capture on MLX → delegate to eager."""
+        return self.block_forward(block_emb, cache_starts), None
+
+    def advance_cache(self, block_len: int) -> None:
+        """Commit ``block_len`` newly-written entries."""
+        for bi in range(self._batch_size):
+            self._cache_lens[bi] += int(block_len)
+
+    def set_shift_ctx(self, ctx: Tensor) -> None:  # noqa: D401
+        """No graph buffer to stage; the caller folds shift_ctx in-line."""
+        return None
+
+    def capture_cuda_graph(
+        self,
+        block_size: int,
+        *,
+        speech_head: Optional[torch.nn.Linear] = None,
+    ) -> None:  # noqa: D401
+        """MLX has no CUDA-graph concept; intentionally a no-op."""
+        logger.debug(
+            "MLXEngine.capture_cuda_graph(block_size=%d) is a no-op.",
+            block_size,
+        )
+        return None
+
+
+__all__ = [
+    "MLXEngine",
+    "MLX_AVAILABLE",
+]

runtime/chatterbox-flash/chatterbox_flash/engines/torch_sdpa.py

@@ -0,0 +1,372 @@
+"""Pure-PyTorch fallback inference engine for Chatterbox-Flash.
+
+Implements the same engine protocol as :class:`FlashInferEngine` but on
+top of the stock HuggingFace transformers ``LlamaModel`` SDPA backend.
+No CUDA-graph capture, no paged buffer, no flashinfer kernels — just
+``transformers.DynamicCache`` with the standard causal/bidirectional
+attention masks.
+
+This is the engine selected automatically when ``flashinfer-python`` is
+not importable, e.g. on Ampere/Turing GPUs or in environments where the
+flashinfer wheels aren't available.  It produces numerically matching
+outputs (up to fp32/bf16 reduction-order differences) at lower
+throughput.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+logger = logging.getLogger(__name__)
+
+
+CacheStarts = Union[int, List[int]]
+
+
+def _crop_dynamic_cache(cache, max_length: int) -> None:
+    """Truncate a :class:`transformers.DynamicCache` to ``max_length`` tokens.
+
+    Compatible with both the new ``cache.crop()`` API (transformers ≥ 4.46)
+    and the older direct-tensor layout.
+    """
+    if max_length is None:
+        return
+    if hasattr(cache, "crop"):
+        cache.crop(max_length)
+    else:
+        for li in range(len(cache.key_cache)):
+            cache.key_cache[li] = cache.key_cache[li][:, :, :max_length, :]
+            cache.value_cache[li] = cache.value_cache[li][:, :, :max_length, :]
+        cache._seen_tokens = max_length
+
+
+class TorchSDPAEngine:
+    """Pure-PyTorch SDPA + :class:`DynamicCache` fallback engine."""
+
+    has_cuda_graph: bool = False
+
+    def __init__(
+        self,
+        model,
+        max_seq_len: int,
+        dtype: torch.dtype,
+        *,
+        batch_size: int = 1,
+        page_size: int = 16,  # ignored; protocol parity
+    ) -> None:
+        from transformers import DynamicCache  # local import; lazy
+
+        self._model = model
+        self._max_seq_len = int(max_seq_len)
+        self._dtype = dtype
+        self._batch_size = max(1, int(batch_size))
+        self._device = model.device
+
+        self._DynamicCache = DynamicCache
+
+        self._dc = DynamicCache()
+        # ``cache_lens[bi]`` is the *committed* (advance_cache-ed) length per
+        # sample.  In the SDPA path the underlying DynamicCache is uniform
+        # across the batch, so the per-sample bookkeeping is mostly for
+        # protocol parity with FlashInfer; the per-sample length is read only
+        # by capacity checks / introspection.
+        self._cache_lens: List[int] = [0] * self._batch_size
+        self._prefix_lens: List[int] = [0] * self._batch_size
+
+    # ------------------------------------------------------------------ #
+    #  Introspection / lifecycle
+    # ------------------------------------------------------------------ #
+
+    @property
+    def batch_size(self) -> int:
+        return self._batch_size
+
+    @property
+    def cache_len(self) -> int:
+        """Committed cache length (uniform across the batch under SDPA)."""
+        return self._cache_lens[0]
+
+    @cache_len.setter
+    def cache_len(self, value: int) -> None:
+        self._cache_lens = [int(value)] * self._batch_size
+
+    @property
+    def cache_lens(self) -> List[int]:
+        return self._cache_lens
+
+    @property
+    def prefix_lens(self) -> List[int]:
+        return self._prefix_lens
+
+    def reset(self) -> None:
+        self._dc = self._DynamicCache()
+        self._cache_lens = [0] * self._batch_size
+        self._prefix_lens = [0] * self._batch_size
+
+    def can_reuse(
+        self,
+        max_seq_len: int,
+        dtype: torch.dtype,
+        *,
+        batch_size: int = 1,
+        page_size: int = 16,
+    ) -> bool:
+        return (
+            self._max_seq_len >= max_seq_len
+            and self._dtype == dtype
+            and self._batch_size == max(1, batch_size)
+        )
+
+    def snapshot_kv(self):
+        """Lightweight snapshot for CFG uncond restore."""
+        ks = [t.clone() for t in self._dc.key_cache]
+        vs = [t.clone() for t in self._dc.value_cache]
+        return (ks, vs, list(self._cache_lens), list(self._prefix_lens))
+
+    def restore_kv(self, snap) -> None:
+        ks, vs, cl, pl = snap
+        self._dc = self._DynamicCache()
+        for k, v in zip(ks, vs):
+            self._dc.key_cache.append(k.clone())
+            self._dc.value_cache.append(v.clone())
+        self._dc._seen_tokens = ks[0].shape[-2] if ks else 0
+        self._cache_lens = list(cl)
+        self._prefix_lens = list(pl)
+
+    # ------------------------------------------------------------------ #
+    #  Prefix forward
+    # ------------------------------------------------------------------ #
+
+    def prefix_forward(
+        self,
+        prefix_emb_or_list: Union[Tensor, List[Tensor]],
+    ) -> Tensor:
+        """Run the cond/text/SOS prefix once and seed the cache."""
+        if isinstance(prefix_emb_or_list, list):
+            return self._prefix_forward_ragged(prefix_emb_or_list)
+
+        prefix_emb = prefix_emb_or_list
+        B_pfx = prefix_emb.size(0)
+        if B_pfx not in (1, self._batch_size):
+            raise ValueError(
+                f"SDPA prefix_forward: prefix batch {B_pfx} must be 1 or "
+                f"engine batch {self._batch_size}",
+            )
+        if B_pfx == 1 and self._batch_size > 1:
+            prefix_emb = prefix_emb.expand(self._batch_size, -1, -1).contiguous()
+
+        lp = prefix_emb.size(1)
+        device = self._device
+        pos_ids = (
+            torch.arange(lp, device=device)
+            .unsqueeze(0)
+            .expand(self._batch_size, -1)
+            .contiguous()
+        )
+
+        cfg = self._model.tfmr.config
+        old_impl = getattr(cfg, "_attn_implementation", "sdpa")
+        cfg._attn_implementation = "sdpa"
+        try:
+            out = self._model.tfmr(
+                input_ids=None,
+                inputs_embeds=prefix_emb.to(self._dtype),
+                position_ids=pos_ids,
+                past_key_values=self._dc,
+                use_cache=True,
+                return_dict=True,
+            )
+        finally:
+            cfg._attn_implementation = old_impl
+
+        self._dc = out.past_key_values
+        for bi in range(self._batch_size):
+            self._cache_lens[bi] = lp
+            self._prefix_lens[bi] = lp
+
+        shift_ctx = out.last_hidden_state[:, lp - 1 : lp, :].contiguous()
+        return shift_ctx
+
+    def _prefix_forward_ragged(self, emb_list: List[Tensor]) -> Tensor:
+        """Cross-text path: pad to max length, mask out padding in attention.
+
+        Each sample gets its own prefix length; the cache is built from a
+        left-padded forward.  Padding tokens are attention-masked so they
+        contribute zero KV, but the cache slots still exist — downstream
+        block forwards know to skip them via per-sample ``cache_starts``.
+        """
+        if len(emb_list) == 1:
+            return self.prefix_forward(emb_list[0])
+
+        B = len(emb_list)
+        if B != self._batch_size:
+            raise ValueError(
+                f"SDPA ragged prefix: list length {B} != engine batch "
+                f"{self._batch_size}",
+            )
+        lp_list = [int(e.size(1)) for e in emb_list]
+        lp_max = max(lp_list)
+        dim = emb_list[0].size(-1)
+        device = self._device
+
+        padded = torch.zeros(B, lp_max, dim, device=device, dtype=self._dtype)
+        for bi, e in enumerate(emb_list):
+            padded[bi, : e.size(1)] = e.to(self._dtype).squeeze(0)
+        # Additive mask: 0 for visible, -inf for padding (additive on attn scores).
+        # transformers SDPA expects 4D mask (B, 1, Q, K).
+        pad_mask = torch.zeros(B, lp_max, device=device, dtype=self._dtype)
+        for bi, lp_i in enumerate(lp_list):
+            if lp_i < lp_max:
+                pad_mask[bi, lp_i:] = float("-inf")
+        attn_mask = pad_mask[:, None, None, :].expand(B, 1, lp_max, lp_max)
+        # Combine with causal: HF SDPA does causal internally if attention_mask
+        # is 2D (B, K); for the 4D additive mask path we have to OR in causal.
+        causal = torch.triu(
+            torch.full((lp_max, lp_max), float("-inf"), device=device, dtype=self._dtype),
+            diagonal=1,
+        )
+        attn_mask = attn_mask + causal[None, None]
+
+        pos_ids = (
+            torch.arange(lp_max, device=device).unsqueeze(0).expand(B, -1).contiguous()
+        )
+
+        cfg = self._model.tfmr.config
+        old_impl = getattr(cfg, "_attn_implementation", "sdpa")
+        cfg._attn_implementation = "sdpa"
+        try:
+            out = self._model.tfmr(
+                input_ids=None,
+                inputs_embeds=padded,
+                position_ids=pos_ids,
+                past_key_values=self._dc,
+                use_cache=True,
+                return_dict=True,
+                attention_mask=attn_mask,
+            )
+        finally:
+            cfg._attn_implementation = old_impl
+
+        self._dc = out.past_key_values
+        for bi, lp_i in enumerate(lp_list):
+            self._cache_lens[bi] = lp_i
+            self._prefix_lens[bi] = lp_i
+
+        # SOS hidden lives at index lp_i - 1 within each row.
+        idx = torch.tensor(
+            [lp_i - 1 for lp_i in lp_list], device=device, dtype=torch.long,
+        )
+        shift_ctx = out.last_hidden_state.gather(
+            1,
+            idx[:, None, None].expand(B, 1, out.last_hidden_state.size(-1)),
+        )
+        return shift_ctx.contiguous()
+
+    # ------------------------------------------------------------------ #
+    #  Block forward
+    # ------------------------------------------------------------------ #
+
+    def block_forward(
+        self,
+        block_emb: Tensor,
+        cache_starts: CacheStarts,
+    ) -> Tensor:
+        """Transient block forward; KV is written but not committed.
+
+        Caller must re-invoke with the same ``cache_starts`` between
+        unmasking steps; the cache is re-cropped each time so the new
+        block's keys/values overwrite the previous step's tentative KV.
+        """
+        B = block_emb.size(0)
+        bl = block_emb.size(1)
+        device = self._device
+        dtype = self._dtype
+
+        if isinstance(cache_starts, int):
+            cs_uniform: Optional[int] = cache_starts
+        else:
+            cs_list = list(cache_starts)
+            if len(set(cs_list)) == 1:
+                cs_uniform = cs_list[0]
+            else:
+                cs_uniform = None
+
+        if cs_uniform is not None:
+            _crop_dynamic_cache(self._dc, cs_uniform)
+            pos_ids = (
+                torch.arange(cs_uniform, cs_uniform + bl, device=device)
+                .unsqueeze(0).expand(B, -1).contiguous()
+            )
+            kv_len = cs_uniform + bl
+            # Full-visible mask within the block (already-causal prefix in cache).
+            attn_mask = torch.zeros((1, 1, bl, kv_len), dtype=dtype, device=device)
+        else:
+            cs_max = max(cs_list)
+            _crop_dynamic_cache(self._dc, cs_max)
+            # Pad each row's positions to cs_max + bl; mask out the prefix
+            # slots beyond cs_i so the block forward only attends to that
+            # sample's true prefix.
+            pos_rows = torch.zeros(B, bl, device=device, dtype=torch.long)
+            for bi, cs in enumerate(cs_list):
+                pos_rows[bi] = torch.arange(cs, cs + bl, device=device)
+            pos_ids = pos_rows
+            kv_len = cs_max + bl
+            attn_mask = torch.zeros(B, 1, bl, kv_len, dtype=dtype, device=device)
+            for bi, cs in enumerate(cs_list):
+                if cs < cs_max:
+                    attn_mask[bi, 0, :, cs:cs_max] = float("-inf")
+
+        cfg = self._model.tfmr.config
+        old_impl = getattr(cfg, "_attn_implementation", "sdpa")
+        cfg._attn_implementation = "sdpa"
+        try:
+            out = self._model.tfmr(
+                input_ids=None,
+                inputs_embeds=block_emb.to(dtype),
+                position_ids=pos_ids,
+                past_key_values=self._dc,
+                use_cache=True,
+                return_dict=True,
+                attention_mask=attn_mask,
+            )
+        finally:
+            cfg._attn_implementation = old_impl
+
+        self._dc = out.past_key_values
+        return out.last_hidden_state
+
+    def block_forward_graph(
+        self,
+        block_emb: Tensor,
+        cache_starts: CacheStarts,
+    ) -> Tuple[Tensor, Optional[Tensor]]:
+        """SDPA has no CUDA-graph fast path → delegate to eager."""
+        return self.block_forward(block_emb, cache_starts), None
+
+    def advance_cache(self, block_len: int) -> None:
+        """Commit ``block_len`` newly-written entries; bump per-sample lens."""
+        for bi in range(self._batch_size):
+            self._cache_lens[bi] += int(block_len)
+
+    def set_shift_ctx(self, ctx: Tensor) -> None:  # noqa: D401
+        """No-op; the SDPA path computes the shift in-line in the caller."""
+        return None
+
+    def capture_cuda_graph(
+        self,
+        block_size: int,
+        *,
+        speech_head: Optional[torch.nn.Linear] = None,
+    ) -> None:  # noqa: D401
+        """No CUDA-graph support; intentionally a no-op."""
+        logger.debug(
+            "TorchSDPAEngine.capture_cuda_graph(block_size=%d) is a no-op "
+            "(SDPA backend has no CUDA-graph fast path).",
+            block_size,
+        )
+        return None

runtime/chatterbox-flash/chatterbox_flash/model.py

@@ -0,0 +1,936 @@
+"""``ChatterboxFlashT3`` — block-diffusion T3 decoder for Chatterbox-Flash.
+
+We extend Chatterbox-TTS's :class:`~chatterbox.models.t3.t3.T3` decoder with a
+single extra ``[MASK]`` token (``id = V``) added to the speech embedding, and
+replace the autoregressive ``inference()`` path with a block-diffusion
+``generate()`` method.
+
+The decoder is otherwise unchanged: same Llama backbone, same conditioning
+encoder, same text / speech heads.  This keeps the released checkpoint
+compatible with the public ``chatterbox-tts`` weights for everything but the
+expanded ``speech_emb`` (which is initialized by copying the V rows from the
+pretrained checkpoint and adding one normal-initialized row for ``[MASK]``).
+
+The ``generate()`` method here implements the production configuration
+described in the paper:
+
+* dual-stream-free single-stream KV cache (``[prefix | speech]``);
+* FlashInfer paged KV cache + CUDA-graph capture of the per-block forward;
+* zero-text-batch classifier-free guidance with ``zero_all`` null prefix
+  (text and conditioning zeroed, speech-x_t duplicated); the only
+  PMI/CFG combination is ``pmi_cfg`` — ``(1+w)·pmi_c − w·pmi_u`` on the
+  PMI scale, no other mode exposed;
+* **pmi_count_early** outlier rule — OmniVoice's r_n count schedule
+  (controlled by ``omnivoice_schedule_t_shift``) as the per-step
+  **floor**, plus the paper's time-shifted PMI-quantile
+  ``q_k = max(0, 1 - τ · (k+1)/K)`` as an early-decoding **ceiling**.
+  When ``time_shift_tau == 0`` the schedule reduces to plain pmi_count
+  (fixed-K, no early decoding);
+* prior-calibrated PMI scoring
+  ``s_i = log p_i(\\hat x_i) - log p̄(\\hat x_i)`` with the **precomputed
+  unconditional block prior** ``p̄`` — a forward on
+  ``[SOS, MASK × block_size]`` with no conditioning, lazy-cached on the
+  model object so it's evaluated at most once per
+  ``(block_size, dtype, device)`` triple over the lifetime of the model;
+* optional Gumbel ``position_temperature`` perturbing the PMI ranking
+  for stochastic top-k position selection (OmniVoice §3.4 uses T=5).
+
+Training / experimental code paths (full block-diffusion forward, dual-stream,
+speaker probe, magi attention, sdpa/flex inference backends, alternative CFG
+modes and outlier methods, remasking, …) are intentionally omitted.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+import math
+from typing import Literal, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from chatterbox.models.t3.t3 import T3 as _ChatterboxT3
+from chatterbox.models.t3.modules.cond_enc import T3Cond
+from chatterbox.models.t3.modules.t3_config import T3Config
+
+from .cfg_guidance import apply_zero_text_cfg_from_logits, pmi_cfg_combine
+from .engines import FLASHINFER_AVAILABLE, InferenceEngine, build_engine
+from .calibration import compute_pmi
+
+
+logger = logging.getLogger(__name__)
+
+
+# --------------------------------------------------------------------------- #
+# Sampling helpers
+# --------------------------------------------------------------------------- #
+
+
+def _gumbel_sample(logits: Tensor, temperature: float) -> Tensor:
+    scaled = logits / temperature
+    u = torch.rand_like(scaled)
+    gumbel = -torch.log(-torch.log(u + 1e-10) + 1e-10)
+    return scaled + gumbel
+
+
+def _sample_at_temperature(
+    logits: Tensor,
+    temperature: float,
+    *,
+    sampling: Literal["multinomial", "gumbel"] = "multinomial",
+) -> Tensor:
+    if sampling == "gumbel":
+        return _gumbel_sample(logits, temperature).argmax(dim=-1)
+    probs = F.softmax(logits / temperature, dim=-1)
+    return torch.multinomial(probs, num_samples=1).squeeze(-1)
+
+
+# --------------------------------------------------------------------------- #
+# Prefix-shaping helpers (CFG ``zero_all`` null branch)
+# --------------------------------------------------------------------------- #
+
+
+def _cond_emb_zero_all(ce: Tensor) -> Tensor:
+    """Zero the entire conditioning prefix (paper's ``zero_all`` null branch)."""
+    return torch.zeros_like(ce)
+
+
+def _zero_text_content_keep_pad(
+    text_emb: Tensor,
+    text_tokens: Tensor,
+    text_token_lens: Tensor | None,
+) -> Tensor:
+    """Null-text replacement: non-pad text positions go to 0, pad positions copy.
+
+    Pad positions are detected from ``text_tokens == -100`` and / or
+    ``text_token_lens`` when supplied. This keeps the sequence length identical
+    so the CFG forward can be batched 2x.
+    """
+    squeeze_batch = text_tokens.ndim == 1
+    toks = text_tokens.unsqueeze(0) if squeeze_batch else text_tokens
+    emb = text_emb.unsqueeze(0) if text_emb.ndim == 2 else text_emb
+    Btxt, Tt = toks.shape[0], toks.size(1)
+    device = toks.device
+    is_pad = toks.eq(-100)
+    if text_token_lens is not None:
+        ttl = text_token_lens.to(device=device, dtype=torch.long).reshape(-1)
+        if ttl.numel() == 1 and Btxt > 1:
+            ttl = ttl.expand(Btxt)
+        is_pad = is_pad | (torch.arange(Tt, device=device)[None, :] >= ttl[:, None])
+    out = emb.clone()
+    out.masked_fill_((~is_pad).unsqueeze(-1), 0.0)
+    return out.squeeze(0) if squeeze_batch else out
+
+
+# --------------------------------------------------------------------------- #
+# OmniVoice §3.4 r_n count schedule
+# --------------------------------------------------------------------------- #
+
+
+def _omnivoice_unmask_schedule(
+    n_total_mask: int,
+    num_steps: int,
+    t_shift: float,
+) -> list[int]:
+    """Per-step unmask **counts** under OmniVoice's time-shifted ``r_n`` curve.
+
+    Cumulative fraction unmasked by the end of step ``s`` is
+    ``t_s = r_n(s/K, τ)`` where
+    ``r_n(s, τ) = τ * s / (1 + (τ - 1) * s)`` (eq. (3) of OmniVoice).
+    We then convert the cumulative target into per-step integer counts by
+    cumulative rounding (``target = round(N * t_{s+1})``,
+    ``num_s = max(0, target - cum)``), capped at the remaining masks.
+    The last step always closes out whatever is left.
+
+    Returns a list of length ``num_steps`` summing to ``n_total_mask``.
+    """
+    K = max(1, int(num_steps))
+    N = max(0, int(n_total_mask))
+    if N == 0:
+        return [0] * K
+    if K == 1:
+        return [N]
+    tau = float(t_shift)
+    if tau <= 0:
+        # Degenerate: split as evenly as possible, last step takes remainder.
+        base = N // K
+        out = [base] * K
+        out[-1] += N - sum(out)
+        return out
+    ts = [tau * (s / K) / (1.0 + (tau - 1.0) * (s / K)) for s in range(K + 1)]
+    counts: list[int] = []
+    cum = 0
+    rem = N
+    for s in range(K):
+        if s == K - 1:
+            num = rem
+        else:
+            target = int(round(N * ts[s + 1]))
+            num = max(0, target - cum)
+            num = min(num, rem)
+        counts.append(num)
+        cum += num
+        rem -= num
+    return counts
+
+
+# --------------------------------------------------------------------------- #
+# Speech-block step (block-diffusion inner loop)
+# --------------------------------------------------------------------------- #
+
+
+def _pmi_count_early_step_unmask(
+    blk_logits_2d: Tensor,
+    blk_toks_1d: Tensor,
+    *,
+    bl: int,
+    k: int,
+    K: int,
+    MASK: int,
+    device: torch.device,
+    temperature: float,
+    temperature_sampling: Literal["multinomial", "gumbel"],
+    omnivoice_unmask_schedule_k: int,
+    time_shift_tau: float,
+    prior_override_v: Tensor | None,
+    block_marginal_prior: Tensor | None,
+    pmi_cfg_probs_c_bl: Tensor | None,
+    pmi_cfg_probs_u_bl: Tensor | None,
+    cfg_scale_for_pmi: float,
+    position_temperature: float,
+) -> dict:
+    """One denoising step for a single sample block — ``pmi_count_early``.
+
+    Per-step unmask count combines OmniVoice's count schedule (floor) with
+    the paper's PMI-quantile (ceiling):
+
+        n_sched  = omnivoice_unmask_schedule_k  (OmniVoice r_n at step k)
+        q_k      = max(0, 1 - time_shift_tau * (k+1) / K)
+        tau_k    = Quantile({pmi_i}_{i in M}, q_k)        (if time_shift_tau>0)
+        n_quant  = |{i in M : pmi_i >= tau_k}|             (else 0)
+        n_unmask = min(max(n_sched, n_quant), m_cur)
+
+    The top ``n_unmask`` masked positions by PMI get unmasked. Properties:
+      * ``time_shift_tau == 0`` reduces to plain pmi_count (fixed-K, no
+        early decoding); per-step count is exactly ``n_sched``.
+      * ``time_shift_tau > 0`` may unmask more than ``n_sched`` when PMI
+        is concentrated, allowing the block to early-finish (n_mask==0)
+        before step ``K-1`` — the *only* source of step savings above the
+        baseline.
+
+    PMI scoring:
+      * Numerator: probabilities from the conditional softmax of the
+        outer block logits (``probs_bl``).
+      * Denominator: ``prior_override_v`` (precomputed unconditional
+        block prior) when provided, otherwise the running per-block
+        marginal frozen at step 0 (``block_marginal_prior``).
+      * Under CFG (``cfg_scale > 0``) the caller passes
+        ``pmi_cfg_probs_{c,u}_bl`` and the CFG-combined PMI
+        ``(1+w) * pmi_c - w * pmi_u`` is used (this is the only PMI/CFG
+        combination supported — no other ``cfg_prior_mode`` is exposed).
+
+    Gumbel position scoring (OmniVoice §3.4 T=5): when
+    ``position_temperature > 0`` we add Gumbel noise scaled by ``T`` to
+    masked positions' PMI before ranking, turning deterministic top-k
+    into stochastic top-k without replacement.
+    """
+    probs_raw = F.softmax(blk_logits_2d, dim=-1)
+    if temperature > 0:
+        sampled = _sample_at_temperature(
+            blk_logits_2d, temperature, sampling=temperature_sampling,
+        )
+    else:
+        sampled = blk_logits_2d.argmax(dim=-1)
+
+    ar = torch.arange(bl, device=device)
+    is_mask = blk_toks_1d == MASK
+
+    # ---- PMI scoring ----
+    probs_bl = probs_raw
+    if k == 0 and block_marginal_prior is None:
+        block_marginal_prior = probs_bl.mean(dim=0)
+
+    if pmi_cfg_probs_c_bl is not None and pmi_cfg_probs_u_bl is not None:
+        mp_c = (
+            prior_override_v
+            if prior_override_v is not None
+            else pmi_cfg_probs_c_bl.mean(dim=0)
+        )
+        mp_u = (
+            prior_override_v
+            if prior_override_v is not None
+            else pmi_cfg_probs_u_bl.mean(dim=0)
+        )
+        pmi_c = compute_pmi(pmi_cfg_probs_c_bl, sampled, marginal_prior=mp_c)
+        pmi_u = compute_pmi(pmi_cfg_probs_u_bl, sampled, marginal_prior=mp_u)
+        pmi = pmi_cfg_combine(pmi_c, pmi_u, cfg_scale_for_pmi)
+    else:
+        _mp = (
+            prior_override_v
+            if prior_override_v is not None
+            else block_marginal_prior
+        )
+        pmi = compute_pmi(probs_bl, sampled, marginal_prior=_mp)
+
+    if position_temperature > 0.0:
+        noise = _gumbel_sample(torch.zeros_like(pmi), position_temperature)
+        pmi = torch.where(is_mask, pmi + noise, pmi)
+
+    # ---- Outlier rule (pmi_count_early): max(count-schedule, quantile) ----
+    if k == K - 1:
+        do_unmask = is_mask
+    else:
+        masked_idx = ar[is_mask]
+        m_cur = int(masked_idx.numel())
+        do_unmask = torch.zeros_like(is_mask, dtype=torch.bool)
+        if m_cur > 0:
+            pmi_m = pmi[masked_idx]
+            n_sched = max(0, min(int(omnivoice_unmask_schedule_k), m_cur))
+            if time_shift_tau > 0.0:
+                q_shifted = max(0.0, 1.0 - time_shift_tau * (k + 1) / K)
+                tau_k = torch.quantile(pmi_m.float(), q_shifted).to(dtype=pmi_m.dtype)
+                n_quant = int((pmi_m >= tau_k).sum().item())
+            else:
+                n_quant = 0
+            n_unmask = min(max(n_sched, n_quant), m_cur)
+            if n_unmask >= m_cur:
+                do_unmask[masked_idx] = True
+            elif n_unmask > 0:
+                _, _topi = pmi_m.topk(n_unmask)
+                do_unmask[masked_idx[_topi]] = True
+
+    xt_step = torch.where(do_unmask, sampled, blk_toks_1d)
+    n_mask = int((xt_step == MASK).sum().item())
+    return dict(
+        xt_step=xt_step,
+        n_mask=n_mask,
+        block_marginal_prior=block_marginal_prior,
+    )
+
+
+# --------------------------------------------------------------------------- #
+# ChatterboxFlashT3
+# --------------------------------------------------------------------------- #
+
+
+class ChatterboxFlashT3(_ChatterboxT3):
+    """Chatterbox-TTS T3 decoder + [MASK] token + block-diffusion generate().
+
+    The Llama backbone, conditioning encoder, text / speech embeddings and
+    heads come straight from the parent ``chatterbox.models.t3.t3.T3``.  We
+    only add the extra mask-token row in :attr:`speech_emb` and override
+    :meth:`generate` with the block-diffusion inference loop.
+    """
+
+    def __init__(
+        self,
+        hp: T3Config | None = None,
+        *,
+        drf_block_size: int = 16,
+    ) -> None:
+        super().__init__(hp=hp)
+        if self.is_gpt:
+            raise NotImplementedError(
+                "ChatterboxFlashT3 currently supports the Llama backbone only "
+                "(hp.llama_config_name='Llama_520M').",
+            )
+        self._mask_token_id: int = self.hp.speech_tokens_dict_size
+        self.drf_block_size = int(drf_block_size)
+
+        # Extend speech_emb by one row for the [MASK] token (id = V).
+        old_speech_emb = self.speech_emb
+        self.speech_emb = nn.Embedding(self.hp.speech_tokens_dict_size + 1, self.dim)
+        with torch.no_grad():
+            self.speech_emb.weight[: self.hp.speech_tokens_dict_size].copy_(
+                old_speech_emb.weight,
+            )
+            nn.init.normal_(
+                self.speech_emb.weight[self.hp.speech_tokens_dict_size :], std=0.02,
+            )
+
+        # Backward-compat alias so :class:`FlashInferEngine` can read
+        # ``model.tfmr.config`` (chatterbox T3 only exposes ``.cfg``).
+        self.config = self.cfg
+        # Inference engine cache (re-used between calls when shape compatible).
+        self._cached_engine: InferenceEngine | None = None
+
+    # ------------------------------------------------------------------ #
+    #  Properties / helpers
+    # ------------------------------------------------------------------ #
+
+    @property
+    def mask_token_id(self) -> int:
+        return self._mask_token_id
+
+    @property
+    def is_llama(self) -> bool:
+        return not self.is_gpt
+
+    def set_block_size(self, block_size: int) -> None:
+        self.drf_block_size = int(block_size)
+
+    def prime_uncond_block_prior(
+        self,
+        dtype: torch.dtype,
+        device: torch.device | str,
+        block_size: int | None = None,
+    ) -> None:
+        """Eagerly compute & cache the PMI unconditional block prior.
+
+        Called at load time so the first ``generate()`` does not pay for the
+        one-shot prior forward — keeping it out of the timed / RTF loop. Uses
+        ``self.drf_block_size`` unless ``block_size`` is given; a later
+        ``generate()`` with a different block size recomputes automatically
+        (the cache is keyed on ``(block_size, dtype, device)``).
+        """
+        device = torch.device(device) if isinstance(device, str) else device
+        bs = int(block_size) if block_size is not None else int(self.drf_block_size)
+        with torch.no_grad():
+            self._compute_uncond_block_prior(
+                bs, int(self.hp.start_speech_token), int(self.mask_token_id),
+                dtype, device,
+            )
+
+    def _embed_speech_tokens(self, tokens: Tensor) -> Tensor:
+        """Embed speech tokens (including the extra ``[MASK]`` id)."""
+        emb = self.speech_emb(tokens)
+        if self.hp.input_pos_emb == "learned":
+            emb = emb + self.speech_pos_emb(tokens)
+        return emb
+
+    # ------------------------------------------------------------------ #
+    #  Precomputed unconditional block prior (PMI denominator)
+    # ------------------------------------------------------------------ #
+
+    def _compute_uncond_block_prior(
+        self,
+        block_size: int,
+        SOS: int,
+        MASK: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> Tensor:
+        """PMI prior from a forward on ``[SOS, MASK × block_size]`` with **no
+        conditioning at all** — no speaker, no text, no prompt speech, not
+        even zero-filled embeddings of those.  Captures the model's intrinsic
+        marginal over speech tokens at the very start of an utterance.
+
+        The result depends only on ``(block_size, SOS, MASK, dtype, device)``
+        and the (frozen, inference-time) model weights, so we **cache** it
+        on ``self._uncond_block_prior_cache`` keyed by that tuple.  The actual
+        forward runs at most once per ``(block_size, dtype, device)`` triple
+        across the entire lifetime of the model object — subsequent
+        ``generate()`` calls reuse the cached ``(V,)`` tensor.  Invalidated
+        implicitly when the model is recreated (e.g. on checkpoint reload).
+        """
+        cache: dict = getattr(self, "_uncond_block_prior_cache", None)
+        if cache is None:
+            cache = {}
+            self._uncond_block_prior_cache = cache
+        key = (int(block_size), int(SOS), int(MASK), str(dtype), str(device))
+        if key in cache:
+            return cache[key]
+
+        n_tok = 1 + int(block_size)
+        tok_seq = torch.cat(
+            [
+                torch.full((1, 1), int(SOS), device=device, dtype=torch.long),
+                torch.full(
+                    (1, int(block_size)), int(MASK),
+                    device=device, dtype=torch.long,
+                ),
+            ],
+            dim=1,
+        )
+        seq_emb = self._embed_speech_tokens(tok_seq).to(dtype)
+        position_ids = torch.arange(n_tok, device=device).unsqueeze(0)
+
+        # The FlashInfer engine swaps the LlamaAttention forward in-place;
+        # we explicitly route this one-shot prior pass through HuggingFace's
+        # SDPA implementation so it never touches the FlashInfer paged cache.
+        cfg = self.tfmr.config
+        old_impl = getattr(cfg, "_attn_implementation", "sdpa")
+        cfg._attn_implementation = "sdpa"
+        try:
+            out = self.tfmr(
+                input_ids=None,
+                inputs_embeds=seq_emb,
+                position_ids=position_ids,
+                use_cache=False,
+                return_dict=True,
+            )
+        finally:
+            cfg._attn_implementation = old_impl
+
+        # Standard token-shift: logits for mask position i come from hidden
+        # at position i (SOS → mask 0, mask k-1 → mask k).
+        shift_hidden = out.last_hidden_state[:, : int(block_size), :]
+        logits = self.speech_head(shift_hidden)
+        probs = F.softmax(logits.float(), dim=-1)
+        prior = probs.mean(dim=1).squeeze(0).detach()
+        cache[key] = prior
+        return prior
+
+    # ------------------------------------------------------------------ #
+    #  State-dict surgery (lets the model accept either a pure chatterbox
+    #  checkpoint of size V or a Chatterbox-Flash checkpoint of size V+1)
+    # ------------------------------------------------------------------ #
+
+    def _load_from_state_dict(
+        self, state_dict, prefix, *args, **kwargs,
+    ):  # noqa: D401
+        key = prefix + "speech_emb.weight"
+        if key in state_dict:
+            ckpt_w = state_dict[key]
+            model_v = self.speech_emb.weight.shape[0]
+            if ckpt_w.shape[0] == model_v - 1:
+                logger.info(
+                    "Expanding speech_emb from %d to %d (adding [MASK] token).",
+                    ckpt_w.shape[0], model_v,
+                )
+                expanded = self.speech_emb.weight.detach().clone()
+                expanded[: ckpt_w.shape[0]] = ckpt_w
+                state_dict[key] = expanded
+        return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
+    # ------------------------------------------------------------------ #
+    #  Prefix forward (run once per generation, cached in FlashInfer)
+    # ------------------------------------------------------------------ #
+
+    def _build_prefix_emb(
+        self,
+        cond_emb: Tensor,
+        text_tokens: Tensor,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> Tensor:
+        """Build ``[cond | text | <start_speech>]`` embedding."""
+        text_emb = self.text_emb(text_tokens)
+        if self.hp.input_pos_emb == "learned":
+            text_emb = text_emb + self.text_pos_emb(text_tokens)
+        sos = torch.full(
+            (1, 1), self.hp.start_speech_token, device=device, dtype=torch.long,
+        )
+        sos_emb = self._embed_speech_tokens(sos)
+        return torch.cat([cond_emb, text_emb, sos_emb], dim=1).to(dtype)
+
+    # ------------------------------------------------------------------ #
+    #  Block-diffusion generate
+    # ------------------------------------------------------------------ #
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        *,
+        t3_cond: "T3Cond | list[T3Cond] | tuple[T3Cond, ...]",
+        text_tokens: "torch.Tensor | list[torch.Tensor]",
+        text_token_lens: "torch.Tensor | None" = None,
+        total_speech_len: "int | list[int]" = 0,
+        num_steps: int = 10,
+        temperature: float = 0.6,
+        temperature_sampling: Literal["multinomial", "gumbel"] = "multinomial",
+        time_shift_tau: float = 0.1,
+        omnivoice_schedule_t_shift: float = 0.5,
+        cfg_scale: float = 1.0,
+        position_temperature: float = 5.0,
+        pmi_uncond_prior_precompute: bool = True,
+        batch_size: int = 1,
+        page_size: int = 16,
+        flashinfer_reserve_max_seq: int | None = None,
+        use_cuda_graph: bool = True,
+        backend: Literal["auto", "flashinfer", "torch"] = "auto",
+    ) -> Tensor:
+        """Block-wise denoising generation matching the paper's best config.
+
+        Inputs accept either a single sample (``text_tokens`` is a tensor and
+        ``t3_cond`` is a :class:`T3Cond`) or a heterogeneous batch
+        (``text_tokens`` is a list, ``t3_cond`` is a list of :class:`T3Cond`).
+        Both modes are dispatched through the same FlashInfer block-decoding
+        loop. Heterogeneous batches pad to ``max(total_speech_len)`` and stop
+        each row once it emits ``hp.stop_speech_token``.
+
+        Decoding-relevant parameters (paper defaults):
+
+        ``num_steps``
+            Maximum denoising iterations per block (``K``). Default 10.
+        ``time_shift_tau``
+            Early-decoding aggressiveness in
+            ``q_k = max(0, 1 - τ·(k+1)/K)``. ``0.0`` disables early decoding
+            (every block uses exactly the OmniVoice count schedule). Default
+            ``0.1`` (paper).
+        ``omnivoice_schedule_t_shift``
+            OmniVoice r_n schedule parameter that controls the per-step
+            unmask count floor. Default ``0.5``.
+        ``cfg_scale``
+            CFG strength. When ``> 0`` the forward batch is doubled with a
+            null branch (``zero_text_batch`` + ``zero_all`` null prefix +
+            text/cond zeroed + speech-x_t duplicated). Token sampling uses
+            the guided distribution; PMI scoring is combined on the PMI
+            scale via ``(1+w)·pmi_c − w·pmi_u`` (``pmi_cfg``) — this is the
+            only supported PMI/CFG combination, no other mode is exposed.
+        ``position_temperature``
+            Gumbel temperature added to PMI for stochastic top-k position
+            ranking. ``0`` = deterministic argmax-top-k; default ``5.0``
+            (OmniVoice §3.4 T=5).
+        ``pmi_uncond_prior_precompute``
+            When ``True`` (default), the PMI denominator is the
+            unconditional block prior computed once via
+            :meth:`_compute_uncond_block_prior` and cached on the model.
+            When ``False``, falls back to the running per-block marginal
+            from step 0.
+
+        Returns
+        -------
+        torch.LongTensor
+            Speech token ids of shape ``(B, T)`` with
+            ``T <= max(total_speech_len)``.
+        """
+        device = self.device
+        BS = self.drf_block_size
+        K = num_steps
+        MASK = self.mask_token_id
+        dtype = next(self.parameters()).dtype
+        stop_tok = self.hp.stop_speech_token
+
+        # ---- cross-text vs same-text batching ----
+        _cross_text = isinstance(text_tokens, list)
+        if _cross_text:
+            B = len(text_tokens)
+            _speech_lens: list[int] = (
+                list(total_speech_len)
+                if isinstance(total_speech_len, list)
+                else [total_speech_len] * B
+            )
+            N = max(_speech_lens) if _speech_lens else 0
+            _per_sample_num_blocks = [math.ceil(sl / BS) for sl in _speech_lens]
+            num_blocks = max(_per_sample_num_blocks) if _per_sample_num_blocks else 0
+        else:
+            B = max(1, batch_size)
+            _speech_lens = []
+            N = total_speech_len if isinstance(total_speech_len, int) else total_speech_len[0]
+            num_blocks = math.ceil(N / BS) if N > 0 else 0
+            _per_sample_num_blocks = []
+
+        if N <= 0:
+            return torch.empty((B, 0), device=device, dtype=torch.long)
+
+        B_usr = B
+        _ztb = cfg_scale > 0
+        B_fwd = (2 * B_usr) if _ztb else B_usr
+
+        # ---- prefix(es) ----
+        t3_is_seq = isinstance(t3_cond, (list, tuple))
+        if t3_is_seq and not _cross_text:
+            raise ValueError(
+                "t3_cond as a list/tuple is only supported when text_tokens is a list.",
+            )
+
+        cond_embs: list[Tensor] | None = None
+        cond_emb_single: Tensor | None = None
+        if _cross_text:
+            if t3_is_seq:
+                if len(t3_cond) != len(text_tokens):
+                    raise ValueError(
+                        f"len(t3_cond) ({len(t3_cond)}) must match "
+                        f"len(text_tokens) ({len(text_tokens)}).",
+                    )
+                cond_embs = [self.prepare_conditioning(c) for c in t3_cond]
+            else:
+                cond_emb_single = self.prepare_conditioning(t3_cond)
+        else:
+            cond_emb_single = self.prepare_conditioning(t3_cond)
+
+        prefix_emb_cond: Tensor | list[Tensor]
+        prefix_emb_null: Tensor | list[Tensor] | None = None
+        _prefix_lens: list[int]
+
+        if _cross_text:
+            cond_list: list[Tensor] = []
+            null_list: list[Tensor] = []
+            lens: list[int] = []
+            for i, tt in enumerate(text_tokens):
+                ce = cond_embs[i] if cond_embs is not None else cond_emb_single
+                p = self._build_prefix_emb(ce, tt, device, dtype)
+                cond_list.append(p)
+                lens.append(p.size(1))
+                if _ztb:
+                    ce_null = _cond_emb_zero_all(ce)
+                    text_emb = self.text_emb(tt)
+                    if self.hp.input_pos_emb == "learned":
+                        text_emb = text_emb + self.text_pos_emb(tt)
+                    text_emb_zero = _zero_text_content_keep_pad(
+                        text_emb, tt, None,
+                    )
+                    sos = torch.full(
+                        (1, 1), self.hp.start_speech_token,
+                        device=device, dtype=torch.long,
+                    )
+                    sos_emb = self._embed_speech_tokens(sos)
+                    p_null = torch.cat(
+                        [ce_null, text_emb_zero, sos_emb], dim=1,
+                    ).to(dtype)
+                    null_list.append(p_null)
+            prefix_emb_cond = cond_list
+            lp = max(lens)
+            _prefix_lens = list(lens)
+            if _ztb:
+                prefix_emb_null = null_list
+                _prefix_lens = _prefix_lens + [p.size(1) for p in null_list]
+        else:
+            assert cond_emb_single is not None
+            prefix_emb = self._build_prefix_emb(
+                cond_emb_single, text_tokens, device, dtype,
+            )
+            lp = prefix_emb.size(1)
+            if B_usr > 1:
+                prefix_emb = prefix_emb.expand(B_usr, -1, -1).contiguous()
+            prefix_emb_cond = prefix_emb
+            _prefix_lens = [lp] * B_fwd
+            if _ztb:
+                cond_for_null = _cond_emb_zero_all(cond_emb_single)
+                text_emb = self.text_emb(text_tokens)
+                if self.hp.input_pos_emb == "learned":
+                    text_emb = text_emb + self.text_pos_emb(text_tokens)
+                text_emb_zero = _zero_text_content_keep_pad(
+                    text_emb, text_tokens, text_token_lens,
+                )
+                sos = torch.full(
+                    (1, 1), self.hp.start_speech_token,
+                    device=device, dtype=torch.long,
+                )
+                sos_emb = self._embed_speech_tokens(sos)
+                p_null = torch.cat(
+                    [cond_for_null, text_emb_zero, sos_emb], dim=1,
+                ).to(dtype)
+                if B_usr > 1:
+                    p_null = p_null.expand(B_usr, -1, -1).contiguous()
+                prefix_emb_null = p_null
+
+        # ---- speech buffer (mask everywhere; we fill it block by block) ----
+        xt = torch.full((B_fwd, N), MASK, device=device, dtype=torch.long)
+
+        # ---- precomputed unconditional block prior (PMI denominator) ----
+        # Single forward on [SOS, MASK × BS] with no conditioning, lazy-cached
+        # on the model. Reused as PMI denominator for every step of every
+        # block in this generate() call (and across subsequent calls).
+        SOS = int(self.hp.start_speech_token)
+        prior_override_v: Tensor | None = None
+        if pmi_uncond_prior_precompute:
+            prior_override_v = self._compute_uncond_block_prior(
+                BS, SOS, MASK, dtype, device,
+            )
+
+        # ---- inference engine (FlashInfer if available, else torch SDPA) ----
+        engine_max_seq = max(lp + N + 64, flashinfer_reserve_max_seq or 0)
+        cached = self._cached_engine
+        if (
+            cached is not None
+            and cached.can_reuse(
+                engine_max_seq, dtype, batch_size=B_fwd, page_size=page_size,
+            )
+        ):
+            engine = cached
+            engine.reset()
+        else:
+            if (
+                use_cuda_graph
+                and cached is not None
+                and getattr(cached, "_max_seq_len", 0) < engine_max_seq
+                and getattr(cached, "_batch_size", -1) == B_fwd
+                and getattr(cached, "_dtype", None) == dtype
+            ):
+                # Grow allocation in 512-token buckets so we don't re-instantiate
+                # the engine every time the sequence length nudges up by one.
+                engine_max_seq = max(
+                    ((engine_max_seq + 511) // 512) * 512,
+                    int(getattr(cached, "_max_seq_len", 0) * 1.5),
+                )
+            engine = build_engine(
+                self, engine_max_seq, dtype,
+                backend=backend,
+                batch_size=B_fwd, page_size=page_size,
+            )
+
+        if _ztb:
+            assert prefix_emb_null is not None
+            if isinstance(prefix_emb_cond, list):
+                assert isinstance(prefix_emb_null, list)
+                pfx_for_engine = [*prefix_emb_cond, *prefix_emb_null]
+            else:
+                assert isinstance(prefix_emb_null, torch.Tensor)
+                pfx_for_engine = [
+                    *[prefix_emb_cond[bi : bi + 1] for bi in range(B_usr)],
+                    *[prefix_emb_null[bi : bi + 1] for bi in range(B_usr)],
+                ]
+            shift_ctx = engine.prefix_forward(pfx_for_engine)
+        else:
+            shift_ctx = engine.prefix_forward(prefix_emb_cond)
+
+        if use_cuda_graph:
+            # ``capture_cuda_graph`` is a no-op on the torch SDPA fallback.
+            engine.capture_cuda_graph(BS, speech_head=self.speech_head)
+        self._cached_engine = engine
+
+        # ---- block-by-block decoding ----
+        # ``_row_eos[bi]`` — row has emitted EOS, skip remaining blocks.
+        # ``_block_done[bi]`` — inner k-loop finished for this row in the
+        # current block (reset at every block boundary; not a sequence-level
+        # flag).
+        _row_eos = (
+            torch.zeros(B_usr, dtype=torch.bool, device=device)
+            if B_usr > 1
+            else None
+        )
+        _block_done = (
+            torch.zeros(B_usr, dtype=torch.bool, device=device)
+            if B_usr > 1
+            else None
+        )
+        block_marginal_priors: list[Tensor | None] = [None] * B_usr
+
+        for b_idx in range(num_blocks):
+            bs_ = b_idx * BS
+            be_ = min(bs_ + BS, N)
+            bl = be_ - bs_
+
+            # OmniVoice r_n count schedule for this block's K steps. Used as
+            # the per-step unmask **floor** by the pmi_count_early outlier
+            # rule. Computed once per block since ``bl`` only changes at the
+            # tail.
+            _omnivoice_sched = _omnivoice_unmask_schedule(
+                n_total_mask=bl,
+                num_steps=K,
+                t_shift=float(omnivoice_schedule_t_shift),
+            )
+
+            if _block_done is not None:
+                _block_done.zero_()
+                if _cross_text:
+                    for bi in range(B_usr):
+                        if (
+                            b_idx >= _per_sample_num_blocks[bi]
+                            or (_row_eos is not None and bool(_row_eos[bi]))
+                        ):
+                            _block_done[bi] = True
+                elif _row_eos is not None:
+                    for bi in range(B_usr):
+                        if bool(_row_eos[bi]):
+                            _block_done[bi] = True
+
+            for k in range(K):
+                # Per-row cache-start positions (cross-text & null-prefix paths
+                # may differ between rows; same-text reuses a scalar).
+                if _cross_text:
+                    cache_start: int | list[int] = [
+                        _prefix_lens[bi] + bs_ for bi in range(B_fwd)
+                    ]
+                else:
+                    cache_start = lp + bs_
+
+                sp_blk = self._embed_speech_tokens(xt[:, bs_:be_]).to(dtype)
+
+                use_graph = engine.has_cuda_graph and bl == BS
+                if use_graph:
+                    if k == 0:
+                        engine.set_shift_ctx(shift_ctx)
+                    blk_hidden, graph_logits = engine.block_forward_graph(
+                        sp_blk, cache_start,
+                    )
+                    if graph_logits is not None:
+                        blk_logits = graph_logits
+                    else:
+                        shift_hidden = torch.cat(
+                            [shift_ctx, blk_hidden[:, : bl - 1]], dim=1,
+                        )
+                        blk_logits = self.speech_head(shift_hidden)
+                else:
+                    blk_hidden = engine.block_forward(sp_blk, cache_start)
+                    shift_hidden = torch.cat(
+                        [shift_ctx, blk_hidden[:, : bl - 1]], dim=1,
+                    )
+                    blk_logits = self.speech_head(shift_hidden)
+
+                pmi_cfg_probs_c: Tensor | None = None
+                pmi_cfg_probs_u: Tensor | None = None
+                if _ztb:
+                    logits_cond = blk_logits[:B_usr]
+                    logits_uncond = blk_logits[B_usr : 2 * B_usr]
+                    # Always pmi_cfg: PMI scores are combined on the PMI
+                    # scale via (1+w)·pmi_c − w·pmi_u. No other CFG/PMI
+                    # combination is exposed.
+                    pmi_cfg_probs_c = F.softmax(logits_cond, dim=-1)
+                    pmi_cfg_probs_u = F.softmax(logits_uncond, dim=-1)
+                    blk_logits = apply_zero_text_cfg_from_logits(
+                        logits_cond, logits_uncond, cfg_scale,
+                    )
+
+                step_break = False
+                for bi in range(B_usr):
+                    if _block_done is not None and _block_done[bi]:
+                        continue
+                    pmi_c_bi = None if pmi_cfg_probs_c is None else pmi_cfg_probs_c[bi]
+                    pmi_u_bi = None if pmi_cfg_probs_u is None else pmi_cfg_probs_u[bi]
+                    r = _pmi_count_early_step_unmask(
+                        blk_logits[bi],
+                        xt[bi, bs_:be_],
+                        bl=bl,
+                        k=k, K=K,
+                        MASK=MASK, device=device,
+                        temperature=temperature,
+                        temperature_sampling=temperature_sampling,
+                        omnivoice_unmask_schedule_k=_omnivoice_sched[k],
+                        time_shift_tau=time_shift_tau,
+                        prior_override_v=prior_override_v,
+                        block_marginal_prior=block_marginal_priors[bi],
+                        pmi_cfg_probs_c_bl=pmi_c_bi,
+                        pmi_cfg_probs_u_bl=pmi_u_bi,
+                        cfg_scale_for_pmi=cfg_scale,
+                        position_temperature=position_temperature,
+                    )
+                    xt[bi, bs_:be_] = r["xt_step"]
+                    block_marginal_priors[bi] = r["block_marginal_prior"]
+                    if _ztb:
+                        xt[B_usr + bi, bs_:be_] = r["xt_step"]
+                    row_hit_eos = bool(
+                        (xt[bi, bs_:be_] == stop_tok).any().item(),
+                    )
+                    if r["n_mask"] == 0 or row_hit_eos:
+                        if _block_done is not None:
+                            _block_done[bi] = True
+                        if row_hit_eos and _row_eos is not None:
+                            _row_eos[bi] = True
+                        if B_usr == 1 and row_hit_eos:
+                            step_break = True
+                        elif B_usr == 1 and r["n_mask"] == 0:
+                            step_break = True
+                if step_break:
+                    break
+                if _block_done is not None and bool(_block_done.all().item()):
+                    break
+
+            # Mid-block finalize: write the committed (possibly partial) tokens
+            # of this block to the KV cache so the next block sees clean context.
+            if b_idx < num_blocks - 1:
+                fin_emb = self._embed_speech_tokens(xt[:, bs_:be_]).to(dtype)
+                if _cross_text:
+                    fin_cs: int | list[int] = [
+                        _prefix_lens[bi] + bs_ for bi in range(B_fwd)
+                    ]
+                else:
+                    fin_cs = lp + bs_
+                fin_hidden = engine.block_forward(fin_emb, fin_cs)
+                engine.advance_cache(bl)
+                shift_ctx = fin_hidden[:, bl - 1 : bl, :].clone()
+
+            # Early termination if all (same-text) rows have finished.
+            if B_usr == 1:
+                eos_pos = (xt[0, bs_:be_] == stop_tok).nonzero(as_tuple=True)[0]
+                if len(eos_pos) > 0:
+                    return xt[:B_usr, : bs_ + eos_pos[0].item()]
+            elif _row_eos is not None and bool(_row_eos.all().item()):
+                break
+
+        return xt[:B_usr]

runtime/chatterbox-flash/chatterbox_flash/nvfp4.py

@@ -0,0 +1,709 @@
+#!/usr/bin/env python3
+"""Canonical, source-free NVFP4 loader for Chatterbox-Flash.
+
+This module is the runtime counterpart to
+``export_chatterbox_flash_w4a16_hf.py``.  It constructs T3 on the meta device,
+loads only the BF16 non-projection tensors, prepares canonical NVFP4 tensors
+one projection at a time, and attaches an engine-local projection provider.
+No BF16 transformer projection is materialized and global ``F.linear`` is
+never patched.
+
+The canonical artifact remains independent of FlashInfer's prepared layout:
+``prepare_bf16_fp4_weights`` runs at load time, while ``nvfp4_quantize`` is
+intentionally neither required nor called.  The provider is attached to the
+model rather than an individual engine, so cache-size engine rebuilds reuse
+the same prepared weights.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import load_file
+
+
+ARTIFACT_FILENAME = "t3_nvfp4.safetensors"
+MANIFEST_FILENAME = "nvfp4_config.json"
+BASE_REPO_ID = "ResembleAI/chatterbox-flash"
+BASE_CHECKPOINT_REVISION = "4385507288b8197e6dab8b4e6b1603328d549d9d"
+BASE_ASSET_FILENAMES = (
+    "ve.safetensors",
+    "s3gen.safetensors",
+    "tokenizer.json",
+)
+SUPPORTED_RELEASE_BLOCK_SIZES = (16, 24)
+PROJECTION_KINDS = ("gate_up", "down", "qkv", "o")
+PROJECTION_COMPONENTS = ("weight_fp4", "scale_128x4_e4m3", "alpha")
+EXPECTED_LAYERS = 30
+EXPECTED_PREPARED_BYTES = 283_116_000
+EXPECTED_NONPROJECTION_TENSORS = 82
+EXPECTED_NONPROJECTION_BYTES = 58_179_584
+PROJECTION_DIMENSIONS = {
+    "gate_up": (8_192, 1_024),
+    "down": (1_024, 4_096),
+    "qkv": (3_072, 1_024),
+    "o": (1_024, 1_024),
+}
+
+
+def _tensor_bytes(tensor: torch.Tensor) -> int:
+    return tensor.numel() * tensor.element_size()
+
+
+def _canonical_cuda_device(device: torch.device | str) -> torch.device:
+    """Resolve bare ``cuda`` to the indexed device used by model tensors."""
+    resolved = torch.device(device)
+    if resolved.type == "cuda" and resolved.index is None:
+        resolved = torch.device("cuda", torch.cuda.current_device())
+    return resolved
+
+
+def _sha256_file(path: Path, chunk_bytes: int = 8 * 1024 * 1024) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        while chunk := handle.read(chunk_bytes):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _is_projection_source_key(key: str) -> bool:
+    parts = key.split(".")
+    if (
+        len(parts) != 6
+        or parts[0] != "tfmr"
+        or parts[1] != "layers"
+        or not parts[2].isdigit()
+        or parts[-1] != "weight"
+    ):
+        return False
+    module, projection = parts[3], parts[4]
+    return (
+        module == "self_attn"
+        and projection in {"q_proj", "k_proj", "v_proj", "o_proj"}
+    ) or (
+        module == "mlp"
+        and projection in {"gate_proj", "up_proj", "down_proj"}
+    )
+
+
+def _canonical_shapes(kind: str) -> dict[str, tuple[int, ...]]:
+    output_features, input_features = PROJECTION_DIMENSIONS[kind]
+    return {
+        "weight_fp4": (output_features, input_features // 2),
+        "scale_128x4_e4m3": (output_features, input_features // 16),
+        "alpha": (1,),
+    }
+
+
+def _resolve_package(
+    package: str | Path,
+    manifest_path: str | Path | None,
+) -> tuple[Path, Path, dict[str, Any]]:
+    package_path = Path(package).expanduser().resolve()
+    if package_path.is_dir():
+        artifact_path = package_path / ARTIFACT_FILENAME
+        default_manifest = package_path / MANIFEST_FILENAME
+    else:
+        artifact_path = package_path
+        default_manifest = package_path.with_name(MANIFEST_FILENAME)
+    config_path = (
+        Path(manifest_path).expanduser().resolve()
+        if manifest_path is not None
+        else default_manifest
+    )
+    if not artifact_path.is_file():
+        raise FileNotFoundError(f"Missing canonical NVFP4 artifact: {artifact_path}")
+    if not config_path.is_file():
+        raise FileNotFoundError(f"Missing canonical NVFP4 manifest: {config_path}")
+    manifest = json.loads(config_path.read_text(encoding="utf-8"))
+    if manifest.get("format") != "chatterbox-flash-canonical-nvfp4":
+        raise RuntimeError(f"Unsupported NVFP4 format in {config_path}")
+    if manifest.get("format_version") != 1:
+        raise RuntimeError(
+            f"Unsupported NVFP4 format version: {manifest.get('format_version')!r}"
+        )
+    declared_filename = manifest.get("artifact", {}).get("filename")
+    if declared_filename != artifact_path.name:
+        raise RuntimeError(
+            f"Manifest declares artifact {declared_filename!r}, "
+            f"but loader received {artifact_path.name!r}"
+        )
+    return artifact_path, config_path, manifest
+
+
+def _verify_base_assets(
+    paths: dict[str, Path],
+    manifest: dict[str, Any],
+) -> None:
+    """Verify every unchanged base asset before it reaches a model loader."""
+    expected_files = manifest.get("base_model", {}).get("files", {})
+    for filename in BASE_ASSET_FILENAMES:
+        path = paths.get(filename)
+        record = expected_files.get(filename)
+        if path is None or not path.is_file():
+            raise FileNotFoundError(f"Missing base-model asset: {filename}")
+        if not isinstance(record, dict):
+            raise RuntimeError(f"Manifest has no fingerprint for {filename}")
+        expected_bytes = record.get("bytes")
+        actual_bytes = path.stat().st_size
+        if actual_bytes != expected_bytes:
+            raise RuntimeError(
+                f"Base asset size mismatch for {filename}: expected "
+                f"{expected_bytes}, found {actual_bytes}"
+            )
+        expected_sha256 = record.get("sha256")
+        actual_sha256 = _sha256_file(path)
+        if not expected_sha256 or actual_sha256 != expected_sha256:
+            raise RuntimeError(
+                f"Base asset SHA-256 mismatch for {filename}: expected "
+                f"{expected_sha256}, found {actual_sha256}"
+            )
+
+
+@dataclass(frozen=True)
+class PreparedProjection:
+    kind: str
+    layer: int
+    input_features: int
+    output_features: int
+    prepared: tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+    packed_bytes: int
+
+
+class CanonicalNVFP4ProjectionProvider:
+    """FlashInfer projection provider prepared directly from canonical FP4."""
+
+    source_free = True
+
+    def __init__(
+        self,
+        artifact_path: str | Path,
+        *,
+        device: torch.device | str,
+        dtype: torch.dtype = torch.bfloat16,
+        autotune_cache: str | Path | None = None,
+    ) -> None:
+        import flashinfer
+
+        self.device = _canonical_cuda_device(device)
+        self.dtype = dtype
+        self.backend = "cute-dsl"
+        self.artifact_path = Path(artifact_path).expanduser().resolve()
+        self.autotune_cache = (
+            Path(autotune_cache).expanduser().resolve()
+            if autotune_cache is not None
+            else None
+        )
+        if self.device.type != "cuda":
+            raise RuntimeError("Canonical NVFP4 projections require a CUDA device")
+        if dtype is not torch.bfloat16:
+            raise RuntimeError(
+                f"Canonical W4A16 requires BF16 activations, found {dtype}"
+            )
+        if not torch.cuda.is_available():
+            raise RuntimeError("CUDA is unavailable")
+        capability = torch.cuda.get_device_capability(self.device)
+        if capability not in {(12, 0), (12, 1)}:
+            raise RuntimeError(
+                "The canonical CuTe-DSL NVFP4 profile is validated only on "
+                f"SM120/SM121, found compute capability {capability}"
+            )
+        required = (
+            "autotune",
+            "prepare_bf16_fp4_weights",
+            "mm_bf16_fp4",
+        )
+        missing = [name for name in required if not hasattr(flashinfer, name)]
+        if missing:
+            raise RuntimeError(
+                f"FlashInfer {getattr(flashinfer, '__version__', 'unknown')} "
+                f"lacks required runtime APIs: {missing}"
+            )
+
+        self._flashinfer = flashinfer
+        self._prepared: dict[str, list[PreparedProjection]] = {
+            kind: [] for kind in PROJECTION_KINDS
+        }
+        self._tuned_shapes: set[tuple[int, int]] = set()
+        self._load_canonical_projections()
+
+    def _load_canonical_projections(self) -> None:
+        expected_dtypes = {
+            "weight_fp4": torch.uint8,
+            "scale_128x4_e4m3": torch.uint8,
+            "alpha": torch.float32,
+        }
+        with safe_open(
+            str(self.artifact_path), framework="pt", device="cpu",
+        ) as handle:
+            available = set(handle.keys())
+            for layer in range(EXPECTED_LAYERS):
+                for kind in PROJECTION_KINDS:
+                    prefix = f"nvfp4.layers.{layer:02d}.{kind}"
+                    keys = {
+                        component: f"{prefix}.{component}"
+                        for component in PROJECTION_COMPONENTS
+                    }
+                    missing = [key for key in keys.values() if key not in available]
+                    if missing:
+                        raise RuntimeError(
+                            f"Canonical artifact is missing projection tensors: {missing}"
+                        )
+                    expected_shapes = _canonical_shapes(kind)
+                    canonical: dict[str, torch.Tensor] = {}
+                    for component in PROJECTION_COMPONENTS:
+                        tensor = handle.get_tensor(keys[component])
+                        if tuple(tensor.shape) != expected_shapes[component]:
+                            raise RuntimeError(
+                                f"{keys[component]} shape {tuple(tensor.shape)} != "
+                                f"{expected_shapes[component]}"
+                            )
+                        if tensor.dtype is not expected_dtypes[component]:
+                            raise RuntimeError(
+                                f"{keys[component]} dtype {tensor.dtype} != "
+                                f"{expected_dtypes[component]}"
+                            )
+                        canonical[component] = tensor.to(
+                            device=self.device, non_blocking=False,
+                        )
+
+                    # This is preparation, not quantization: canonical FP4
+                    # codes/scales are consumed directly, one projection at a
+                    # time.  Synchronizing before advancing bounds transient
+                    # device ownership to this projection plus prepared output.
+                    prepared_raw = self._flashinfer.prepare_bf16_fp4_weights(
+                        canonical["weight_fp4"],
+                        canonical["scale_128x4_e4m3"],
+                        canonical["alpha"],
+                        backend=self.backend,
+                    )
+                    if len(prepared_raw) != 3 or any(
+                        tensor is None for tensor in prepared_raw
+                    ):
+                        raise RuntimeError(
+                            f"prepare_bf16_fp4_weights returned an invalid tuple "
+                            f"for {kind}[{layer}]"
+                        )
+                    prepared = tuple(prepared_raw)
+                    torch.cuda.synchronize(self.device)
+                    output_features, input_features = PROJECTION_DIMENSIONS[kind]
+                    record = PreparedProjection(
+                        kind=kind,
+                        layer=layer,
+                        input_features=input_features,
+                        output_features=output_features,
+                        prepared=prepared,
+                        packed_bytes=sum(_tensor_bytes(tensor) for tensor in prepared),
+                    )
+                    self._prepared[kind].append(record)
+                    del canonical, prepared_raw, prepared
+
+        packed_bytes = sum(
+            record.packed_bytes
+            for records in self._prepared.values()
+            for record in records
+        )
+        if packed_bytes != EXPECTED_PREPARED_BYTES:
+            raise RuntimeError(
+                f"Prepared projection payload {packed_bytes} bytes != "
+                f"{EXPECTED_PREPARED_BYTES}"
+            )
+
+    def validate_engine(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        num_layers: int,
+    ) -> None:
+        if _canonical_cuda_device(device) != self.device:
+            raise RuntimeError(
+                f"Projection provider is on {self.device}, engine requested {device}"
+            )
+        if dtype is not self.dtype:
+            raise RuntimeError(
+                f"Projection provider uses {self.dtype}, engine requested {dtype}"
+            )
+        if num_layers != EXPECTED_LAYERS or any(
+            len(records) != num_layers for records in self._prepared.values()
+        ):
+            raise RuntimeError(
+                f"Projection provider does not cover all {num_layers} layers"
+            )
+
+    def prepare_for_engine(self, *, batch_size: int, block_size: int) -> None:
+        shape = (int(batch_size), int(block_size))
+        if shape in self._tuned_shapes:
+            return
+        if self.autotune_cache is not None:
+            self.autotune_cache.parent.mkdir(parents=True, exist_ok=True)
+        m = shape[0] * shape[1]
+        representatives = [self._prepared[kind][0] for kind in PROJECTION_KINDS]
+        with self._flashinfer.autotune(
+            True,
+            cache=(str(self.autotune_cache) if self.autotune_cache else None),
+            tuning_buckets=(m,),
+            round_up=False,
+        ):
+            for record in representatives:
+                activation = torch.zeros(
+                    (m, record.input_features),
+                    device=self.device,
+                    dtype=self.dtype,
+                )
+                self._run(record, activation)
+        torch.cuda.synchronize(self.device)
+        self._tuned_shapes.add(shape)
+
+    def _run(
+        self,
+        record: PreparedProjection,
+        activation: torch.Tensor,
+    ) -> torch.Tensor:
+        original_shape = activation.shape
+        flattened = activation.reshape(-1, record.input_features)
+        output = self._flashinfer.mm_bf16_fp4(
+            flattened,
+            *record.prepared,
+            backend=self.backend,
+            out_dtype=self.dtype,
+        )
+        return output.view(*original_shape[:-1], record.output_features)
+
+    def project(
+        self,
+        kind: str,
+        layer_idx: int,
+        activation: torch.Tensor,
+    ) -> torch.Tensor:
+        try:
+            record = self._prepared[kind][layer_idx]
+        except (KeyError, IndexError) as exc:
+            raise RuntimeError(
+                f"Unknown canonical projection {kind}[{layer_idx}]"
+            ) from exc
+        if activation.device != self.device or activation.dtype is not self.dtype:
+            raise RuntimeError(
+                f"{kind}[{layer_idx}] expected {self.device}/{self.dtype}, "
+                f"found {activation.device}/{activation.dtype}"
+            )
+        if activation.shape[-1] != record.input_features:
+            raise RuntimeError(
+                f"{kind}[{layer_idx}] expected K={record.input_features}, "
+                f"found {activation.shape[-1]}"
+            )
+        return self._run(record, activation)
+
+    def report(self) -> dict[str, Any]:
+        return {
+            "format": "chatterbox-flash-canonical-nvfp4",
+            "source_free": True,
+            "artifact": str(self.artifact_path),
+            "backend": self.backend,
+            "device": str(self.device),
+            "dtype": str(self.dtype),
+            "layers": EXPECTED_LAYERS,
+            "packed_linear_count": sum(map(len, self._prepared.values())),
+            "prepared_bytes": sum(
+                record.packed_bytes
+                for records in self._prepared.values()
+                for record in records
+            ),
+            "tuned_engine_shapes": sorted(self._tuned_shapes),
+            "calls_nvfp4_quantize": False,
+        }
+
+
+def _load_source_free_model_state(
+    artifact_path: Path,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+    drf_block_size: int,
+):
+    from chatterbox_flash.model import ChatterboxFlashT3
+
+    with torch.device("meta"):
+        model = ChatterboxFlashT3(drf_block_size=drf_block_size)
+
+    with safe_open(
+        str(artifact_path), framework="pt", device="cpu",
+    ) as handle:
+        keys = set(handle.keys())
+        forbidden = sorted(key for key in keys if _is_projection_source_key(key))
+        if forbidden:
+            raise RuntimeError(
+                "Canonical artifact unexpectedly contains BF16 projection sources: "
+                f"{forbidden[:3]}"
+            )
+        nonprojection_keys = sorted(
+            key for key in keys if not key.startswith("nvfp4.")
+        )
+        if len(nonprojection_keys) != EXPECTED_NONPROJECTION_TENSORS:
+            raise RuntimeError(
+                f"Expected {EXPECTED_NONPROJECTION_TENSORS} non-projection tensors, "
+                f"found {len(nonprojection_keys)}"
+            )
+        state = {
+            key: handle.get_tensor(key).to(
+                device=device,
+                dtype=dtype,
+                non_blocking=False,
+            )
+            for key in nonprojection_keys
+        }
+        state_bytes = sum(_tensor_bytes(tensor) for tensor in state.values())
+        if state_bytes != EXPECTED_NONPROJECTION_BYTES:
+            raise RuntimeError(
+                f"Non-projection payload {state_bytes} bytes != "
+                f"{EXPECTED_NONPROJECTION_BYTES} bytes"
+            )
+        incompatible = model.load_state_dict(state, strict=False, assign=True)
+        if incompatible.unexpected_keys:
+            raise RuntimeError(
+                f"Unexpected non-projection keys: {incompatible.unexpected_keys}"
+            )
+        unexpected_missing = [
+            key for key in incompatible.missing_keys
+            if not _is_projection_source_key(key)
+        ]
+        if unexpected_missing:
+            raise RuntimeError(
+                f"Artifact is missing non-projection model keys: {unexpected_missing}"
+            )
+        del state
+
+        rotary = model.tfmr.rotary_emb
+        rotary.inv_freq = handle.get_tensor(
+            "nvfp4.runtime.rotary.inv_freq"
+        ).to(device=device, dtype=dtype)
+        rotary.original_inv_freq = handle.get_tensor(
+            "nvfp4.runtime.rotary.original_inv_freq"
+        ).to(device=device, dtype=dtype)
+
+        prior_cache: dict[tuple[Any, ...], torch.Tensor] = {}
+        sos = int(model.hp.start_speech_token)
+        mask = int(model.mask_token_id)
+        prior_prefix = "nvfp4.priors.block_"
+        for key in sorted(key for key in keys if key.startswith(prior_prefix)):
+            block_size = int(key.removeprefix(prior_prefix))
+            prior = handle.get_tensor(key)
+            if tuple(prior.shape) != (8_194,) or prior.dtype is not torch.float32:
+                raise RuntimeError(
+                    f"Invalid cached prior {key}: {tuple(prior.shape)}/{prior.dtype}"
+                )
+            cache_key = (
+                block_size,
+                sos,
+                mask,
+                str(dtype),
+                str(device),
+            )
+            prior_cache[cache_key] = prior.to(device=device, non_blocking=False)
+
+    if drf_block_size not in {
+        int(key[0]) for key in prior_cache
+    }:
+        raise RuntimeError(
+            f"Artifact has no cached prior for D={drf_block_size}; source-free "
+            "models cannot compute a missing prior through the BF16 transformer"
+        )
+
+    meta_projection_parameters = [
+        name for name, parameter in model.named_parameters()
+        if parameter.is_meta and _is_projection_source_key(name)
+    ]
+    unexpected_meta = [
+        name for name, parameter in model.named_parameters()
+        if parameter.is_meta and not _is_projection_source_key(name)
+    ]
+    unexpected_meta.extend(
+        name for name, buffer in model.named_buffers() if buffer.is_meta
+    )
+    if len(meta_projection_parameters) != EXPECTED_LAYERS * 7:
+        raise RuntimeError(
+            f"Expected {EXPECTED_LAYERS * 7} meta projection sentinels, "
+            f"found {len(meta_projection_parameters)}"
+        )
+    if unexpected_meta:
+        raise RuntimeError(f"Unexpected meta model tensors: {unexpected_meta}")
+
+    model._uncond_block_prior_cache = prior_cache
+    model._nvfp4_source_free = True
+    model.eval()
+    return model
+
+
+def load_prepacked_t3(
+    package: str | Path,
+    *,
+    device: torch.device | str = "cuda",
+    dtype: torch.dtype = torch.bfloat16,
+    drf_block_size: int = 16,
+    autotune_cache: str | Path | None = None,
+    manifest_path: str | Path | None = None,
+    verify_sha256: bool = True,
+):
+    """Load a source-free T3 plus its reusable FlashInfer projection provider."""
+    if int(drf_block_size) not in SUPPORTED_RELEASE_BLOCK_SIZES:
+        raise RuntimeError(
+            f"Release profiles support D={SUPPORTED_RELEASE_BLOCK_SIZES}; "
+            f"D={drf_block_size} was not accepted for release use"
+        )
+    device = _canonical_cuda_device(device)
+    artifact_path, _config_path, manifest = _resolve_package(
+        package, manifest_path,
+    )
+    if verify_sha256:
+        expected_sha256 = manifest.get("artifact", {}).get("sha256")
+        actual_sha256 = _sha256_file(artifact_path)
+        if not expected_sha256 or actual_sha256 != expected_sha256:
+            raise RuntimeError(
+                f"Artifact SHA-256 mismatch: expected {expected_sha256}, "
+                f"found {actual_sha256}"
+            )
+    model = _load_source_free_model_state(
+        artifact_path,
+        device=device,
+        dtype=dtype,
+        drf_block_size=int(drf_block_size),
+    )
+    provider = CanonicalNVFP4ProjectionProvider(
+        artifact_path,
+        device=device,
+        dtype=dtype,
+        autotune_cache=autotune_cache,
+    )
+    model._flashinfer_projection_provider = provider
+    model._nvfp4_manifest = manifest
+    return model
+
+
+def _assemble_tts(
+    t3,
+    *,
+    ve_path: Path,
+    s3gen_path: Path,
+    tokenizer_path: Path,
+    device: torch.device,
+):
+    from chatterbox.models.s3gen import S3Gen
+    from chatterbox.models.tokenizers import EnTokenizer
+    from chatterbox.models.voice_encoder import VoiceEncoder
+    from chatterbox_flash.tts import ChatterboxFlashTTS
+
+    ve = VoiceEncoder()
+    ve.load_state_dict(load_file(str(ve_path)))
+    ve.to(device).eval()
+
+    s3gen = S3Gen(meanflow=True)
+    s3gen.load_state_dict(load_file(str(s3gen_path)), strict=False)
+    s3gen.to(device).eval()
+    tokenizer = EnTokenizer(str(tokenizer_path))
+    return ChatterboxFlashTTS(t3, s3gen, ve, tokenizer, device)
+
+
+def load_prepacked_tts_local(
+    package: str | Path,
+    base_model_dir: str | Path,
+    *,
+    device: torch.device | str = "cuda",
+    dtype: torch.dtype = torch.bfloat16,
+    drf_block_size: int = 16,
+    autotune_cache: str | Path | None = None,
+    verify_sha256: bool = True,
+    verify_base_sha256: bool = True,
+):
+    """Build the complete pipeline without loading ``t3_flash.safetensors``."""
+    device = _canonical_cuda_device(device)
+    base = Path(base_model_dir).expanduser().resolve()
+    t3 = load_prepacked_t3(
+        package,
+        device=device,
+        dtype=dtype,
+        drf_block_size=drf_block_size,
+        autotune_cache=autotune_cache,
+        verify_sha256=verify_sha256,
+    )
+    base_paths = {
+        filename: base / filename for filename in BASE_ASSET_FILENAMES
+    }
+    if verify_base_sha256:
+        _verify_base_assets(base_paths, t3._nvfp4_manifest)
+    return _assemble_tts(
+        t3,
+        ve_path=base_paths["ve.safetensors"],
+        s3gen_path=base_paths["s3gen.safetensors"],
+        tokenizer_path=base_paths["tokenizer.json"],
+        device=device,
+    )
+
+
+def load_prepacked_tts_from_hub(
+    repo_id: str,
+    *,
+    device: torch.device | str = "cuda",
+    dtype: torch.dtype = torch.bfloat16,
+    drf_block_size: int = 16,
+    autotune_cache: str | Path | None = None,
+    revision: str | None = None,
+    base_repo_id: str = BASE_REPO_ID,
+    base_revision: str | None = BASE_CHECKPOINT_REVISION,
+    verify_sha256: bool = True,
+    verify_base_sha256: bool = True,
+):
+    """Download a derivative artifact plus unmodified base pipeline assets."""
+    from huggingface_hub import hf_hub_download
+
+    artifact_path = Path(hf_hub_download(
+        repo_id=repo_id,
+        filename=ARTIFACT_FILENAME,
+        revision=revision,
+    ))
+    manifest_path = Path(hf_hub_download(
+        repo_id=repo_id,
+        filename=MANIFEST_FILENAME,
+        revision=revision,
+    ))
+    base_paths = {
+        filename: Path(hf_hub_download(
+            repo_id=base_repo_id,
+            filename=filename,
+            revision=base_revision,
+        ))
+        for filename in BASE_ASSET_FILENAMES
+    }
+    device = _canonical_cuda_device(device)
+    t3 = load_prepacked_t3(
+        artifact_path,
+        manifest_path=manifest_path,
+        device=device,
+        dtype=dtype,
+        drf_block_size=drf_block_size,
+        autotune_cache=autotune_cache,
+        verify_sha256=verify_sha256,
+    )
+    if verify_base_sha256:
+        _verify_base_assets(base_paths, t3._nvfp4_manifest)
+    return _assemble_tts(
+        t3,
+        ve_path=base_paths["ve.safetensors"],
+        s3gen_path=base_paths["s3gen.safetensors"],
+        tokenizer_path=base_paths["tokenizer.json"],
+        device=device,
+    )
+
+
+__all__ = [
+    "CanonicalNVFP4ProjectionProvider",
+    "PreparedProjection",
+    "load_prepacked_t3",
+    "load_prepacked_tts_local",
+    "load_prepacked_tts_from_hub",
+]

runtime/chatterbox-flash/chatterbox_flash/text_norm/__init__.py

@@ -0,0 +1,15 @@
+"""English text normalization for Chatterbox-Flash inference.
+
+This is a stripped-down port of Resemble's ``en_us_cleaner`` pipeline that
+produces identical cleaned text but drops the ``SliceMap`` alignment plumbing
+required for training. It expands numbers, currencies, units, time, phone
+numbers and a small set of abbreviations so the input text matches the form
+the model was trained on, which is essential for WER consistency on the
+LibriSpeech-PC and Seed-TTS evaluation benchmarks.
+
+The single entry point is :func:`en_us_cleaner`.
+"""
+
+from .cleaners import en_us_cleaner
+
+__all__ = ["en_us_cleaner"]

runtime/chatterbox-flash/chatterbox_flash/text_norm/cleaners.py

@@ -0,0 +1,129 @@
+"""``en_us_cleaner`` — exact-output port of Resemble's training-time text
+normalization pipeline (number / time / phone / abbreviation expansion) with
+the ``SliceMap`` alignment plumbing removed.
+
+This pipeline runs at inference time before the speech tokenizer; matching the
+training-time form is critical for WER on LibriSpeech-PC / Seed-TTS test-en.
+"""
+
+from __future__ import annotations
+
+import re
+import unicodedata
+
+from .number_norm import normalize_numbers
+from .phone_norm import convert_phone_numbers_to_words
+from .time_norm import convert_time_to_words
+
+
+_whitespace_re = re.compile(r"\s+")
+
+_abbreviations = [
+    (re.compile(r"\b%s\b\.?" % abbrev, re.IGNORECASE), expanded)
+    for abbrev, expanded in [
+        ('mrs', 'misess'),
+        ('mr', 'mister'),
+        ('dr', 'doctor'),
+        ('st', 'saint'),
+        ('jr', 'junior'),
+        ('maj', 'major'),
+        ('gen', 'general'),
+        ('drs', 'doctors'),
+        ('rev', 'reverend'),
+        ('lt', 'lieutenant'),
+        ('hon', 'honorable'),
+        ('sgt', 'sergeant'),
+        ('capt', 'captain'),
+        ('esq', 'esquire'),
+        ('ltd', 'limited'),
+        ('col', 'colonel'),
+        ('ft', 'feet'),
+        ('abbrev', 'abbreviation'),
+        ('ave', 'avenue'),
+        ('abstr', 'abstract'),
+        ('addr', 'address'),
+        ('jan', 'january'),
+        ('feb', 'february'),
+        ('mar', 'march'),
+        ('apr', 'april'),
+        ('jul', 'july'),
+        ('aug', 'august'),
+        ('sep', 'september'),
+        ('sept', 'september'),
+        ('oct', 'october'),
+        ('nov', 'november'),
+        ('dec', 'december'),
+        ('mon', 'monday'),
+        ('tue', 'tuesday'),
+        ('wed', 'wednesday'),
+        ('thur', 'thursday'),
+        ('fri', 'friday'),
+        ('sec', 'second'),
+        ('min', 'minute'),
+        ('mo', 'month'),
+        ('yr', 'year'),
+        ('cal', 'calorie'),
+        ('dept', 'department'),
+        ('gal', 'gallon'),
+        ('kg', 'kilogram'),
+        ('km', 'kilometer'),
+        ('mt', 'mount'),
+        ('oz', 'ounce'),
+        ('vol', 'volume'),
+        ('vs', 'versus'),
+        ('yd', 'yard'),
+        (r'e\.g', 'eg'),
+        (r'i\.e', 'ie'),
+        ('etc', 'etc'),
+        ('ai', 'a i'),
+        ('cms', 'c m s'),
+        ('cdn', 'c d n'),
+        ('ceo', 'c e o'),
+        ('tts', 't t s'),
+    ]
+]
+
+_custom_symbol_replacements = [
+    (re.compile('–'), '-'),
+    (re.compile('—'), '-'),
+    (re.compile(r'(?<=\d\s)>(?=\s\d)', re.IGNORECASE), 'is greater than'),
+    (re.compile(r'(?<=\d\s)<(?=\s\d)', re.IGNORECASE), 'is less than'),
+    (re.compile(r'(?<=\D)\s*-\s+(?=\D)'), ', '),
+    (re.compile(r'(?<=\D)\s+-\s*(?=\D)'), ', '),
+    (re.compile(r'(?<=[a-z])-(?=[a-z])', re.IGNORECASE), ' '),
+    (re.compile(r'(?<=[a-z])-(?=\d)', re.IGNORECASE), ' '),
+    (re.compile(r'(?<=\d)-(?=[a-z])', re.IGNORECASE), ' '),
+    (re.compile(r'(?<=\d)-(?=\d)', re.IGNORECASE), ' - '),
+]
+
+_abbreviations += _custom_symbol_replacements
+
+
+def _expand_abbreviations(text: str) -> str:
+    for regex, replacement in _abbreviations:
+        text = regex.sub(replacement, text)
+    return text
+
+
+def _standardize_characters(text: str) -> str:
+    return unicodedata.normalize("NFKC", text)
+
+
+def _collapse_whitespace(text: str) -> str:
+    return _whitespace_re.sub(" ", text)
+
+
+def en_us_cleaner(text: str) -> str:
+    """English text normalization pipeline used for both training and the
+    paper's evaluations.
+
+    Matches the cleaned-text output of Resemble's ``transforms.text.cleaners.en_us_cleaner``
+    for any normal English input.
+    """
+    text = _standardize_characters(text)
+    text = _collapse_whitespace(text)
+    text = _expand_abbreviations(text)
+    text = convert_time_to_words(text)
+    text = convert_phone_numbers_to_words(text)
+    text = normalize_numbers(text)
+    return text

runtime/chatterbox-flash/chatterbox_flash/text_norm/number_norm.py

@@ -0,0 +1,366 @@
+"""Number / currency / unit / year / ordinal expansion.
+
+Verbatim logic ported from Resemble's training-time ``transforms.text.number_norm``
+with the ``SliceMap`` alignment machinery stripped. The resulting cleaned text
+is byte-identical to the original pipeline for any normal English input.
+"""
+
+from __future__ import annotations
+
+import re
+
+
+_comma_number_re = re.compile(r'(\(?[A-Z]{2,3})?([\$|£|¥|€|#|\(]*[0-9][0-9\,\.]+[0-9])([^\s]+)?')
+_decimal_number_re = re.compile(r'(number\s)?([0-9]+\.[0-9]+)(\.|,|\?|!)?')
+_hash_number_re = re.compile(r'(#)([0-9]+(?:\.[0-9]+)?)(\.|,|\?|!)?')
+
+_pounds_re = re.compile(r'(\(?£)([0-9\.]*[0-9]+)(\.|,|\?|\!)?')
+_yen_re = re.compile(r'(\(?¥)([0-9]+)(\.|,|\?|\!)?')
+_euro_re = re.compile(r'(\(?€)([0-9\.]*[0-9]+)(\.|,|\?|\!)?')
+_dollars_re = re.compile(
+    r'(?P<curr>\(?\$)'
+    r'(?P<val>[0-9,]*\.?[0-9]+)'
+    r'(?P<punc>[\.|,|\?|\!|\)]+)?'
+)
+
+_curr_abbrev_re = re.compile(
+    r'(?P<curr>\(?[$£¥€])'
+    r'(?P<val>[0-9]*\.?[0-9]+)'
+    r'(?:(?P<unit1>[BKMT])|(?:\s(?P<unit2>[BMbmTtr]+illion)))'
+    r'(?P<punc>[.,?!)]+)?'
+)
+
+_ml_re = re.compile(r'([0-9\.]*[0-9]+)(ml)(\.|,|\?|!)?')
+_cl_re = re.compile(r'([0-9\.]*[0-9]+)(cl)(\.|,|\?|!)?')
+_g_re = re.compile(r'([0-9\.]*[0-9]+)(g)(\.|,|\?|!)?')
+_l_re = re.compile(r'([0-9\.]*[0-9]+)(l)(\.|,|\?|!)?')
+_m_re = re.compile(r'([0-9\.]*[0-9]+)(m)(\.|,|\?|!)?')
+_kg_re = re.compile(r'([0-9\.]*[0-9]+)(kg)(\.|,|\?|!)?')
+_mm_re = re.compile(r'([0-9\.]*[0-9]+)(mm)(\.|,|\?|!)?')
+_cm_re = re.compile(r'([0-9\.]*[0-9]+)(cm)(\.|,|\?|!)?')
+_km_re = re.compile(r'([0-9\.]*[0-9]+)(km)(\.|,|\?|!)?')
+_in_re = re.compile(r'([0-9\.]*[0-9]+)(in)(\.|,|\?|!)?')
+_ft_re = re.compile(r'([0-9\.]*[0-9]+)(ft)(\.|,|\?|!)?')
+_yd_re = re.compile(r'([0-9\.]*[0-9]+)(yd[s]?)(\.|,|\?|!)?')
+_s_re = re.compile(r'([0-9\.]*[0-9]+)(s[ecs]*)(\.|,|\?|!)?')
+_percent = re.compile(r'([0-9]{1,3}(?:,[0-9]{3})*(\.[0-9]+)?)(%)([.,?!])?')
+
+_ordinal_re = re.compile(r'([0-9]+)(st|nd|rd|th)')
+_number_re = re.compile(r'([0-9]+)(\.|,|\?|!)?')
+_year_re = re.compile(
+    r'(to|years|year|from|after|before|by|until|in|since|around|circa|during)'
+    r'(\s)?'
+    r'(?<!\$|£|¥|€|₩|₹|₽|฿|₫|₦|₪|₱|₴)'
+    r'(1[1-9]|20)'
+    r'([0-9]{2})'
+    r'($|\-|\s\-|\D)',
+    re.IGNORECASE,
+)
+
+_units = [
+    '', 'one', 'two', 'three', 'four', 'five', 'six', 'seven', 'eight', 'nine',
+    'ten', 'eleven', 'twelve', 'thirteen', 'fourteen', 'fifteen', 'sixteen',
+    'seventeen', 'eighteen', 'nineteen',
+]
+_tens = [
+    '', 'ten', 'twenty', 'thirty', 'forty', 'fifty', 'sixty', 'seventy',
+    'eighty', 'ninety',
+]
+_digit_groups = ['', 'thousand', 'million', 'billion', 'trillion', 'quadrillion']
+
+_ordinal_suffixes = [
+    ('one', 'first'),
+    ('two', 'second'),
+    ('three', 'third'),
+    ('five', 'fifth'),
+    ('eight', 'eighth'),
+    ('nine', 'ninth'),
+    ('twelve', 'twelfth'),
+    ('ty', 'tieth'),
+]
+
+_sub_ten_nums = {"00", "01", "02", "03", "04", "05", "06", "07", "08", "09"}
+
+_curr_dict = {"$": "dollars", "£": "pounds", "¥": "yen", "€": "euros"}
+_unit_dict = {"B": "billion", "K": "thousand", "M": "million", "T": "trillion"}
+
+
+def standard_number_to_words(n: int, digit_group: int) -> str:
+    parts = []
+    if n >= 1000:
+        parts.append(standard_number_to_words(n // 1000, digit_group + 1))
+        n = n % 1000
+    if n >= 100:
+        parts.append('%s hundred' % _units[n // 100])
+    if n % 100 >= len(_units):
+        parts.append(_tens[(n % 100) // 10])
+        parts.append(_units[(n % 100) % 10])
+    else:
+        parts.append(_units[n % 100])
+    if n > 0:
+        parts.append(_digit_groups[digit_group])
+    return ' '.join(x for x in parts if x)
+
+
+def _number_to_words(n: int) -> str:
+    if n >= 1000000000000000000:
+        return str(n)
+    if n == 0:
+        return 'zero'
+    if n % 100 == 0 and n % 1000 != 0 and n < 3000:
+        return standard_number_to_words(n // 100, 0) + ' hundred'
+    return standard_number_to_words(n, 0)
+
+
+def _remove_commas(match):
+    _, m1, m2 = match.groups(default="")
+    return ''.join([m1.replace(",", ""), m2])
+
+
+def _expand_year(match):
+    prep, _, mill_cent, dec_year, post = match.groups(default="")
+    if mill_cent == "20" and dec_year in _sub_ten_nums:
+        year_out = mill_cent + dec_year
+    elif dec_year in _sub_ten_nums:
+        year = dec_year[-1]
+        year_out = mill_cent + " " + "oh" + " " + year
+    else:
+        year_out = _number_to_words(int(mill_cent)) + " " + _number_to_words(int(dec_year))
+
+    if "-" in post:
+        year_out += " to"
+    else:
+        year_out += post
+    return prep + " " + year_out
+
+
+def _re_sub(pattern, repl, text: str) -> str:
+    """Backwards (end-to-beginning) regex substitution, identical match order
+    to the training-time ``aligned_re_sub`` so non-overlapping match decisions
+    line up exactly."""
+    matches = sorted(re.finditer(pattern, text), key=lambda m: m.start(), reverse=True)
+    for m in matches:
+        middle = repl(m) if callable(repl) else repl
+        text = text[:m.start()] + middle + text[m.end():]
+    return text
+
+
+# --- helpers operating on (word, _) mapping list, but with the SliceMap dropped ---
+
+
+def _to_word_list(text: str):
+    words = re.split(r"(\s+)", text)
+    return [[w] for w in words]
+
+
+def _from_word_list(mapping) -> str:
+    return ''.join(t[0] for t in mapping)
+
+
+def _convert_hash(mapping):
+    text = _from_word_list(mapping)
+    matches = re.findall(_hash_number_re, text)
+    for result in matches:
+        m = ''.join(result)
+        r_clean = ' '.join(["number", result[1]])
+        if len(result) == 3:
+            r_clean = r_clean + result[2]
+        for i, t in enumerate(mapping):
+            if m == t[0]:
+                mapping[i] = [r_clean]
+    return mapping
+
+
+def _expand_decimal_point(mapping):
+    text = _from_word_list(mapping)
+    matches = re.findall(_decimal_number_re, text)
+    for m in matches:
+        integral, fractional = m[1].split('.')
+        expanded_frac = ""
+        frac_it = iter(list(fractional))
+        for digit in frac_it:
+            if digit == "0":
+                expanded_frac += "zero "
+            else:
+                expanded_frac += digit + " "
+                for rest in frac_it:
+                    expanded_frac += rest + " "
+        out = m[0] + integral + ' point ' + expanded_frac.strip()
+        if len(m) == 3:
+            out = out + m[2]
+        raw = ''.join(m)
+        for i, t in enumerate(mapping):
+            if t[0] == raw:
+                mapping[i] = [out]
+    return mapping
+
+
+def _expand_dollars(mapping):
+    text = _from_word_list(mapping)
+    matches = re.findall(_dollars_re, text)
+    for m in matches:
+        val = m[1]
+        punc = m[2] if m[2] else ''
+        parts = val.split('.')
+        whole_num = int(parts[0]) if parts[0] else 0
+        fraction = int(parts[1]) if len(parts) > 1 and parts[1] else 0
+        dollar_unit = 'dollar' if whole_num == 1 else 'dollars'
+        cent_unit = 'cent' if fraction == 1 else 'cents'
+        if whole_num and fraction:
+            out = '%s %s, %s %s' % (whole_num, dollar_unit, fraction, cent_unit)
+        elif whole_num:
+            out = '%s %s' % (whole_num, dollar_unit)
+        elif fraction:
+            out = '%s %s' % (fraction, cent_unit)
+        else:
+            out = 'zero dollars'
+        if punc:
+            out = out + punc
+        raw = ''.join(m)
+        for i, t in enumerate(mapping):
+            if t[0] == raw:
+                mapping[i] = [out]
+    return mapping
+
+
+def _expand_other_currency(mapping, regex, one, many):
+    text = _from_word_list(mapping)
+    matches = re.findall(regex, text)
+    for m in matches:
+        parts = m[1].split(".")
+        curr = one if int(parts[0]) == 1 else many
+        try:
+            out = parts[0] + " " + curr + " " + parts[1]
+        except IndexError:
+            out = parts[0] + " " + curr
+        raw = ''.join(m)
+        for i, t in enumerate(mapping):
+            if t[0] == raw:
+                mapping[i] = [out]
+    return mapping
+
+
+def _expand_abbreviated_currency_unit(mapping):
+    text = _from_word_list(mapping)
+    matches = re.findall(_curr_abbrev_re, text)
+    to_remove = []
+    for m in matches:
+        curr = m[0]
+        val = m[1]
+        unit = m[2] if m[2] else m[3]
+        punc = m[4] if m[4] else ''
+        curr = curr.strip("(")
+        val_parts = val.split(".")
+        if len(val_parts) > 1:
+            val_out = val_parts[0] + ' ' + 'point ' + ' '.join(val_parts[1])
+        else:
+            val_out = val_parts[0]
+        try:
+            out = ' '.join([val_out, _unit_dict[unit], _curr_dict[curr]])
+            out = out + punc
+            raw = ''.join(m)
+        except KeyError:
+            out = ' '.join([val_out, unit, _curr_dict[curr]])
+            out = out + punc
+            raw = ''.join(m[:-2]) + " " + unit + punc
+        for i, t in enumerate(mapping):
+            if t[0] == raw:
+                mapping[i] = [out]
+            try:
+                join_text = ''.join([mapping[i][0], mapping[i + 1][0], mapping[i + 2][0]])
+                if join_text == raw:
+                    mapping[i] = [out]
+                    to_remove.append(i + 1)
+                    to_remove.append(i + 2)
+            except IndexError:
+                continue
+    return [v for i, v in enumerate(mapping) if i not in to_remove]
+
+
+def _expand_other_unit(mapping, regex, one, many):
+    text = _from_word_list(mapping)
+    matches = re.findall(regex, text)
+    for m in matches:
+        parts = re.split(r"\.", m[0])
+        unit = one if parts[0] == "1" else many
+        if len(parts) > 1:
+            dec = ''.join([i + " " for i in parts[1]])
+            parts = parts[0] + " " + "point" + " " + dec
+            unit = many
+            out = parts + unit
+        else:
+            out = parts[0] + " " + unit
+        out = out + m[-1]
+        raw = ''.join(m)
+        for i, t in enumerate(mapping):
+            if t[0] == raw:
+                mapping[i] = [out]
+    return mapping
+
+
+def _expand_ordinal(mapping):
+    text = _from_word_list(mapping)
+    matches = re.findall(_ordinal_re, text)
+    for m in matches:
+        num = _number_to_words(int(m[0]))
+        out = num + 'th'
+        for suffix, replacement in _ordinal_suffixes:
+            if num.endswith(suffix):
+                out = num[:-len(suffix)] + replacement
+                break
+        raw = ''.join(m)
+        for i, t in enumerate(mapping):
+            if t[0] == raw:
+                mapping[i] = [out]
+    return mapping
+
+
+def _expand_number(mapping):
+    text = _from_word_list(mapping)
+    matches = re.findall(_number_re, text)
+    for m in matches:
+        out = _number_to_words(int(m[0])) + m[1]
+        raw = ''.join(m)
+        for i, t in enumerate(mapping):
+            text_re_nums = re.search(r"\d+", t[0])
+            if text_re_nums is None:
+                continue
+            if m[0] == text_re_nums.group():
+                rep = re.escape(raw)
+                j = re.sub(rep, out, t[0])
+                mapping[i] = [j]
+    return mapping
+
+
+def normalize_numbers(text: str) -> str:
+    text = _re_sub(_comma_number_re, _remove_commas, text)
+    # year norm is applied twice on purpose (matches training pipeline)
+    text = _re_sub(_year_re, _expand_year, text)
+    text = _re_sub(_year_re, _expand_year, text)
+
+    mapping = _to_word_list(text)
+    mapping = _expand_abbreviated_currency_unit(mapping)
+    mapping = _expand_other_currency(mapping, _pounds_re, "pound", "pounds")
+    mapping = _expand_other_currency(mapping, _yen_re, "yen", "yen")
+    mapping = _expand_other_currency(mapping, _euro_re, "euro", "euros")
+    mapping = _expand_other_unit(mapping, _ml_re, "milliliter", "milliliters")
+    mapping = _expand_other_unit(mapping, _cl_re, "centiliter", "centiliters")
+    mapping = _expand_other_unit(mapping, _g_re, "gram", "grams")
+    mapping = _expand_other_unit(mapping, _kg_re, "kilogram", "kilograms")
+    mapping = _expand_other_unit(mapping, _mm_re, "millimeter", "millimeters")
+    mapping = _expand_other_unit(mapping, _cm_re, "centimeter", "centimeters")
+    mapping = _expand_other_unit(mapping, _km_re, "kilometer", "kilometers")
+    mapping = _expand_other_unit(mapping, _in_re, "inch", "inches")
+    mapping = _expand_other_unit(mapping, _ft_re, "foot", "feet")
+    mapping = _expand_other_unit(mapping, _l_re, "liter", "liters")
+    mapping = _expand_other_unit(mapping, _m_re, "meter", "meters")
+    mapping = _expand_other_unit(mapping, _yd_re, "yard", "yards")
+    mapping = _expand_other_unit(mapping, _s_re, "second", "seconds")
+    mapping = _expand_other_unit(mapping, _percent, "%", "percent")
+    mapping = _expand_dollars(mapping)
+    mapping = _convert_hash(mapping)
+    mapping = _expand_decimal_point(mapping)
+    mapping = _expand_ordinal(mapping)
+    mapping = _expand_number(mapping)
+
+    return _from_word_list(mapping)

runtime/chatterbox-flash/chatterbox_flash/text_norm/phone_norm.py

@@ -0,0 +1,107 @@
+"""Phone number expansion (e.g. "555-1234" -> "five five five, one two three four")
+ported from ``transforms.text.phone_num_norm`` without SliceMap."""
+
+from __future__ import annotations
+
+import re
+
+
+_NUM_WORDS = {
+    0: 'zero', 1: 'one', 2: 'two', 3: 'three', 4: 'four',
+    5: 'five', 6: 'six', 7: 'seven', 8: 'eight', 9: 'nine',
+}
+
+_NON_ALPHA = re.compile(r"(\s?[^a-zA-Z]+\s?)")
+_PHONE_NUM_CHARS = set(list("0123456789.+-()[] "))
+
+
+def is_phone_num(text: str) -> bool:
+    text = text.rstrip('.!?,;')
+    if text.count(".") == 1:
+        return False
+    if any(c.isalpha() for c in text):
+        return False
+    if any(c not in _PHONE_NUM_CHARS for c in text):
+        return False
+    nums = [c for c in text if c.isdigit()]
+    if not (5 < len(nums) < 14):
+        return False
+    if "+" not in text and "." not in text and list(text).count("-") == 1:
+        nums = [[c for c in subtext if c.isdigit()] for subtext in text.split("-")]
+        if len(nums) == 2 and all(
+            subnums[0] in ("1", "2") and len(subnums) == 4 for subnums in nums
+        ):
+            return False
+    return True
+
+
+def _group_num(text):
+    prev_was_digit = False
+    group = []
+    num_groups = []
+    for c in text:
+        if c.isdigit():
+            if len(group) == 0 or prev_was_digit:
+                group.append(c)
+            else:
+                num_groups.append(group)
+                group = [c]
+            prev_was_digit = True
+        else:
+            prev_was_digit = False
+    if group:
+        num_groups.append(group)
+    return num_groups
+
+
+def _natural_phone_num_grouping(nums, max_natural_group=4):
+    if len(nums) == 0:
+        return []
+    if len(nums) <= max_natural_group:
+        return [nums]
+    if len(nums) == 5:
+        return [nums[:2], nums[2:]]
+    if len(nums) == 8:
+        return [nums[:2], nums[2:5], nums[5:]]
+    return _natural_phone_num_grouping(nums[:-max_natural_group], 3) + [nums[-max_natural_group:]]
+
+
+def _convert_phone_num_to_words(text: str) -> str:
+    original_punctuation = ''
+    while text and text[-1] in '.!?,;':
+        original_punctuation = text[-1] + original_punctuation
+        text = text[:-1]
+
+    num_groups = _group_num(text)
+    num_word_groups = [
+        [" " + _NUM_WORDS[int(c)] for c in nums if c.isdigit()] for nums in num_groups
+    ]
+
+    output = []
+    for num_word_group in reversed(num_word_groups):
+        grouplen = len(num_word_group)
+        regroupped = [num_word_group]
+        if grouplen > 4:
+            regroupped = reversed(_natural_phone_num_grouping(num_word_group))
+        for new_group in regroupped:
+            output = ["".join(new_group)] + output
+
+    output = ",".join(output).strip()
+    parts = re.split(r"\s", text)
+    output = (" " if parts[0] == "" else "") + output + (" " if parts[-1] == "" else "")
+    output += original_punctuation
+    return output
+
+
+def _convert_phone_num_match_to_words(match: re.Match) -> str:
+    text = match.group()
+    if is_phone_num(text):
+        return _convert_phone_num_to_words(text)
+    return text
+
+
+def convert_phone_numbers_to_words(text: str) -> str:
+    matches = sorted(_NON_ALPHA.finditer(text), key=lambda m: m.start(), reverse=True)
+    for m in matches:
+        text = text[: m.start()] + _convert_phone_num_match_to_words(m) + text[m.end():]
+    return text

runtime/chatterbox-flash/chatterbox_flash/text_norm/time_norm.py

@@ -0,0 +1,49 @@
+"""Time-of-day expansion (12:30, 3.45 pm, ...) ported from
+``transforms.text.time_norm`` without the SliceMap return value."""
+
+from __future__ import annotations
+
+import re
+
+from .number_norm import standard_number_to_words
+
+
+_time_pattern = re.compile(
+    r'(?P<hour>\d{1,2})(?P<delim>[\.\:])?\s*(?P<min>\d{2})?\s*(?P<am_pm>[AaPp]\.?[Mm]\.?)?'
+)
+
+
+def _time_to_words(match: re.Match) -> str:
+    gd = match.groupdict()
+    hours = int(gd['hour'])
+    delim = gd.get("delim")
+    minutes = gd.get("min")
+    am_pm = gd.get("am_pm")
+
+    has_minutes = minutes is not None
+    has_am_pm = am_pm is not None
+    is_time = has_am_pm or (delim == ":" and has_minutes) or (delim == "." and has_am_pm)
+    if not is_time:
+        return match.group()
+
+    minutes_i = int(minutes) if has_minutes else 0
+    am_pm = (' ' + am_pm.replace('.', '').strip().upper()) if has_am_pm else ''
+    hour_word = standard_number_to_words(hours, 0)
+
+    if minutes_i < 10 and minutes_i != 0:
+        minute_word = standard_number_to_words(minutes_i, 0)
+        minute_word = f"oh {minute_word}"
+    else:
+        minute_word = standard_number_to_words(minutes_i, 0) if minutes_i > 0 else ""
+
+    if minutes_i > 0:
+        minute_word = f" {minute_word}"
+
+    return f"{hour_word}{minute_word}{am_pm}"
+
+
+def convert_time_to_words(text: str) -> str:
+    matches = sorted(_time_pattern.finditer(text), key=lambda m: m.start(), reverse=True)
+    for m in matches:
+        text = text[: m.start()] + _time_to_words(m) + text[m.end():]
+    return text

runtime/chatterbox-flash/chatterbox_flash/tts.py

@@ -0,0 +1,398 @@
+"""User-facing API for Chatterbox-Flash."""
+
+from __future__ import annotations
+
+import dataclasses
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable, Literal
+
+import librosa
+import numpy as np
+import torch
+import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+from chatterbox.models.s3gen import S3GEN_SR, S3Gen
+from chatterbox.models.s3tokenizer import S3_SR
+from chatterbox.models.t3.modules.cond_enc import T3Cond
+from chatterbox.models.tokenizers import EnTokenizer
+from chatterbox.models.voice_encoder import VoiceEncoder
+
+from .model import ChatterboxFlashT3
+from .text_norm import en_us_cleaner
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_REPO_ID = "ResembleAI/chatterbox-flash"
+
+# Per-file names expected inside a checkpoint dir / HF repo.
+T3_FILENAME = "t3_flash.safetensors"
+S3GEN_FILENAME = "s3gen.safetensors"
+VE_FILENAME = "ve.safetensors"
+TOKENIZER_FILENAME = "tokenizer.json"
+
+# Conditioning durations (mirrored from chatterbox-tts).
+ENC_COND_LEN = 6 * S3_SR
+DEC_COND_LEN = 10 * S3GEN_SR
+
+
+@dataclass
+class Conditionals:
+    """Prepared reference-speech conditioning for one speaker.
+
+    Bundles the T3 conditioning (speaker embedding + prompt speech tokens)
+    and the S3Gen reference dict produced by :meth:`S3Gen.embed_ref`. Use
+    :meth:`save` / :meth:`load` to cache these for repeated generations.
+    """
+
+    t3: T3Cond
+    gen: dict
+
+    def to(self, device) -> "Conditionals":
+        self.t3 = self.t3.to(device=device)
+        for k, v in self.gen.items():
+            if torch.is_tensor(v):
+                self.gen[k] = v.to(device=device)
+        return self
+
+    def save(self, fpath: Path) -> None:
+        torch.save({"t3": self.t3.__dict__, "gen": self.gen}, fpath)
+
+    @classmethod
+    def load(cls, fpath, map_location="cpu") -> "Conditionals":
+        if isinstance(map_location, str):
+            map_location = torch.device(map_location)
+        kwargs = torch.load(fpath, map_location=map_location, weights_only=True)
+        return cls(T3Cond(**kwargs["t3"]), kwargs["gen"])
+
+
+def _speech_len_for_text_tokens(n_text_tokens: int) -> int:
+    """Heuristic max-speech-token budget per utterance (matches paper config)."""
+    return max(n_text_tokens * 6, 300)
+
+
+class ChatterboxFlashTTS:
+    """Zero-shot TTS pipeline matching the configuration in the paper.
+
+    Stage 1 is the :class:`ChatterboxFlashT3` block-diffusion decoder.
+    Stage 2 is the unmodified Chatterbox-TTS S3Gen flow-matching vocoder.
+    Reference encoding uses the GE2E-trained voice encoder from Chatterbox.
+    """
+
+    def __init__(
+        self,
+        t3: ChatterboxFlashT3,
+        s3gen: S3Gen,
+        ve: VoiceEncoder,
+        tokenizer: EnTokenizer,
+        device: str | torch.device,
+        conds: Conditionals | None = None,
+    ) -> None:
+        self.sr = S3GEN_SR
+        self.t3 = t3
+        self.s3gen = s3gen
+        self.ve = ve
+        self.tokenizer = tokenizer
+        self.device = torch.device(device) if isinstance(device, str) else device
+        self.conds = conds
+
+    # -------- construction ------------------------------------------------
+
+    @classmethod
+    def from_local(
+        cls,
+        ckpt_dir: str | Path,
+        device: str | torch.device,
+        *,
+        dtype: torch.dtype = torch.bfloat16,
+        drf_block_size: int = 16,
+    ) -> "ChatterboxFlashTTS":
+        """Load a Chatterbox-Flash checkpoint from a local directory.
+
+        Expects ``t3_flash.safetensors`` (with the expanded MASK row),
+        ``s3gen.safetensors``, ``ve.safetensors`` and ``tokenizer.json``.
+        """
+        ckpt_dir = Path(ckpt_dir)
+        device = torch.device(device) if isinstance(device, str) else device
+
+        ve = VoiceEncoder()
+        ve.load_state_dict(load_file(str(ckpt_dir / VE_FILENAME)))
+        ve.to(device).eval()
+
+        t3 = ChatterboxFlashT3(drf_block_size=drf_block_size)
+        t3_state = load_file(str(ckpt_dir / T3_FILENAME))
+        if "model" in t3_state:  # legacy layout
+            t3_state = t3_state["model"][0]
+        t3.load_state_dict(t3_state)
+        t3.to(device=device, dtype=dtype).eval()
+        # Prime the PMI unconditional block prior now (at load) so the first
+        # generate() doesn't fold the one-shot prior forward into the timed /
+        # RTF loop. A generate() with a different block size recomputes it
+        # automatically (cache keyed on block_size/dtype/device).
+        t3.prime_uncond_block_prior(dtype, device)
+
+        # ``meanflow=True``: the released ``s3gen.safetensors`` is the
+        # meanflow-distilled vocoder, sampled with the meanflow ODE path at 2
+        # CFM steps (matches the internal ``S3GenStreamer`` defaults). On a
+        # standard-CFM checkpoint this would degrade quality — flip to False.
+        s3gen = S3Gen(meanflow=True)
+        s3gen.load_state_dict(
+            load_file(str(ckpt_dir / S3GEN_FILENAME)), strict=False,
+        )
+        s3gen.to(device).eval()
+
+        tokenizer = EnTokenizer(str(ckpt_dir / TOKENIZER_FILENAME))
+
+        return cls(t3, s3gen, ve, tokenizer, device)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        repo_id: str = DEFAULT_REPO_ID,
+        device: str | torch.device = "cuda",
+        *,
+        dtype: torch.dtype = torch.bfloat16,
+        drf_block_size: int = 16,
+    ) -> "ChatterboxFlashTTS":
+        """Fetch a Chatterbox-Flash release from the Hugging Face Hub."""
+        last = None
+        for fname in (
+            VE_FILENAME, T3_FILENAME, S3GEN_FILENAME, TOKENIZER_FILENAME,
+        ):
+            last = hf_hub_download(repo_id=repo_id, filename=fname)
+        return cls.from_local(
+            Path(last).parent, device,
+            dtype=dtype, drf_block_size=drf_block_size,
+        )
+
+    # -------- conditioning -----------------------------------------------
+
+    def prepare_conditionals(
+        self,
+        wav_fpath: str | Path,
+        exaggeration: float = 0.5,
+    ) -> Conditionals:
+        """Build the speaker / prompt-speech conditioning for a reference wav.
+
+        The reference is resampled and segmented identically to chatterbox-tts:
+        the first 6 s feed S3's speech tokenizer to produce prompt speech
+        tokens, the first 10 s feed S3Gen's :meth:`embed_ref`, and the GE2E
+        voice encoder produces the global speaker embedding.
+        """
+        s3gen_ref_wav, _sr = librosa.load(str(wav_fpath), sr=S3GEN_SR)
+        ref_16k_wav = librosa.resample(s3gen_ref_wav, orig_sr=S3GEN_SR, target_sr=S3_SR)
+
+        s3gen_ref_wav = s3gen_ref_wav[:DEC_COND_LEN]
+        s3gen_ref_dict = self.s3gen.embed_ref(
+            s3gen_ref_wav, S3GEN_SR, device=self.device,
+        )
+
+        t3_cond_prompt_tokens = None
+        if plen := self.t3.hp.speech_cond_prompt_len:
+            tokens, _ = self.s3gen.tokenizer.forward(
+                [ref_16k_wav[:ENC_COND_LEN]], max_len=plen,
+            )
+            t3_cond_prompt_tokens = torch.atleast_2d(tokens).to(self.device)
+
+        t3_dtype = next(self.t3.parameters()).dtype
+        ve_embed = torch.from_numpy(
+            self.ve.embeds_from_wavs([ref_16k_wav], sample_rate=S3_SR),
+        )
+        ve_embed = ve_embed.mean(axis=0, keepdim=True).to(device=self.device, dtype=t3_dtype)
+
+        t3_cond = T3Cond(
+            speaker_emb=ve_embed,
+            cond_prompt_speech_tokens=t3_cond_prompt_tokens,
+            emotion_adv=exaggeration * torch.ones(1, 1, 1, dtype=t3_dtype),
+        ).to(device=self.device)
+        self.conds = Conditionals(t3_cond, s3gen_ref_dict)
+        return self.conds
+
+    # -------- tokenization -----------------------------------------------
+
+    def _encode_text(
+        self, text: str, normalize_text: bool = True,
+    ) -> torch.Tensor:
+        if normalize_text:
+            text = en_us_cleaner(text)
+        text_tokens = self.tokenizer.text_to_tokens(text).to(self.device)
+        text_tokens = F.pad(text_tokens, (1, 0), value=self.t3.hp.start_text_token)
+        text_tokens = F.pad(text_tokens, (0, 1), value=self.t3.hp.stop_text_token)
+        return text_tokens
+
+    # -------- inference --------------------------------------------------
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        text: str,
+        *,
+        audio_prompt_path: str | Path | None = None,
+        exaggeration: float = 0.5,
+        normalize_text: bool = True,
+        num_steps: int = 10,
+        temperature: float = 0.6,
+        time_shift_tau: float = 0.1,
+        omnivoice_schedule_t_shift: float = 0.5,
+        cfg_scale: float = 1.0,
+        position_temperature: float = 5.0,
+        pmi_uncond_prior_precompute: bool = True,
+        use_cuda_graph: bool = True,
+        backend: Literal["auto", "flashinfer", "torch", "mlx"] = "auto",
+        max_speech_tokens: int | None = None,
+        n_cfm_timesteps: int = 2,
+    ) -> torch.Tensor:
+        """Synthesize a single utterance.
+
+        Parameters mirror the inference defaults from the paper (best
+        configuration). Returns the synthesized waveform as a 1xT float
+        tensor at :attr:`sr` Hz. See
+        :meth:`ChatterboxFlashT3.generate` for the full parameter spec.
+        """
+        if audio_prompt_path is not None:
+            self.prepare_conditionals(audio_prompt_path, exaggeration=exaggeration)
+        if self.conds is None:
+            raise ValueError(
+                "Conditioning not prepared — pass audio_prompt_path or call "
+                "prepare_conditionals() first.",
+            )
+
+        text_tokens = self._encode_text(text, normalize_text=normalize_text)
+        n_tokens = int(text_tokens.size(1))
+        total_speech_len = _speech_len_for_text_tokens(n_tokens)
+        if max_speech_tokens is not None:
+            total_speech_len = min(total_speech_len, max_speech_tokens)
+
+        speech_tokens = self.t3.generate(
+            t3_cond=self.conds.t3,
+            text_tokens=text_tokens,
+            text_token_lens=torch.tensor([n_tokens], device=self.device),
+            total_speech_len=total_speech_len,
+            num_steps=num_steps,
+            temperature=temperature,
+            time_shift_tau=time_shift_tau,
+            omnivoice_schedule_t_shift=omnivoice_schedule_t_shift,
+            cfg_scale=cfg_scale,
+            position_temperature=position_temperature,
+            pmi_uncond_prior_precompute=pmi_uncond_prior_precompute,
+            use_cuda_graph=use_cuda_graph,
+            backend=backend,
+            batch_size=1,
+        )
+        speech_tokens = self._trim_to_eos(speech_tokens[0])
+        wav, _ = self.s3gen.inference(
+            speech_tokens=speech_tokens.to(self.device),
+            ref_dict=self.conds.gen,
+            n_cfm_timesteps=n_cfm_timesteps,
+        )
+        return wav.detach().cpu().squeeze(0)
+
+    @torch.inference_mode()
+    def generate_batch(
+        self,
+        texts: Iterable[str],
+        *,
+        audio_prompt_paths: Iterable[str | Path] | None = None,
+        conds_list: Iterable[Conditionals] | None = None,
+        exaggeration: float = 0.5,
+        normalize_text: bool = True,
+        num_steps: int = 10,
+        temperature: float = 0.6,
+        time_shift_tau: float = 0.1,
+        omnivoice_schedule_t_shift: float = 0.5,
+        cfg_scale: float = 1.0,
+        position_temperature: float = 5.0,
+        pmi_uncond_prior_precompute: bool = True,
+        use_cuda_graph: bool = True,
+        backend: Literal["auto", "flashinfer", "torch", "mlx"] = "auto",
+        max_speech_tokens: int | None = None,
+        n_cfm_timesteps: int = 2,
+    ) -> list[torch.Tensor]:
+        """Synthesize a heterogeneous batch (one prompt + text per row).
+
+        Exactly one of ``audio_prompt_paths`` and ``conds_list`` must be
+        provided. Both ``texts`` and the conditioning iterable must have the
+        same length. Returns a list of 1D waveforms (one per row).
+        """
+        texts = list(texts)
+        if audio_prompt_paths is not None and conds_list is not None:
+            raise ValueError("Pass exactly one of audio_prompt_paths / conds_list.")
+        if audio_prompt_paths is None and conds_list is None:
+            raise ValueError("Pass either audio_prompt_paths or conds_list.")
+
+        if conds_list is None:
+            conds_list = [
+                dataclasses.replace(
+                    self.prepare_conditionals(p, exaggeration=exaggeration),
+                ) if False else self.prepare_conditionals(p, exaggeration=exaggeration)
+                for p in audio_prompt_paths
+            ]
+        else:
+            conds_list = list(conds_list)
+        if len(conds_list) != len(texts):
+            raise ValueError(
+                f"texts ({len(texts)}) and conditioning ({len(conds_list)}) "
+                "must have equal length.",
+            )
+
+        text_tokens_list: list[torch.Tensor] = []
+        speech_lens: list[int] = []
+        for t in texts:
+            tt = self._encode_text(t, normalize_text=normalize_text)
+            text_tokens_list.append(tt)
+            sl = _speech_len_for_text_tokens(int(tt.size(1)))
+            if max_speech_tokens is not None:
+                sl = min(sl, max_speech_tokens)
+            speech_lens.append(sl)
+
+        speech_tokens_batch = self.t3.generate(
+            t3_cond=[c.t3 for c in conds_list],
+            text_tokens=text_tokens_list,
+            text_token_lens=None,
+            total_speech_len=speech_lens,
+            num_steps=num_steps,
+            temperature=temperature,
+            time_shift_tau=time_shift_tau,
+            omnivoice_schedule_t_shift=omnivoice_schedule_t_shift,
+            cfg_scale=cfg_scale,
+            position_temperature=position_temperature,
+            pmi_uncond_prior_precompute=pmi_uncond_prior_precompute,
+            use_cuda_graph=use_cuda_graph,
+            backend=backend,
+        )
+
+        wavs: list[torch.Tensor] = []
+        for i, sl in enumerate(speech_lens):
+            row = speech_tokens_batch[i : i + 1, :sl]
+            row = self._trim_to_eos(row[0]).unsqueeze(0)
+            wav, _ = self.s3gen.inference(
+                speech_tokens=row.to(self.device),
+                ref_dict=conds_list[i].gen,
+                n_cfm_timesteps=n_cfm_timesteps,
+            )
+            wavs.append(wav.detach().cpu().squeeze(0))
+        return wavs
+
+    # -------- internal ---------------------------------------------------
+
+    def _trim_to_eos(self, speech_tokens: torch.Tensor) -> torch.Tensor:
+        """Truncate at the first stop-speech token and drop non-codec ids.
+
+        T3 generates ids in ``[0, hp.speech_tokens_dict_size)``, but S3Gen's
+        flow embedding is sized to the raw S3 codec vocab
+        (``hp.start_speech_token``). Anything ``>= start_speech_token`` —
+        the SOS / EOS specials, the ``[MASK]`` row, or never-unmasked slots —
+        must be removed before vocoding.
+        """
+        if speech_tokens.numel() == 0:
+            return speech_tokens.long()
+        stop = self.t3.hp.stop_speech_token
+        eos = (speech_tokens == stop).nonzero(as_tuple=True)[0]
+        if len(eos) > 0:
+            speech_tokens = speech_tokens[: eos[0].item()]
+        codec_vocab = int(self.t3.hp.start_speech_token)  # 6561 for chatterbox
+        speech_tokens = speech_tokens[speech_tokens < codec_vocab]
+        return speech_tokens.long()

runtime/chatterbox-flash/pyproject.toml

@@ -0,0 +1,137 @@
+[build-system]
+requires = ["setuptools>=61", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "chatterbox-flash-nvfp4"
+version = "0.1.0+nvfp4.1"
+description = "Unofficial source-free NVFP4 W4A16 runtime for Chatterbox-Flash T3 on NVIDIA Blackwell."
+readme = "README.md"
+license = { text = "MIT" }
+requires-python = ">=3.10"
+authors = [
+    { name = "Resemble AI" },
+    { name = "Community NVFP4 contributors" },
+]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Multimedia :: Sound/Audio :: Speech",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    # T3 / S3Gen / VoiceEncoder / EnTokenizer come from chatterbox-tts. Note that
+    # chatterbox-tts 0.1.7 hard-pins ``torch==2.6.0``, ``torchaudio==2.6.0``,
+    # ``transformers==5.2.0``, ``librosa==0.11.0`` and ``safetensors==0.5.3``.
+    # We deliberately leave those ranges loose here and let the chatterbox-tts
+    # solve drive the install — otherwise uv flips them back-and-forth between
+    # each invocation.
+    "chatterbox-tts>=0.1.7",
+    "torch>=2.6",
+    "torchaudio>=2.6",
+    "transformers>=4.46",
+    "safetensors>=0.4",
+    "huggingface-hub>=0.24",
+    "numpy>=1.24",
+    "tqdm>=4.66",
+    "soundfile>=0.12",
+    "librosa>=0.10",
+    "scipy>=1.11",
+    "regex",
+    "inflect>=7.0",
+    "Unidecode>=1.3",
+    # Used by the top-level ``synthesize.py`` CLI to parse YAML configs in
+    # ``configs/`` (best / best_no_early_dec / fast). The library itself
+    # does not import PyYAML.
+    "pyyaml>=6.0",
+]
+
+[project.optional-dependencies]
+# FlashInfer kernels (required for batched / CUDA-graph inference).
+# Install matching your CUDA build: pip install flashinfer-python
+flashinfer = [
+    # The canonical loader depends on the CuTe-DSL W4A16 preparation and GEMM
+    # APIs validated for this release.  Older 0.6.11 builds resolve from the
+    # upstream lock but do not reproduce the accepted runtime.
+    "flashinfer-python==0.6.14",
+    "nvidia-cutlass-dsl==4.5.2",
+]
+
+# Apple Silicon (Mac) native Metal backend. Pulls in mlx + mlx-lm; the
+# mlx-lm LLaMA implementation is used as the backbone with our custom
+# inputs_embeds bypass. CPU/GPU paths do not import these.
+mlx = [
+    "mlx>=0.13; sys_platform == 'darwin'",
+    "mlx-lm>=0.10; sys_platform == 'darwin'",
+]
+
+# Evaluation suite (SIM-o / WER / UTMOS) via the OmniVoice pip package.
+eval = [
+    "omnivoice>=0.1.5",
+    # omnivoice's WER eval imports the legacy ``jiwer.compute_measures``, which
+    # was removed in jiwer 4.0. Cap below 4 so the import keeps working.
+    "jiwer>=3.0,<4",
+    "zhon>=2.0",
+    "whisper-normalizer>=0.0.10",
+    # omnivoice's SIM-o (speaker similarity) loads the WavLM-large upstream via
+    # torch.hub, whose hubconf imports ``s3prl`` (s3prl.util.download +
+    # s3prl.upstream.wavlm.expert). omnivoice does not declare it, so pull it in
+    # here — otherwise Stage 2/3 SIM-o crashes with ModuleNotFoundError: s3prl.
+    "s3prl>=0.4",
+]
+
+[project.urls]
+Upstream = "https://github.com/resemble-ai/chatterbox-flash"
+BaseModel = "https://huggingface.co/ResembleAI/chatterbox-flash"
+Model = "https://huggingface.co/ajh-code/chatterbox-flash-t3-w4a16-nvfp4"
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["chatterbox_flash*"]
+
+[tool.setuptools.package-data]
+chatterbox_flash = ["py.typed"]
+
+# ----------------------------------------------------------------------------
+# uv-specific configuration
+# ----------------------------------------------------------------------------
+#
+# ``chatterbox-tts==0.1.7`` hard-pins ``torch==2.6.0`` / ``torchaudio==2.6.0`` /
+# ``transformers==5.2.0``, but most modern wheels installed alongside it
+# (``torchvision``, ``xformers``, ``flash-attn``, ``flashinfer-python``) ship
+# only torch-2.7 ABI binaries. Letting uv honour the chatterbox-tts pin causes
+# a cascade of ABI mismatches: ``torchvision::nms does not exist``,
+# ``flash_attn_2_cuda.so: undefined symbol``, etc.
+#
+# We therefore override the upper-stream pins to ``>=2.7`` for the install
+# itself. The chatterbox-tts code paths we use (T3 / S3Gen / VE / EnTokenizer)
+# work correctly against torch 2.7 — only the install-time metadata is stale.
+[tool.uv]
+override-dependencies = [
+    # Pin to the 2.7.x ABI that the rest of the wheel ecosystem (torchvision,
+    # xformers, flash-attn, flashinfer-python) currently ships matched binaries
+    # against. Bumping the upper bound requires re-checking those wheels.
+    "torch==2.7.1",
+    "torchaudio==2.7.1",
+    "torchvision==0.22.1",
+    "transformers==5.9.0",
+]
+
+# Use the exact CUDA 12.8 wheels validated on Blackwell.  Marking this index
+# explicit keeps unrelated packages on PyPI while routing only the PyTorch
+# family below to the official cu128 wheel repository.
+[[tool.uv.index]]
+name = "pytorch-cu128"
+url = "https://download.pytorch.org/whl/cu128"
+explicit = true
+
+[tool.uv.sources]
+torch = { index = "pytorch-cu128" }
+torchaudio = { index = "pytorch-cu128" }
+torchvision = { index = "pytorch-cu128" }

t3_nvfp4.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5ea2d1d1441f657b383d8dad54bbae3778447d675822fb15918cb99b3587f1e
+size 341439040