voice_audio_multi/docs/archive/echo_tts_dit_status.md

Echo-TTS DiT: implementation status

What exists

Component	State	Verification
Byte tokenizer + WhisperD normalisation	complete	executed, output checked by hand
PCA forward / inverse	complete	executed, cross-checked against numpy
Flattening-point crop	complete	executed
Euler dual-CFG sampler	complete	executed, matches a numpy transcription of inference.py to 6.6e-07
Timestep embedding	complete	matches model.py exactly (0.0 diff)
Attention mask construction	complete	layout checked against upstream cat() semantics and ggml constraints
DiT graph (encoders, joint attention, adaLN, blocks)	written	compiles against real headers; never executed
Weight loading (1,117 tensors)	written	compiles; tensor names unconfirmed against a real checkpoint
Conditioning / denoiser graph execution	written	compiles; never executed
Fish codec z_q seam	not started	—
Session integration	not started	—

The distinction in that last column is the important one. Everything above the line was run and compared against a reference. Everything below it has only been type-checked. A clean compile here means the framework APIs are used correctly; it says nothing about whether the numbers are right.

Design decisions worth reviewing

Flash attention in the DiT joint attention

joint_attention uses ggml_flash_attn_ext, which never materialises the (lanes, heads, seq, keys) scores tensor. That tensor was the largest per-request allocation in the model:

Case	Keys	Scores tensor removed
Typical (64 text bytes, 10 s reference)	793	97 MB per attention
Long text, 30 s reference	1569	193 MB per attention
Worst case (768 text, 5 min reference)	3008	370 MB per attention

Live across 24 blocks with ggml_gallocr reuse, the practical saving is a few hundred MB to over a gigabyte, and flash attention is also faster.

This was initially written with the explicit lowering on the belief that the speaker-unconditional CFG lane produces fully masked rows, which would make -inf softmax to NaN. That was wrong: make_denoiser_mask leaves the self block of every row unmasked, so a query always attends to at least its own 640 positions and no row can be fully masked.

Two details the flash path requires. The mask must be F16, so the masked value is -65000 rather than -1e9; the latter converts to -inf in F16, which would reintroduce exactly the NaN hazard the explicit path was chosen to avoid. And q->ne[2] % mask->ne[2] and q->ne[3] % mask->ne[3] must both be zero, which holds because the mask carries a singleton head axis and matches the lane count.

Set AUDIOCPP_ECHO_TTS_NO_FLASH=1 to fall back to the explicit lowering and F32 mask, for A/B comparison without a rebuild.

The two encoders still use the explicit lowering. Their sequences are short (a few hundred tokens at most) so the scores tensors are small, and the speaker encoder is causal with no explicit mask, which the flash path rejects.

Speaker references are encoded in chunks

encode_speaker splits the reference into ~29.7 s chunks (640 latents x 2048 samples), zero-pads the last one, and concatenates the per-chunk latents, following inference.py::get_speaker_latent_and_mask. Upstream's comment calls that the longest chunk seen in training, so this is a fidelity matter as much as a memory one -- encoding several minutes in a single pass is a different computation from what the model saw.

The memory difference is large, because the Fish encoder's first stages run at the full 44.1 kHz rate. A single 64-channel activation is 0.34 GB for one chunk against 3.04 GB for a 4m29s reference encoded in one pass, and several such tensors are live at once. Fixed-size chunks also mean one encode graph is built and reused across all chunks.

After chunking, the dominant per-request allocation at long reference lengths is the persistent KV cache: 0.59 GB at 4m29s and 0.65 GB at the 297 s maximum, stored F32. Halving it to F16 is the obvious next step if that ever matters.

KV cache as a separate backend buffer

The conditioning encoders and the denoiser are separate graphs so the encoders run once per request rather than once per sampler step. They share the cached projections through tensors allocated in their own ggml_context and backend buffer, referenced as leaves by both graphs. ggml_gallocr leaves already-allocated tensors alone, so the conditioning graph writes into them with ggml_cpy and the denoiser graph reads them directly.

Consequence: changing text length or speaker length invalidates the cache and every graph built against it. prepare_conditioning tears all of it down and rebuilds, which is correct but means a request with new conditioning pays full graph construction. Acceptable given that a 40-step sample dominates.

Speaker KV scaling round-trips through the host

scale_speaker_kv reads the cached tensors back, scales, and re-uploads, because the cache has no graph attached. This runs at most twice per request (once to apply, once to undo at the threshold) and touches at most 24 layers x 2 tensors. It is not on the per-step path. If it ever shows up in a profile, the fix is a tiny scaling graph rather than a host round trip.

Things most likely to be wrong

Listed in rough order of how much damage they would do and how hard they would be to spot without a parity run:

1. Tensor names. Derived from model.py's module structure, corroborated by a parameter count matching the published file size to ten digits, but never resolved against an actual checkpoint. convert_echo_tts.py --model-dir ... settles this in seconds and prints exactly what is wrong if anything is.
2. Half-head RoPE. Heads 0-7 rotate, 8-15 do not. Implemented as slice/rope/concat on the head axis. Wrong here means plausible-sounding but incorrect audio, with no shape error.
3. Rotary pairing convention. GGML_ROPE_TYPE_NORMAL (interleaved), matching upstream's complex view of adjacent pairs. The in-tree rf_dit.cpp uses NEOX, so copying from it would be wrong.
4. Speaker patchify reshape. Folding patch_size frames into the feature axis assumes row-major frame-then-channel ordering. A transposed reading would still produce correct shapes.
5. adaLN chunk order. shift, scale, gate from cond_embed.chunk(3, -1). A permutation here is silent.

Items 2-5 are all caught by the per-block parity dumps from tools/community_models/echo_tts/echo_tts_reference.py, which is why that script dumps per-block activations at a fixed timestep rather than only the final output.

Next steps

1. Run convert_echo_tts.py against the real checkpoint to confirm item 1.
2. Build on a machine with a GPU and run the parity comparison per block.
3. Split fish_audio/codec.cpp's build_decode_quantizer at the post_module boundary and expose the summed quantized term from build_encode_quantizer, giving Echo decode and native speaker encoding.
4. Wire the session: tokenize, encode speaker, sample, PCA inverse, decode, crop.