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/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")