From 6a98d7f651b8c2cfc5237e83ca9b0b2c1ad04464 Mon Sep 17 00:00:00 2001 From: Jordi Luque Date: Tue, 27 Oct 2026 10:24:33 +0100 Subject: [PATCH 1/2] fix(audio): size RoPE table to the reachable position range --- requirements-audio-test.txt | 1 + rope_capacity_0.5.3.patch | 193 ++++++++++++++++++ rope_capacity_0.5.3.patch.old | 175 ++++++++++++++++ .../audio/tts/dia2/_dia2/core/layers.py | 14 ++ .../audio/tts/dia2/_dia2/core/model.py | 24 +++ .../audio/tts/dia2/_dia2/runtime/generator.py | 6 +- tests/test_audio_rope_capacity.py | 87 ++++++++ 7 files changed, 499 insertions(+), 1 deletion(-) create mode 100644 rope_capacity_0.5.3.patch create mode 100644 rope_capacity_0.5.3.patch.old create mode 100644 tests/test_audio_rope_capacity.py 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/rope_capacity_0.5.3.patch b/rope_capacity_0.5.3.patch new file mode 100644 index 0000000..4aafb6c --- /dev/null +++ b/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") diff --git a/rope_capacity_0.5.3.patch.old b/rope_capacity_0.5.3.patch.old new file mode 100644 index 0000000..f611328 --- /dev/null +++ b/rope_capacity_0.5.3.patch.old @@ -0,0 +1,175 @@ +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..2b2ecc8 +--- /dev/null ++++ b/tests/test_audio_rope_capacity.py +@@ -0,0 +1,78 @@ ++""" ++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 pytest ++ ++torch = pytest.importorskip("torch") ++ ++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") 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 @@ 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) 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 ..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) 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") From 27ae937b8d64ef7540020787deed5245201b3669 Mon Sep 17 00:00:00 2001 From: Jordi Luque Date: Tue, 27 Oct 2026 10:43:39 +0100 Subject: [PATCH 2/2] fix(audio): size RoPE table to the reachable position range --- rope_capacity_0.5.3.patch.old | 175 ---------------------------------- 1 file changed, 175 deletions(-) delete mode 100644 rope_capacity_0.5.3.patch.old diff --git a/rope_capacity_0.5.3.patch.old b/rope_capacity_0.5.3.patch.old deleted file mode 100644 index f611328..0000000 --- a/rope_capacity_0.5.3.patch.old +++ /dev/null @@ -1,175 +0,0 @@ -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..2b2ecc8 ---- /dev/null -+++ b/tests/test_audio_rope_capacity.py -@@ -0,0 +1,78 @@ -+""" -+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 pytest -+ -+torch = pytest.importorskip("torch") -+ -+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")