Port T3 backbone to MLX with performance improvements and bug fixes - #532
Draft
PeteHaughie wants to merge 3 commits into
Draft
Port T3 backbone to MLX with performance improvements and bug fixes#532PeteHaughie wants to merge 3 commits into
PeteHaughie wants to merge 3 commits into
Conversation
- T3MLX model: 30-layer LLaMA decoder with LLaMA-3 RoPE scaling, SwiGLU MLP, learned position embeddings, KV cache autoregressive loop - Weight conversion: safetensors → .npz with key remapping - Achieves 62 tok/s (2.5× real-time) — 60× faster than PyTorch MPS - End-to-end pipeline: PyTorch conditioning → MLX generate → S3Gen decode - Logit correlation >0.98 vs PyTorch reference - 5 ADRs documenting architecture decisions
Bug fixes identified during pre-PR review: - Add causal mask (tril) during prefill forward pass to prevent future-token leakage into earlier positions via KV cache - Add position embeddings to autoregressive tokens via offset param in embed_speech_with_pos (BOS at pos 0, step i at pos i+1) - Guard division by zero in _top_p_sample when softmax underflows - Guard temperature near-zero in _sample_logits (clamp at 1e-6) - Remove unused top_p parameter from _sample_logits Housekeeping: - .gitignore: assets/*.npz and venv* directories - pyproject.toml: add mlx>=0.31.0 as optional dependency
PeteHaughie
marked this pull request as draft
June 21, 2026 15:36
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request introduces the initial proof-of-concept (PoC) for porting the T3 text-to-speech model's backbone to MLX, Apple's Metal-native ML framework. It establishes the technical boundary between PyTorch and MLX, provides scripts for weight conversion and benchmarking, and documents architectural decisions and success criteria. The changes focus on enabling fast, real-time inference on Apple Silicon, while keeping peripheral components in PyTorch for simplicity.
Key changes:
MLX Port Architecture and Documentation
CONTEXT.md) and five architectural decision records (ADRs) detailing the MLX port boundary, backbone construction, weight conversion, PoC scope, and success criteria. These documents clarify what is ported, the rationale, and how success is measured. [1] [2] [3] [4] [5] [6]Weight Conversion and Loading
scripts/convert_weights_mlx.py, a script to convert T3 model weights from HuggingFace safetensors format to MLX-compatible.npzfiles, with explicit key mapping for all backbone and embedding layers.MLX Inference and Benchmarking
scripts/poc_mlx.pyfor benchmarking the MLX T3 backbone, measuring forward pass and autoregressive generation speed, and verifying correct weight loading.scripts/poc_end_to_end.pyfor an end-to-end PoC that runs PyTorch-based conditioning, MLX-based T3 generation, and PyTorch S3Gen audio decoding, demonstrating the cross-framework pipeline.Dependency Management
pyproject.tomlto add MLX as an optional dependency, enabling easy installation for users with Apple Silicon.References:
[1] [2] [3] [4] [5] [6] [7] [8] [9] [10]