fix(audio): size RoPE table to the reachable position range

requirements-audio-test.txt

@@ -6,3 +6,4 @@ pyloudnorm
 pyroomacoustics
 huggingface_hub[cli]
 dscaper>=1.7.7
+sphn>=0.2.0; python_version >= "3.10"

rope_capacity_0.5.3.patch

@@ -0,0 +1,193 @@
+diff --git a/requirements-audio-test.txt b/requirements-audio-test.txt
+index 792e1d8..7c7eae3 100644
+--- a/requirements-audio-test.txt
++++ b/requirements-audio-test.txt
+@@ -6,3 +6,4 @@ pyloudnorm
+ pyroomacoustics
+ huggingface_hub[cli]
+ dscaper>=1.7.7
++sphn>=0.2.0; python_version >= "3.10"
+diff --git a/src/sdialog/audio/tts/dia2/_dia2/core/layers.py b/src/sdialog/audio/tts/dia2/_dia2/core/layers.py
+index ee495ff..866d3c7 100644
+--- a/src/sdialog/audio/tts/dia2/_dia2/core/layers.py
++++ b/src/sdialog/audio/tts/dia2/_dia2/core/layers.py
+@@ -39,6 +39,20 @@ class RotaryEmbedding(nn.Module):
+         self.register_buffer("cos_cached", emb.cos(), persistent=False)
+         self.register_buffer("sin_cached", emb.sin(), persistent=False)
+
++    def ensure_capacity(self, max_position: int) -> bool:
++        """Grow the cached table so positions up to ``max_position`` are valid.
++
++        Row ``i`` of the table depends only on ``i`` (it is ``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 unaffected.
++        """
++        needed = int(max_position) + 1
++        if int(self.max_seq_len) >= needed:
++            return False
++        self.max_seq_len = needed
++        self._update_cos_sin_cache()
++        return True
++
+     def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor:
+         cos = self.cos_cached[position_ids].unsqueeze(2).to(dtype=x.dtype)
+         sin = self.sin_cached[position_ids].unsqueeze(2).to(dtype=x.dtype)
+diff --git a/src/sdialog/audio/tts/dia2/_dia2/core/model.py b/src/sdialog/audio/tts/dia2/_dia2/core/model.py
+index 74494f5..ebcbc58 100644
+--- a/src/sdialog/audio/tts/dia2/_dia2/core/model.py
++++ b/src/sdialog/audio/tts/dia2/_dia2/core/model.py
+@@ -9,6 +9,7 @@ from torch import nn
+ from ..config import DiaConfig
+ from .cache import KVCache
+ from .depformer import Depformer
++from .layers import RotaryEmbedding
+ from .precision import Precision
+ from .transformer import TransformerDecoder
+
+@@ -33,6 +34,29 @@ class Dia2Model(nn.Module):
+         self.depformer = Depformer(config, precision, device)
+         self._cast_norms_to_compute()
+
++    def ensure_position_capacity(self, max_position: int) -> int:
++        """Grow every rotary table so absolute positions up to ``max_position`` are valid.
++
++        ``max_context_steps`` bounds how many steps are *generated*, but with a
++        voice-clone prefix the generation loop runs *absolute* positions from
++        ``prefix_len - 1`` up to ``prefix_len + max_context_steps - 2``. The rotary
++        table is indexed by that absolute position, so it must cover the whole
++        range, not just ``max_context_steps``.
++        """
++        grown = 0
++        for module in self.modules():
++            if isinstance(module, RotaryEmbedding) and module.ensure_capacity(max_position):
++                grown += 1
++        self._position_capacity = max(
++            int(getattr(self, "_position_capacity", 0)), int(max_position) + 1
++        )
++        return grown
++
++    @property
++    def position_capacity(self) -> int:
++        """Largest absolute position the rotary tables can currently serve."""
++        return int(getattr(self, "_position_capacity", self.config.runtime.max_context_steps))
++
+     def init_state(self, batch_size: int, device: torch.device, max_steps: int) -> DecodeState:
+         transformer_cache = self.transformer.init_cache(batch_size, device, max_steps)
+         depformer_cache = self.depformer.init_cache(batch_size, device, self.depformer.num_depth)
+diff --git a/src/sdialog/audio/tts/dia2/_dia2/runtime/generator.py b/src/sdialog/audio/tts/dia2/_dia2/runtime/generator.py
+index f7708e8..f43953b 100644
+--- a/src/sdialog/audio/tts/dia2/_dia2/runtime/generator.py
++++ b/src/sdialog/audio/tts/dia2/_dia2/runtime/generator.py
+@@ -106,6 +106,10 @@ def build_initial_state(
+         prefix_len = delayed.shape[1]
+     limit = runtime.config.runtime.max_context_steps
+     total_steps = max(limit + prefix_len + 1, limit)
++    # The KV cache is sized for the prefix, but the rotary table is not: the loop
++    # below reaches absolute position prefix_len + limit - 2. Size it to the range
++    # that is actually reachable (no-op when it already fits).
++    runtime.model.ensure_position_capacity(prefix_len + limit)
+     decode_state = runtime.model.init_state(branches, runtime.device, total_steps)
+     audio_buf = torch.full(
+         (branches, dep_q, total_steps),
+@@ -314,7 +318,7 @@ def run_generation_loop(
+                 break
+             if t + 1 >= audio_buf.shape[-1]:
+                 break
+-            if t >= max_context:  # a voice prefix shifts t: positions must stay within the RoPE/KV cache
++            if t >= runtime.model.position_capacity:  # defensive: never index past the rotary table
+                 break
+             generation.reset_dep_cache()
+             positions.fill_(t)
+diff --git a/tests/test_audio_rope_capacity.py b/tests/test_audio_rope_capacity.py
+new file mode 100644
+index 0000000..180cb1f
+--- /dev/null
++++ b/tests/test_audio_rope_capacity.py
+@@ -0,0 +1,87 @@
++"""
++Regression tests for the rotary (RoPE) position-table capacity.
++
++``max_context_steps`` bounds how many steps are *generated*, but with a voice-clone
++prefix the generation loop runs *absolute* positions from ``prefix_len - 1`` up to
++``prefix_len + max_context_steps - 2``. ``RotaryEmbedding.forward`` gathers
++``cos_cached[position_ids]``, so the table must cover that whole absolute range --
++otherwise long references index past the end and CUDA raises a device-side assert
++(``index out of bounds``), which also poisons the CUDA context for every later
++generation in the process.
++
++These tests are CPU-only and need no model weights.
++"""
++import sys
++
++import pytest
++
++if sys.version_info < (3, 10):
++    pytest.skip("Dia2 requires Python >= 3.10", allow_module_level=True)
++
++torch = pytest.importorskip("torch")
++# Importing anything under the vendored package runs its ``__init__``, which
++# pulls in runtime.audio_io -> sphn. Skip when that optional dependency is
++# absent (same guard as tests/test_audio_synthesis.py::test_vendored_dia2_imports).
++pytest.importorskip("sphn")
++
++from sdialog.audio.tts.dia2._dia2.core.layers import RotaryEmbedding  # noqa: E402
++
++
++def _rope(max_seq_len=1500, head_dim=64):
++    return RotaryEmbedding(head_dim, 10000, 10000, max_seq_len=max_seq_len)
++
++
++def test_ensure_capacity_grows_table():
++    rope = _rope(max_seq_len=1500)
++    assert rope.ensure_capacity(2265) is True
++    assert rope.max_seq_len == 2266
++    assert rope.cos_cached.shape[0] == 2266
++    assert rope.sin_cached.shape[0] == 2266
++
++
++def test_ensure_capacity_is_noop_when_it_already_fits():
++    rope = _rope(max_seq_len=1500)
++    cos_before, sin_before = rope.cos_cached.clone(), rope.sin_cached.clone()
++    assert rope.ensure_capacity(1499) is False
++    assert rope.max_seq_len == 1500
++    assert torch.equal(rope.cos_cached, cos_before)
++    assert torch.equal(rope.sin_cached, sin_before)
++
++
++def test_ensure_capacity_is_idempotent():
++    rope = _rope(max_seq_len=1500)
++    assert rope.ensure_capacity(2265) is True
++    assert rope.ensure_capacity(2265) is False
++    assert rope.max_seq_len == 2266
++
++
++def test_growing_leaves_existing_rows_bit_identical():
++    """Row ``i`` is ``cos(i * inv_freq)``, so extending the range must not change
++    rows that already existed: generations that used to fit stay byte-identical."""
++    small, big = _rope(max_seq_len=1500), _rope(max_seq_len=2266)
++    assert torch.equal(small.cos_cached, big.cos_cached[:1500])
++    assert torch.equal(small.sin_cached, big.sin_cached[:1500])
++
++
++def test_forward_accepts_positions_past_max_context_steps():
++    """The actual failure mode: a position at/after ``max_context_steps`` used to
++    index past the table."""
++    rope = _rope(max_seq_len=1500)
++    x = torch.randn(1, 1, 1, 64)
++
++    # before growing, position 1500 is out of range for a 1500-row table
++    assert rope.cos_cached.shape[0] == 1500
++    with pytest.raises(IndexError):
++        rope(x, torch.tensor([[1500]]))
++
++    assert rope.ensure_capacity(2265) is True
++    out = rope(x, torch.tensor([[1500]]))  # now valid
++    assert out.shape == x.shape
++    assert torch.isfinite(out).all()
++
++
++def test_model_exposes_position_capacity():
++    from sdialog.audio.tts.dia2._dia2.core.model import Dia2Model
++
++    assert hasattr(Dia2Model, "ensure_position_capacity")
++    assert hasattr(Dia2Model, "position_capacity")

src/sdialog/audio/tts/dia2/_dia2/core/layers.py

@@ -39,6 +39,20 @@ def _update_cos_sin_cache(self) -> None:
         self.register_buffer("cos_cached", emb.cos(), persistent=False)
         self.register_buffer("sin_cached", emb.sin(), persistent=False)
 
+    def ensure_capacity(self, max_position: int) -> bool:
+        """Grow the cached table so positions up to ``max_position`` are valid.
+
+        Row ``i`` of the table depends only on ``i`` (it is ``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 unaffected.
+        """
+        needed = int(max_position) + 1
+        if int(self.max_seq_len) >= needed:
+            return False
+        self.max_seq_len = needed
+        self._update_cos_sin_cache()
+        return True
+
     def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor:
         cos = self.cos_cached[position_ids].unsqueeze(2).to(dtype=x.dtype)
         sin = self.sin_cached[position_ids].unsqueeze(2).to(dtype=x.dtype)

src/sdialog/audio/tts/dia2/_dia2/core/model.py

@@ -9,6 +9,7 @@
 from ..config import DiaConfig
 from .cache import KVCache
 from .depformer import Depformer
+from .layers import RotaryEmbedding
 from .precision import Precision
 from .transformer import TransformerDecoder
 
@@ -33,6 +34,29 @@ def __init__(
         self.depformer = Depformer(config, precision, device)
         self._cast_norms_to_compute()
 
+    def ensure_position_capacity(self, max_position: int) -> int:
+        """Grow every rotary table so absolute positions up to ``max_position`` are valid.
+
+        ``max_context_steps`` bounds how many steps are *generated*, but with a
+        voice-clone prefix the generation loop runs *absolute* positions from
+        ``prefix_len - 1`` up to ``prefix_len + max_context_steps - 2``. The rotary
+        table is indexed by that absolute position, so it must cover the whole
+        range, not just ``max_context_steps``.
+        """
+        grown = 0
+        for module in self.modules():
+            if isinstance(module, RotaryEmbedding) and module.ensure_capacity(max_position):
+                grown += 1
+        self._position_capacity = max(
+            int(getattr(self, "_position_capacity", 0)), int(max_position) + 1
+        )
+        return grown
+
+    @property
+    def position_capacity(self) -> int:
+        """Largest absolute position the rotary tables can currently serve."""
+        return int(getattr(self, "_position_capacity", self.config.runtime.max_context_steps))
+
     def init_state(self, batch_size: int, device: torch.device, max_steps: int) -> DecodeState:
         transformer_cache = self.transformer.init_cache(batch_size, device, max_steps)
         depformer_cache = self.depformer.init_cache(batch_size, device, self.depformer.num_depth)

src/sdialog/audio/tts/dia2/_dia2/runtime/generator.py

@@ -106,6 +106,10 @@ def build_initial_state(
         prefix_len = delayed.shape[1]
     limit = runtime.config.runtime.max_context_steps
     total_steps = max(limit + prefix_len + 1, limit)
+    # The KV cache is sized for the prefix, but the rotary table is not: the loop
+    # below reaches absolute position prefix_len + limit - 2. Size it to the range
+    # that is actually reachable (no-op when it already fits).
+    runtime.model.ensure_position_capacity(prefix_len + limit)
     decode_state = runtime.model.init_state(branches, runtime.device, total_steps)
     audio_buf = torch.full(
         (branches, dep_q, total_steps),
@@ -314,7 +318,7 @@ def run_generation_loop(
                 break
             if t + 1 >= audio_buf.shape[-1]:
                 break
-            if t >= max_context:  # a voice prefix shifts t: positions must stay within the RoPE/KV cache
+            if t >= runtime.model.position_capacity:  # defensive: never index past the rotary table
                 break
             generation.reset_dep_cache()
             positions.fill_(t)

tests/test_audio_rope_capacity.py

@@ -0,0 +1,87 @@
+"""
+Regression tests for the rotary (RoPE) position-table capacity.
+
+``max_context_steps`` bounds how many steps are *generated*, but with a voice-clone
+prefix the generation loop runs *absolute* positions from ``prefix_len - 1`` up to
+``prefix_len + max_context_steps - 2``. ``RotaryEmbedding.forward`` gathers
+``cos_cached[position_ids]``, so the table must cover that whole absolute range --
+otherwise long references index past the end and CUDA raises a device-side assert
+(``index out of bounds``), which also poisons the CUDA context for every later
+generation in the process.
+
+These tests are CPU-only and need no model weights.
+"""
+import sys
+
+import pytest
+
+if sys.version_info < (3, 10):
+    pytest.skip("Dia2 requires Python >= 3.10", allow_module_level=True)
+
+torch = pytest.importorskip("torch")
+# Importing anything under the vendored package runs its ``__init__``, which
+# pulls in runtime.audio_io -> sphn. Skip when that optional dependency is
+# absent (same guard as tests/test_audio_synthesis.py::test_vendored_dia2_imports).
+pytest.importorskip("sphn")
+
+from sdialog.audio.tts.dia2._dia2.core.layers import RotaryEmbedding  # noqa: E402
+
+
+def _rope(max_seq_len=1500, head_dim=64):
+    return RotaryEmbedding(head_dim, 10000, 10000, max_seq_len=max_seq_len)
+
+
+def test_ensure_capacity_grows_table():
+    rope = _rope(max_seq_len=1500)
+    assert rope.ensure_capacity(2265) is True
+    assert rope.max_seq_len == 2266
+    assert rope.cos_cached.shape[0] == 2266
+    assert rope.sin_cached.shape[0] == 2266
+
+
+def test_ensure_capacity_is_noop_when_it_already_fits():
+    rope = _rope(max_seq_len=1500)
+    cos_before, sin_before = rope.cos_cached.clone(), rope.sin_cached.clone()
+    assert rope.ensure_capacity(1499) is False
+    assert rope.max_seq_len == 1500
+    assert torch.equal(rope.cos_cached, cos_before)
+    assert torch.equal(rope.sin_cached, sin_before)
+
+
+def test_ensure_capacity_is_idempotent():
+    rope = _rope(max_seq_len=1500)
+    assert rope.ensure_capacity(2265) is True
+    assert rope.ensure_capacity(2265) is False
+    assert rope.max_seq_len == 2266
+
+
+def test_growing_leaves_existing_rows_bit_identical():
+    """Row ``i`` is ``cos(i * inv_freq)``, so extending the range must not change
+    rows that already existed: generations that used to fit stay byte-identical."""
+    small, big = _rope(max_seq_len=1500), _rope(max_seq_len=2266)
+    assert torch.equal(small.cos_cached, big.cos_cached[:1500])
+    assert torch.equal(small.sin_cached, big.sin_cached[:1500])
+
+
+def test_forward_accepts_positions_past_max_context_steps():
+    """The actual failure mode: a position at/after ``max_context_steps`` used to
+    index past the table."""
+    rope = _rope(max_seq_len=1500)
+    x = torch.randn(1, 1, 1, 64)
+
+    # before growing, position 1500 is out of range for a 1500-row table
+    assert rope.cos_cached.shape[0] == 1500
+    with pytest.raises(IndexError):
+        rope(x, torch.tensor([[1500]]))
+
+    assert rope.ensure_capacity(2265) is True
+    out = rope(x, torch.tensor([[1500]]))  # now valid
+    assert out.shape == x.shape
+    assert torch.isfinite(out).all()
+
+
+def test_model_exposes_position_capacity():
+    from sdialog.audio.tts.dia2._dia2.core.model import Dia2Model
+
+    assert hasattr(Dia2Model, "ensure_position_capacity")
+    assert hasattr(Dia2Model, "position_capacity")