Learnings: PersonaPlex/Moshi full-duplex S2S export + real-time ORT inference (#368)

[justinchuby]

Summary

Tracking/knowledge issue capturing the learnings from adding PersonaPlex / Moshi
full-duplex speech-to-speech support to mobius (Kyutai Mimi codec + Moshi temporal/depformer LM),
plus the real-time browser voice-chat demo. Work landed in #368.

The goal is to make these (sometimes hard-won) findings discoverable for the next
person who exports a streaming codec / dual-stream speech LM, or who needs ORT
real-time inference in Python.

---

1. Architecture (what the model actually is)

• Mimi codec (SEANet + codec transformer): encoder (B,1,T)->(B,8,Tf),
  decoder (B,8,Tf)->(B,1,T), 12.5 Hz frame rate, 8 codebooks.
  Decoder uses ELU residual blocks (not SnakeBeta) + downsample/upsample.
• Temporal transformer: dim 4096, 32 layers, 32 heads, head 128, RoPE θ=1e4,
  3000-frame sliding causal context. RMSNorm(eps=1e-8, fp32), SwiGLU
  (linear_in 4096->2*11264, linear_out 11264->4096). Fused QKV in_proj_weight (3*4096,4096).
  17-channel dual stream: ch0=text, ch1..8=agent audio, ch9..16=user audio,
  per-channel delays [0,0,1,1,1,1,1,1,1,0,1,1,1,1,1,1,1].
• Depformer: dim 1024, 6 layers, 16 heads, context 8, no RoPE,
  weights_per_step=16 (stacked per-substep linear weights via multi_linear).
  Resets KV every temporal frame; 16 substeps generate the 8 agent + 8 user codes.

2. Export gotchas (codec + LM)

• LayerNorm epsilon: Kyutai/Mimi codec transformer uses eps=1e-5
  (create_norm_fn), but mobius LayerNorm defaults to 1e-6. The default was
  the single root cause of an early codec parity failure — pass layer_norm_eps
  explicitly. (stored as a repo memory)
• Causal Conv1d padding: Kyutai SEANet causal convs pad
  left = effective_kernel - stride, not kernel - 1. The latter is wrong for
  strided Mimi convs. (stored as a repo memory)
• Mimi fp16 Conv bug -> fix: exporting Mimi in fp16 broke because the conv I/O
  dtype didn't match. Fix (5a25dd0): cast audio I/O float32<->compute dtype inside
  MimiEncoder/DecoderModel while keeping ONNX graph I/O as float32 PCM. fp16 decoder
  then hits ~71 dB SNR. This lets the codec run fp16 alongside the fp16 LM.
• TF32: CUDA parity requires use_tf32=0; otherwise the matmuls diverge from the
  PyTorch fp32 reference enough to fail atol=1e-4.

3. Real-time inference learnings (the big ones)

3a. The temporal KV cache must be device-resident (ORT IO binding)

Feeding the temporal KV cache as numpy each frame makes ORT copy the entire
growing cache host->device every step — an O(N) cost that dominates per-frame
time as the conversation lengthens (per-frame climbed 60ms -> 200ms+, crossing the
80ms real-time budget after ~20s). GQA compute is only ~6ms; the host copy was the
whole problem.

Fix: ORT IO binding — bind present.* outputs to device and reuse them as the
next frame's past.* inputs (double-buffering). The cache never touches host memory;
hidden/text_logits still come back to host for the depformer + sampling.
Result: per-frame temporal cost flat ~6ms regardless of length; end-to-end flat
~49ms/frame (was 60->200), p90 53ms, RTF ~0.61. (landed in #368, commit 7928a81)

Caveat: the device cache grows ~0.5MB/frame (fine for multi-minute demos; very long
sessions would need a periodic reset+re-prime).

3b. GroupQueryAttention >> standard op.Attention on CUDA (even for MHA)

A sliding-window alternative was prototyped: re-export the temporal model with
position_ids (non-GQA path) so the cache could be windowed. Rejected — standard
opset-23 op.Attention on the CUDA EP is ~13x slower than the fused
GroupQueryAttention contrib op (no fused Flash / no in-place KV; unfused MatMul+Add+Attention
~80ms vs ~6ms at ctx=200). GQA gives a large benefit even for MHA (32 q heads = 32 kv heads).
Lesson: prefer the GQA path on CUDA; IO binding (3a), not windowing, is the right fix for the host-copy cost.

3c. GQA can't express a sliding window

With GQA, the RoPE position is derived from seqlens_k = sum(attention_mask) - 1,
bounded by the physical cache length, and cached keys are pre-rotated at absolute
positions. So you cannot slice/window the cache without corrupting RoPE — the temporal
model has no position_ids input on the GQA path. This is why 3b's windowing
required a (slower) re-export.

3d. Env gotcha — CPU-only ORT silently fakes "CUDA"

The conda onnx env has CPU-only onnxruntime; it silently runs on
CPUExecutionProvider even when CUDA is requested, producing misleading "CUDA"
benchmarks. Any GPU perf measurement must use a real onnxruntime-gpu env.
(personal env note; worth a doc line for contributors)

4. Browser real-time demo (examples/personaplex/)

• moshi_ort.py — MoshiORT inference class (shared reset_stream()/process_frame()
  front door); offline / --stream / --mic CLIs. Imports only numpy+onnxruntime
  (no mobius) so it runs under the GPU venv.
• server.py — aiohttp WebSocket server: config-first handshake, background warmup,
  4-phase voice/persona priming, single-user lock, per-frame budget logging,
  per-second stats push (frame_ms rolling avg, RTF, over_budget).
• static/index.html — 24kHz mic client with an AudioWorklet ring-buffer player
  • latency slider + collapsible live-metrics panel (jitter buffer, rebuffers, RTF).
• Persona: record / upload / preset voice priming supported.

References

• PR: Add PersonaPlex / Moshi full-duplex S2S support (Mimi codec + Moshi LM + ORT example) #368
• Branch: justinchu/personaplex-moshi
• Related memories already stored: Mimi LayerNorm eps, SEANet causal conv padding,
  Moshi depformer weights_per_step, codec transformer eps.

Possible follow-ups

• Optionally IO-bind the depformer 16-substep loop for extra jitter headroom
  (marginal: sampling needs host logits each substep).
• Periodic temporal cache reset+re-prime for very long (>10 min) sessions.
• Document the GQA-vs-op.Attention CUDA perf gap + the CPU-only-ORT env trap in
  the relevant skills.

[justinchuby]

@titaiwangms still some issues with Attention

[justinchuby]

Update: real-time perf verified end-to-end (browser demo smooth)

Two more learnings from live testing after the IO-binding fix landed:

3e. Verify the fix isolates the variable — diagnose bimodal timing as contention, not growth.
After IO binding, the live browser session still stuttered after ~30s (compute/frame ~134ms,
RTF 1.67x). But the per-frame log was bimodal (~48ms good frames alternating with ~130-159ms
bad ones), not the monotonic ramp of the original cache-growth bug. Replaying the exact
inference path headless on an idle GPU for 1500 frames showed the median rock-flat at 48ms
through cache_len 1499 (~2 min) — confirming IO binding eliminated the O(N) growth. The bad frames
were GPU contention: the live server shared a GPU with another job (and a 133GB job sat on a
sibling GPU). Moving the server to an idle GPU made it smooth (RTF ~0.6, flat). Lesson: a bimodal
latency distribution is a contention signature; a monotonic ramp is a growth/copy signature —
reproduce headless on an isolated device before "fixing" code.

3f. Downstream jitter-buffer symptoms are consequences, not causes.
The client's adaptive jitter buffer grew to 320ms with 768 rebuffers — these were consequences of
sustained RTF>1, not independent bugs. Once per-frame went flat again, the buffer settled and
rebuffers stopped. Don't chase the buffer metrics directly.

Residual (minor, deferred): on an isolated GPU, ~3.6% of frames still spike to ~115-130ms from
the CUDA allocator growing the on-device IO-bound KV buffer each frame. The browser jitter buffer
absorbs these. Mitigations if needed: chunked pre-allocation (amortize realloc) or periodic
context reset+re-prime.

[justinchuby]

Follow-up: re-examining "op.Attention 13× slower than GQA" (3b) + corrections to 3c

Measured the attention kernels in isolation on an idle GPU (real onnxruntime-gpu, fp16, B=1, 32 heads, head_dim=128, decode seq_len=1) and traced the ORT 1.27 dispatch. Summary: the 13× is mostly graph-level, not an intrinsic attention-kernel gap, and a couple of the original conclusions are now outdated.

Isolated single-op latency (device-resident I/O, kernel compute only)

config	past=200	past=1500
GQA (contrib, fused rotary)	0.029 ms	0.044 ms
ONNX Attention + float bias mask	0.040 ms	0.083 ms
ONNX Attention is_causal=1 (no mask)	0.040 ms	0.075 ms

The pure-kernel gap is only ~1.4–1.9×, and it widens with context (the masked path is O(context), GQA is not). The "80 ms vs 6 ms" in 3b was a layer-region measurement; the bulk of it came from three things around the kernel, not the kernel itself.

Why the non-GQA path was actually ~13× (structural, verified)

1. No in-place KV (the big one). The standard opset-24 Attention op has no past_present_share_buffer — its only attributes are qk_matmul_output_mode, softcap, softmax_precision, kv_num_heads, q_num_heads, scale, is_causal. So every decode step emits a brand-new present of length past+1 → re-alloc + full copy of the growing cache each step (O(N)/step, O(N²) over a stream). This is the same growth cost as 3a's host-copy bug, just on-device. GQA writes the single new token in-place into a pre-allocated max-seq buffer.
2. A sliding-window mask kills Flash → falls to MEA/unfused (scales with context).
3. Unfused QKV projections — the GQA rewrite also packs Q/K/V into one MatMul; the prototype left separate MatMul+Add.

Correction to 3a/3b implication — nonpad_kv_seqlen does not rescue decode

The attention-optimization skill says nonpad_kv_seqlens "enables Flash + past_present_share_buffer for efficient KV cache." In ORT 1.27 the op explicitly rejects it whenever past KV is present:

nonpad_kv_seqlen should not be used together with past_key and past_value inputs

So nonpad_kv_seqlen is prefill-only. On the autoregressive decode path (which always has past_key/past_value) you cannot use it, so ONNX Attention decode is limited to is_causal=1 (Flash, but still a growing separate present — no in-place buffer).

Correction to 3c — GQA can express a sliding window

3c says "GQA can't express a sliding window." The contrib GroupQueryAttention op has a local_window_size attribute; I built a GQA node with local_window_size=256 and it loads and runs on CUDAExecutionProvider. So windowing can stay on the fast GQA decode path (in-place buffer preserved) — no need to re-export to the slower op.Attention+position_ids path. The RoPE-absolute-position caveat is real (cached keys are pre-rotated), but local_window_size masks to the last W keys without slicing the physical cache, so it doesn't corrupt RoPE.

What the shipped PersonaPlex temporal model actually does (clarifying in-place KV)

To avoid confusion: the temporal transformer ships on the GQA path (the rewrite folds op.Attention → GroupQueryAttention(do_rotary=1); src/mobius/rewrite_rules/_group_query_attention.py). The "slow Attention" only ever existed in the rejected sliding-window re-export prototype.

The current example (examples/personaplex/moshi_ort.py) is device-resident but not in-place:

• The GQA rewrite never sets past_present_share_buffer, so the temporal cache is a growing separate past → present (the kv_len in _temporal_step starts at 0 and grows each frame).
• IO-binding (3a) keeps that growing cache on-device and double-buffers present.* → past.*, eliminating the host↔device copy — but ORT still allocates a one-frame-longer present on the GPU each step (~0.5 MB/frame; the residual ~3.6% spikes are the CUDA allocator growing that buffer).
• The depformer uses plain sess.run with a numpy growing cache and no IO-binding — fine because its context is 8 and resets every frame.

True in-place KV would need GQA past_present_share_buffer=1 + a fixed max-seq buffer + present aliased to the past OrtValue + seqlens_k.

Actionable follow-ups

• Fix the attention-optimization skill: nonpad_kv_seqlens is prefill-only (rejected with past KV); the standard Attention op has no in-place KV buffer.
• Optionally teach the GQA rewrite to emit local_window_size (from config.sliding_window) and past_present_share_buffer=1 for a bounded, in-place streaming cache — this would let the PersonaPlex temporal model keep memory flat without the slow re-export, and remove the residual allocator spikes.

[justinchuby]

Correction to the previous comment — the 13× is real; I found the dominant factor

My earlier isolated benchmark measured the attention op with device-resident, fixed-context I/O, which is why it only showed ~2×. That deliberately excluded the factor that actually produces the 13×. I reproduced the demo conditions (streaming loop, KV cache growing 0→N) and the 13× falls right out:

Per-frame, single temporal layer (fp16, 32 heads, head_dim=128, hidden=4096), CUDA, idle GPU:

cache_len	A: op.Attention, growing present, numpy feeds (no IO-binding)	B: op.Attention, growing present, IO-bound on device	C: GQA, in-place shared buffer
50	0.249 ms	0.028 ms	0.022 ms
200	0.721 ms	0.038 ms	0.036 ms
499	1.492 ms	0.047 ms	0.035 ms

×32 layers (full temporal per-frame estimate), ratio vs GQA:

cache_len	A_numpy ×32	B_iobind ×32	C_gqa ×32	A/C	B/C
50	8.0 ms	0.9 ms	0.7 ms	11.5×	1.3×
200	23.1 ms	1.2 ms	1.2 ms	20.1×	1.1×
499	47.7 ms	1.5 ms	1.1 ms	42.9×	1.3×

Root cause: the 13× was the KV-cache host↔device copy, not the attention kernel

The op.Attention windowing prototype was measured before IO-binding existed (IO-binding was added later as the 3a fix for the GQA path). So that prototype fed the entire growing cache as numpy every frame, and ORT copied it host→device ×32 layers each step — an O(N) cost (column A: 0.25 → 0.72 → 1.49 ms/layer, growing with the conversation). That's exactly the 11×→20×→43× ramp, and it's the same mechanism as the original 60→200ms growth in 3a.

It was effectively an apples-to-oranges comparison: un-bound op.Attention with a growing numpy cache (A) vs IO-bound GQA (C). The fair comparison is B vs C ≈ 1.1–1.3× — i.e. once the op.Attention cache is also kept device-resident, the kernel is at near-parity with GQA. This reconciles the demo's 13× with the isolated ~2× kernel gap from the previous comment: both are correct, they just measure different things.

Why GQA still "wins" in practice

GQA doesn't have a magically faster attention kernel (it's ~1.2× here). It wins because it forces a fixed-size, pre-allocated KV buffer (past_present_share_buffer), which is naturally device-resident — there's no growing-numpy-cache footgun. The standard op.Attention op allows a growing separate present, which makes it very easy to fall into the per-frame host-copy trap unless you explicitly IO-bind it.

Takeaway (revises 3b)

• The intrinsic op.Attention-vs-GQA kernel gap is ~1.2–1.3× for this MHA decode, not 13×.
• The 13× was the growing KV-cache host↔device copy from feeding numpy each frame with no IO-binding (O(N), worsens over the conversation).
• Fix is orthogonal to op choice: keep the cache device-resident (IO-binding). GQA gets this "for free" via its fixed shared buffer; op.Attention needs explicit IO-binding (and still pays for a growing on-device buffer unless you add local_window_size/bounded reuse).

[justinchuby]

The residual ~1.2–1.7× (IO-bound op.Attention vs GQA): present materialization + decode kernel

After removing the host-copy factor, IO-bound op.Attention (B) is still ~1.2–1.7× slower than GQA (C), and crucially the gap grows with context — so it's O(context), not a fixed kernel constant:

context (fixed length, fully device-resident)	GQA	op.Attention	ratio
past=200	0.029 ms	0.040 ms	1.4×
past=1500	0.044 ms	0.075 ms	1.7×

Two reasons:

1. present materialization (dominant, O(context)).
The standard op.Attention op outputs present = concat(past, new) every step — it reads the entire length-N past and writes it into a fresh present buffer, then attends over N+1. This O(N) copy is in the op's semantics: it happens every step even when the cache is IO-bound and even at fixed context. GQA, via past_present_share_buffer, writes only the single new token's K/V in-place at offset N in a fixed buffer (O(1)) and attends over the existing buffer — no concat. So the gap is "copy the whole cache each step" vs "write one token", which scales linearly with N (hence 1.4× → 1.7×).

2. Dedicated single-token decode kernel.
For seq==1 + past_present_share_buffer + head_size ∈ {64,128,256} + SM≥8.0, GQA dispatches to a hand-tuned decode kernel (XQA / Flash-decode FastDecode). head_dim=128 here hits it. The standard op.Attention decode path (is_causal=1, separate growing present) uses a more general attention kernel without that seq=1 in-place fast path.

Takeaway: GQA isn't faster because of better attention math — it's the memory model. Its in-place fixed buffer avoids the per-step O(N) present concat and unlocks the decode kernel. You can't close this last gap with op.Attention because the opset-24 op has no past_present_share_buffer (its only attrs are qk_matmul_output_mode, softcap, softmax_precision, kv_num_heads, q_num_heads, scale, is_causal). This is the fundamental reason mobius rewrites decoder attention to GQA.

Full attribution of the original demo 13×

1. Dominant: growing KV-cache host↔device copy (numpy feeds, no IO-binding) — O(N), ~11–43× over a conversation.
2. Residual ~1.2–1.7×: O(N) present concat + general (non-decode) kernel.
3. (Earlier prototype also lost Flash via the sliding-window mask and had unfused QKV projections.)