§09
Architecture swap
5—10%
LayerNorm → RMSNorm
Speedup 5—10% decode · Effort 2—5 days (includes fine-tuning) · Risk Medium
What it is
Standard LayerNorm computes mean and variance, then normalises:
LayerNorm(x) = γ · (x − mean) / sqrt(variance + ε) + β
RMSNorm (Root Mean Square Normalisation) drops the mean-centering:
RMSNorm(x) = γ · x / sqrt(mean(x²) + ε)
Why it matters for AR speech
- One fewer reduction operation per layer per token. In a 30-layer model doing 300 decode steps, that's 18,000 fewer reduction ops.
- 5—10% end-to-end speedup in our benchmarks on AR decode. Small per-op saving, but it compounds.
- Reduction operations (computing mean) are expensive on GPUs because they require cross-thread synchronisation within a warp.
How to swap
class RMSNorm(nn.Module):
def __init__(self, dim, eps=1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(dim))
self.eps = eps
def forward(self, x):
norm = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
return x * norm * self.weight
for name, module in model.named_modules():
if isinstance(module, nn.LayerNorm):
rms = RMSNorm(module.normalized_shape[0], module.eps)
rms.weight.data = module.weight.data
parent = dict(model.named_modules())[name.rsplit('.', 1)[0]]
setattr(parent, name.split('.')[-1], rms)
Why this is "High" effort
The swap itself takes 30 minutes. But:
- Dropping the bias term (β) means this is not a drop-in replacement. You need 1,000—5,000 fine-tuning steps to recover quality — that means training infrastructure, data pipeline, and quality evaluation.
- Some models (GPT-2 style) rely on the mean-centering for training stability. Test on a validation set before shipping.
- Modern architectures (LLaMA, Mistral) already use RMSNorm. This optimisation applies mainly to older GPT-2/BERT-style speech models like XTTS.
- The 5—10% speedup is modest compared to the effort. Prioritise this only after exhausting medium-effort options.
§10
Training · teacher—student
3—8×
Knowledge distillation
Speedup 3—8× · Effort 1—3 weeks · Risk High
What it is
Train a smaller "student" model to mimic a larger "teacher" model's output distribution, rather than training on ground-truth data directly. The student learns from the teacher's soft probability distributions, which contain richer information than hard labels.
Why it matters for AR speech
AR speech models are often over-parameterised. A 30-layer, 1024-dim GPT can be distilled into a 12-layer, 512-dim model with 80—90% of the quality at 5—8× the speed.
Distillation strategies for AR speech
1. Logit-level distillation. The student learns to match the teacher's token-level probability distribution.
loss_kd = F.kl_div(
F.log_softmax(student_logits / temperature, dim=-1),
F.softmax(teacher_logits / temperature, dim=-1),
reduction='batchmean'
) * (temperature ** 2)
loss = alpha * loss_kd + (1 - alpha) * loss_ce
2. Hidden-state distillation. Map student hidden states to teacher hidden states at corresponding layers. This transfers internal representations, not just outputs.
proj = nn.Linear(student_dim, teacher_dim)
loss_hidden = F.mse_loss(proj(student_hidden), teacher_hidden)
3. Speculative decoding (inference-time distillation). Use a small draft model to generate candidate tokens in parallel, then verify with the large model in a single forward pass. This gives you the quality of the large model at closer to the speed of the small model.
Draft model (fast): generates 4 candidate tokens
Teacher model: verifies all 4 in one forward pass
Accept/reject: keep accepted tokens, resample from rejection point
Speculative decoding is especially promising for AR TTS because speech token distributions tend to be peaky (low entropy), meaning the draft model's acceptance rate is high.
Why this is "High" effort
- Requires full training infrastructure — data pipeline, GPU compute, hyperparameter tuning.
- Student architecture design is non-trivial — how many layers? What dimension? Which layers to align?
- Quality evaluation is expensive — you need MOS tests, not just loss curves.
- Training time: 1—2 weeks of GPU compute for a full distillation run.
- The student might fail on edge cases (rare languages, unusual speaker characteristics) that the teacher handles.
Practical tips
- Start with 50% layer reduction and 75% dimension — this is the sweet spot for quality/speed.
- Use the teacher's generated audio (not ground truth) as training data — this reduces train/inference mismatch.
- Fine-tune the student on real data after distillation for 10—20% of the original training steps.
Custom Triton kernels
Speedup 10—20% on specific bottlenecks · Effort 3—7 days per kernel · Risk Medium-high
What it is
Triton is a Python-based GPU programming language from OpenAI that lets you write custom CUDA kernels without touching C++/CUDA directly. You write Python-like code that compiles to PTX (GPU assembly).
Why it matters for AR speech
Standard PyTorch ops leave performance on the table in two ways:
- Unfused operations: A sequence like
layernorm → linear → gelu → linear launches 4 separate kernels, each reading/writing to HBM.
- Suboptimal kernels: Generic kernels can't exploit model-specific properties (fixed head dim, known sparsity patterns, etc.).
Triton lets you write fused, model-specific kernels that keep data in SRAM.
Key kernels worth writing for AR speech
1. Fused RMSNorm + Linear
@triton.jit
def fused_rmsnorm_linear_kernel(
X, W, Out,
stride_xm, stride_xn, stride_wn, stride_wk,
M, N, K,
eps: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
):
row = tl.program_id(0)
x_ptrs = X + row * stride_xm + tl.arange(0, BLOCK_N)
x = tl.load(x_ptrs, mask=tl.arange(0, BLOCK_N) < N)
rms = tl.sqrt(tl.sum(x * x) / N + eps)
x_norm = x / rms
for k in range(0, K, BLOCK_K):
w = tl.load(W + tl.arange(0, BLOCK_N)[:, None] * stride_wn
+ (k + tl.arange(0, BLOCK_K))[None, :])
acc = tl.sum(x_norm[:, None] * w, axis=0)
tl.store(Out + row * K + k + tl.arange(0, BLOCK_K), acc)
2. Fused softmax + top-K sampling. AR decoding always ends with softmax → sampling. Fusing this saves one full read/write of the vocabulary logits.
@triton.jit
def fused_softmax_topk_kernel(logits, output_token, temperature, top_k, V: tl.constexpr):
offs = tl.arange(0, V)
x = tl.load(logits + offs) / temperature
x = x - tl.max(x)
exp_x = tl.exp(x)
probs = exp_x / tl.sum(exp_x)
3. KV cache update kernel. The KV cache append operation (insert new K, V at position t) is a memory-bound operation that benefits from a custom kernel.
@triton.jit
def kv_cache_append(
cache, new_kv, position,
num_heads, head_dim,
BLOCK: tl.constexpr,
):
head = tl.program_id(0)
offs = tl.arange(0, BLOCK)
mask = offs < head_dim
src = tl.load(new_kv + head * head_dim + offs, mask=mask)
tl.store(cache + head * MAX_SEQ * head_dim + position * head_dim + offs, src, mask=mask)
Triton vs hand-written CUDA
| Aspect | Triton | CUDA C++ |
|---|
| Development speed | 5—10× faster | Slow |
| Performance | 85—95% of CUDA | 100% |
| Maintainability | Python, readable | Complex |
| Auto-tuning | Built-in triton.autotune | Manual |
| Debugging | Easier | Painful |
Why this is "High" effort
- Each kernel takes 3—7 days to write, test, and tune.
- Requires GPU programming knowledge (memory hierarchy, tiling, occupancy).
- Gains are incremental (10—20% per kernel) — you need to profile first to identify the actual bottleneck.
- Triton kernels can't be used inside TensorRT engines — choose one path or the other.
- For AR speech inference, Triton gets you 90% of the way with 10% of the CUDA effort. But it's still significant effort.