[jordiluque]

Problem

With voice-clone prefixes, long turns crash with a CUDA device-side assert
(IndexKernel.cu:94 index out of bounds), which also poisons the CUDA context so
every later generation in the same process fails.

The rotary table is sized with max_seq_len = config.runtime.max_context_steps
(core/layers.py), but the generation loop runs absolute positions:
build_initial_state allocates the KV cache as max_context_steps + prefix_len + 1
and warmup_with_prefix returns start_step = prefix_len - 1, so t reaches
prefix_len + max_context_steps - 2. RotaryEmbedding.forward gathers
cos_cached[position_ids], so any t >= max_context_steps indexes past the table.

The 0.5.2/0.5.3 guard (if t >= max_context: break) prevents the crash but
silently truncates generation whenever a window reaches absolute position
max_context_steps. With the default reference_seconds=8.0 the reachable range is
already 100 ... 1598, i.e. past 1500 — the guard is only unreachable in practice
because max_window_words=170 caps a window at ~790 generated frames, well short of
the ~1400 needed to hit it. Raising reference_seconds or max_window_words, or
passing a long (array, sr) reference, makes the guard fire and the tail of the
utterance is dropped with no warning.

Fix

Size the rotary table to the range that is actually reachable, at the one place
where it is known — right next to the existing total_steps computation:

runtime.model.ensure_position_capacity(prefix_len + limit)

RotaryEmbedding.ensure_capacity() extends max_seq_len and rebuilds the cache.
Row i depends only on i (cos(i * inv_freq), built from arange(max_seq_len)),
so extending the range leaves every existing row bit-identical — generations that
already fit are byte-for-byte unchanged. The loop guard becomes
t >= runtime.model.position_capacity, a defensive bound that never fires in normal
use.

Files changed (3 source + 1 test, +121/-1):

• src/sdialog/audio/tts/dia2/_dia2/core/layers.py — RotaryEmbedding.ensure_capacity()
• src/sdialog/audio/tts/dia2/_dia2/core/model.py — Dia2Model.ensure_position_capacity() + position_capacity
• src/sdialog/audio/tts/dia2/_dia2/runtime/generator.py — call it in build_initial_state; guard uses the real capacity
• tests/test_audio_rope_capacity.py — CPU-only regression tests (no weights, no GPU)

Tested in H100, torch 2.6.0+cu124, sdialog 0.5.3

Memory overhead is negligible: one extra cos/sin pair of (Δlen, head_dim) fp32
per attention layer (~22 MB for 1500 → 2266).

Notes

• Output is unchanged for every configuration that works today. The table does
  grow (e.g. 1500 -> 1601 at the default reference_seconds=8.0), but because row i
  depends only on i, the added rows are never read by a generation that previously
  fit — so its audio is byte-identical. Verified numerically: max |cos_1500[i] − cos_1601[i]| = 0.0 for all i < 1500. The patch only affects generations that
  would otherwise have been truncated.

[codecov]

Codecov Report

❌ Patch coverage is 55.00000% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 50.14%. Comparing base (5c77282) to head (27ae937).

Files with missing lines	Patch %	Lines
src/sdialog/audio/tts/dia2/_dia2/core/model.py	36.36%	7 Missing ⚠️
.../sdialog/audio/tts/dia2/_dia2/runtime/generator.py	0.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #136      +/-   ##
==========================================
+ Coverage   44.92%   50.14%   +5.22%     
==========================================
  Files          84       84              
  Lines       10213    10232      +19     
==========================================
+ Hits         4588     5131     +543     
+ Misses       5625     5101     -524     

Files with missing lines	Coverage Δ
src/sdialog/audio/tts/dia2/_dia2/core/layers.py	39.47% <100.00%> (+39.47%)	⬆️
.../sdialog/audio/tts/dia2/_dia2/runtime/generator.py	16.78% <0.00%> (+16.78%)	⬆️
src/sdialog/audio/tts/dia2/_dia2/core/model.py	47.05% <36.36%> (+47.05%)	⬆️

... and 22 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[sergioburdisso]

@all-contributors please, add @jordiluque as a contributor for code 👍

[allcontributors]

@sergioburdisso

I've put up a pull request to add @jordiluque! 🎉

[sergioburdisso]

Thanks @jordiluque for your fix, a new version 0.5.4 has been just released with your patch :)