460 lines
22 KiB
Python
460 lines
22 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Check depth attention against HF Llama and request-local sampling state."""
|
|
|
|
from collections import defaultdict
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from transformers import LlamaConfig, MimiConfig
|
|
from transformers.models.llama.modeling_llama import LlamaModel
|
|
from transformers.models.mimi.modeling_mimi import MimiConv1d
|
|
|
|
from tests.model_executor.models.breeze_tts_2_helpers import _request_info, _small_talker
|
|
from vllm_omni.engine.serialization import deserialize_additional_information, serialize_additional_information
|
|
from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import (
|
|
BreezeDepthDecoder,
|
|
sample_graph_logits,
|
|
sample_logits,
|
|
)
|
|
from vllm_omni.model_executor.models.breeze_tts_2.prompt import CFG_UNCOND_SUFFIX, build_breeze_prompt
|
|
from vllm_omni.model_executor.models.breeze_tts_2.reference_encoder import BreezeReferenceConv
|
|
from vllm_omni.model_executor.stage_input_processors.breeze_tts_2 import expand_cfg_prompts, talker2code2wav_async_chunk
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
@pytest.mark.parametrize("parameter", ["temperature", "top_p", "repetition_penalty"])
|
|
@pytest.mark.parametrize("value", [float("nan"), float("inf")])
|
|
def test_offline_prompt_rejects_nonfinite_sampling(mocker, parameter, value):
|
|
with pytest.raises(ValueError, match="sampling"):
|
|
build_breeze_prompt(mocker.Mock(), "Hello.", **{parameter: value})
|
|
|
|
|
|
@pytest.mark.parametrize("stride", [2, 3])
|
|
def test_reference_convolution_batch_preserves_partial_final_frames(stride):
|
|
torch.manual_seed(7)
|
|
config = MimiConfig(use_causal_conv=True, pad_mode="constant")
|
|
first = MimiConv1d(config, 1, 4, kernel_size=2 * stride, stride=stride)
|
|
last = MimiConv1d(config, 4, 4, kernel_size=4, stride=2, pad_mode="replicate", bias=False)
|
|
batched_first = BreezeReferenceConv(first, stride)
|
|
batched_last = BreezeReferenceConv(last, 2 * stride)
|
|
lengths = torch.tensor([37, 62])
|
|
batched_first.sample_lengths = batched_last.sample_lengths = lengths
|
|
waves = [torch.randn(1, 1, int(length)) for length in lengths]
|
|
padded = torch.cat([F.pad(wave, (0, 96 - wave.shape[-1])) for wave in waves])
|
|
actual = batched_last(batched_first(padded))
|
|
for row, waveform in enumerate(waves):
|
|
expected = last(first(waveform))
|
|
torch.testing.assert_close(actual[row : row + 1, :, : expected.shape[-1]], expected)
|
|
assert not actual[row, :, expected.shape[-1] :].count_nonzero()
|
|
|
|
|
|
def test_reference_conditioning_survives_transport_and_cfg_keeps_true_prompt_lengths(mocker) -> None:
|
|
tokenizer = mocker.Mock()
|
|
tokenizer.encode.side_effect = lambda text, **kwargs: list(text.encode("utf-8"))
|
|
waveform = np.zeros(1921, dtype=np.float32)
|
|
prompt = build_breeze_prompt(
|
|
tokenizer, "Hello.", "Speak happily.", ref_audio=(waveform, 24000), ref_text="Reference.", guidance_scale=4
|
|
)
|
|
companion = expand_cfg_prompts(prompt, None)[0].prompt
|
|
assert len(prompt["prompt_token_ids"]) > len(companion["prompt_token_ids"])
|
|
positive = deserialize_additional_information(serialize_additional_information(prompt["additional_information"]))
|
|
negative = deserialize_additional_information(serialize_additional_information(companion["additional_information"]))
|
|
assert positive["breeze_prompt"]["target_ids"] == list(b"[S0]<ins_bos>Speak happily.<ins_eos>Hello.")
|
|
assert negative["breeze_prompt"]["target_ids"] == list(b"[S0]Hello.")
|
|
assert positive["breeze_prompt"]["reference_ids"] == negative["breeze_prompt"]["reference_ids"]
|
|
assert positive["breeze_prompt"]["reference_frames"] == 2
|
|
assert positive["cfg_group"]["role"] == "cond"
|
|
assert negative["cfg_group"]["role"] == "uncond"
|
|
torch.testing.assert_close(positive["reference_waveform"], torch.from_numpy(waveform), atol=0, rtol=0)
|
|
torch.testing.assert_close(negative["reference_waveform"], positive["reference_waveform"], atol=0, rtol=0)
|
|
|
|
|
|
@pytest.fixture
|
|
def depth() -> BreezeDepthDecoder:
|
|
torch.manual_seed(42)
|
|
config = LlamaConfig(
|
|
hidden_size=32,
|
|
intermediate_size=64,
|
|
num_hidden_layers=2,
|
|
num_attention_heads=4,
|
|
num_key_value_heads=2,
|
|
head_dim=8,
|
|
vocab_size=35,
|
|
max_position_embeddings=16,
|
|
rope_theta=500000.0,
|
|
rope_scaling={
|
|
"rope_type": "llama3",
|
|
"factor": 32.0,
|
|
"low_freq_factor": 0.001953125,
|
|
"high_freq_factor": 0.0078125,
|
|
"original_max_position_embeddings": 16,
|
|
},
|
|
)
|
|
config.num_codebooks = 4
|
|
config.audio_embed_size = 48
|
|
config._attn_implementation = "eager"
|
|
model = BreezeDepthDecoder(config).eval()
|
|
torch.nn.init.normal_(model.codebooks_head, std=0.1)
|
|
return model
|
|
|
|
|
|
@pytest.fixture
|
|
def hf_depth(depth: BreezeDepthDecoder) -> LlamaModel:
|
|
reference = LlamaModel(depth.config).eval()
|
|
weights = {}
|
|
for name, parameter in depth.named_parameters():
|
|
if name.startswith("layers."):
|
|
if name.endswith("qkv.weight"):
|
|
layer = depth.layers[int(name.split(".")[1])]
|
|
for part, tensor in zip(
|
|
("q", "k", "v"), parameter.split([layer.q_size, layer.kv_size, layer.kv_size]), strict=True
|
|
):
|
|
weights[name.replace("qkv.weight", f"self_attn.{part}_proj.weight")] = tensor
|
|
elif name.endswith("gate_up.weight"):
|
|
for part, tensor in zip(("gate", "up"), parameter.chunk(2), strict=True):
|
|
weights[name.replace("gate_up.weight", f"mlp.{part}_proj.weight")] = tensor
|
|
else:
|
|
weights[name.replace(".o_proj.", ".self_attn.o_proj.").replace(".down_proj.", ".mlp.down_proj.")] = (
|
|
parameter
|
|
)
|
|
elif name == "norm.weight":
|
|
weights[name] = parameter
|
|
result = reference.load_state_dict(weights, strict=False)
|
|
assert result.missing_keys == ["embed_tokens.weight"]
|
|
assert result.unexpected_keys == []
|
|
return reference
|
|
|
|
|
|
@pytest.mark.parametrize("batch", [1, 2])
|
|
@torch.inference_mode()
|
|
def test_depth_cached_attention_matches_hf_full_sequence(
|
|
depth: BreezeDepthDecoder, hf_depth: LlamaModel, batch: int
|
|
) -> None:
|
|
"""Wrong cache positions or fused-QKV ordering change the reference hidden states."""
|
|
embeddings = torch.randn(batch, 4, depth.config.hidden_size)
|
|
expected = hf_depth(inputs_embeds=embeddings, use_cache=False).last_hidden_state
|
|
shape = (batch, 2, 4, 8)
|
|
caches = [(torch.zeros(shape), torch.zeros(shape)) for _ in depth.layers]
|
|
actual = []
|
|
for start, end in [(0, 2), (2, 3), (3, 4)]:
|
|
x = embeddings[:, start:end]
|
|
for layer, cache in zip(depth.layers, caches, strict=True):
|
|
x = layer(
|
|
x,
|
|
depth.positions[start:end],
|
|
cache,
|
|
depth.causal_mask[start:end, :4][None, None],
|
|
depth.rope_cos[start:end][None, None],
|
|
depth.rope_sin[start:end][None, None],
|
|
)
|
|
actual.append(depth.norm(x))
|
|
torch.testing.assert_close(torch.cat(actual, 1), expected, atol=2e-6, rtol=2e-5)
|
|
|
|
|
|
@pytest.mark.parametrize("guidance_scale", [0.5, 4.0])
|
|
@torch.inference_mode()
|
|
def test_cfg_depth_matches_hf_full_prefix_at_every_codebook(
|
|
depth: BreezeDepthDecoder, hf_depth: LlamaModel, guidance_scale: float
|
|
) -> None:
|
|
hidden = torch.randn(4, depth.config.audio_embed_size)
|
|
first = torch.tensor([7, 11])
|
|
prefix = torch.cat((hidden[:, None], depth.embed_tokens(first.repeat(2)[:, None])), dim=1)
|
|
expected = [first]
|
|
for codebook in range(depth.num_codebooks - 1):
|
|
states = hf_depth(inputs_embeds=depth.inputs_embeds_projector(prefix), use_cache=False).last_hidden_state
|
|
logits = F.linear(states[:, -1].float(), depth.codebooks_head[codebook].T)
|
|
cond, uncond = logits.chunk(2)
|
|
token = (uncond + guidance_scale * (cond - uncond))[:, : depth.vocab_size - 3].argmax(-1)
|
|
expected.append(token)
|
|
embed = depth.embed_tokens((token.repeat(2) + (codebook + 1) * depth.vocab_size)[:, None])
|
|
prefix = torch.cat((prefix, embed), dim=1)
|
|
actual = depth._generate_frame(
|
|
hidden, first, depth._allocate_cache(hidden), 0.0, 0, 1.0, torch.Generator(), guidance_scale=guidance_scale
|
|
)
|
|
torch.testing.assert_close(actual, torch.stack(expected, dim=-1), atol=0, rtol=0)
|
|
|
|
|
|
@torch.inference_mode()
|
|
def test_frame_generation_does_not_reuse_another_requests_cache(depth: BreezeDepthDecoder) -> None:
|
|
hidden = torch.randn(1, 48)
|
|
first = torch.tensor([7])
|
|
|
|
def generate(value: torch.Tensor) -> torch.Tensor:
|
|
return depth.generate_frame(
|
|
value, first, temperature=0.9, top_k=10, top_p=0.8, generator=torch.Generator().manual_seed(99)
|
|
)
|
|
|
|
expected = generate(hidden)
|
|
generate(hidden * -3)
|
|
torch.testing.assert_close(generate(hidden), expected, atol=0, rtol=0)
|
|
|
|
|
|
def test_top_k_keeps_tied_candidates_and_nucleus_keeps_crossing_token() -> None:
|
|
generator = torch.Generator().manual_seed(42)
|
|
logits = torch.tensor([[0.0, 0.0, -100.0]]).expand(128, -1)
|
|
sampled = sample_logits(logits, 1.0, 1, 1.0, generator)
|
|
assert set(sampled.tolist()) == {0, 1}
|
|
sampled = sample_logits(torch.tensor([[10.0, 0.0]]), 1.0, 0, 0.1, generator)
|
|
assert sampled.item() == 0
|
|
|
|
|
|
@pytest.mark.parametrize("parameters", [(0.0, 0, 1.0), (0.9, 50, 1.0), (0.7, 100, 0.8), (1.2, 0, 0.95)])
|
|
def test_graph_sampler_matches_scalar_filtering_with_identical_noise(parameters) -> None:
|
|
generator = torch.Generator().manual_seed(42)
|
|
logits = torch.randn(32, 2051, generator=generator)
|
|
logits[:, -3:] = -torch.inf
|
|
logits[:, 10:14] = 4.0
|
|
noise = torch.empty_like(logits).exponential_(generator=generator)
|
|
expected = sample_logits(logits, *parameters, generator, noise)
|
|
actual = sample_graph_logits(logits, torch.tensor(parameters), noise)
|
|
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
|
|
|
|
|
|
@pytest.mark.parametrize("top_p", [0.25, 0.5, 0.75])
|
|
@pytest.mark.parametrize("top_k", [0, 2])
|
|
def test_graph_sampler_excludes_token_after_exact_nucleus_boundary(top_p, top_k) -> None:
|
|
# Equal probabilities make the CDF boundary exact in float32. Bias the
|
|
# shared noise toward the first token outside that nucleus so retaining
|
|
# one extra candidate deterministically changes the sampled token.
|
|
logits = torch.tensor([[0.0, 0.0, 0.0, 0.0, -torch.inf, -torch.inf, -torch.inf]])
|
|
noise = torch.ones_like(logits)
|
|
first_excluded = int(top_p * 4)
|
|
noise[:, first_excluded] = 0.001
|
|
parameters = (1.0, top_k, top_p)
|
|
expected = sample_logits(logits, *parameters, torch.Generator().manual_seed(42), noise)
|
|
actual = sample_graph_logits(logits, torch.tensor(parameters), noise)
|
|
|
|
assert expected.item() != first_excluded
|
|
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
|
|
|
|
|
|
def test_codec_chunks_preserve_zero_frames_and_flush_each_request_once(mocker) -> None:
|
|
manager = mocker.Mock()
|
|
manager.connector.config = {"extra": {"codec_chunk_frames": 5, "initial_codec_chunk_frames": 5}}
|
|
manager.code_prompt_token_ids = defaultdict(list)
|
|
manager.request_payload = {}
|
|
first = mocker.Mock(external_req_id="first")
|
|
first.is_finished.return_value = False
|
|
second = mocker.Mock(external_req_id="second")
|
|
second.is_finished.return_value = False
|
|
zeros = {"codes": {"audio": torch.zeros(3, 16, dtype=torch.long)}}
|
|
assert talker2code2wav_async_chunk(manager, zeros, first) is None
|
|
assert talker2code2wav_async_chunk(manager, zeros, second) is None
|
|
codes = torch.arange(64).reshape(4, 16)
|
|
output = talker2code2wav_async_chunk(manager, {"codes": {"audio": codes}}, first)
|
|
expected = torch.cat((torch.zeros(3, 16, dtype=torch.long), codes))
|
|
torch.testing.assert_close(output.codes.audio.reshape(16, -1).T, expected)
|
|
final = talker2code2wav_async_chunk(manager, None, first, is_finished=True)
|
|
assert final.codes.audio.numel() == 0
|
|
assert final.meta.finished.item()
|
|
other = talker2code2wav_async_chunk(manager, None, second, is_finished=True)
|
|
torch.testing.assert_close(other.codes.audio, torch.zeros(48, dtype=torch.long))
|
|
|
|
|
|
def test_codec_chunk_ramp_keeps_frame_order_and_flushes_partial_tail(mocker) -> None:
|
|
manager = mocker.Mock()
|
|
manager.connector.config = {"extra": {"codec_chunk_frames": 5, "codec_chunk_ramp": [1, 2, 4, 5]}}
|
|
manager.code_prompt_token_ids = defaultdict(list)
|
|
manager.request_payload = {}
|
|
request = mocker.Mock(external_req_id="ramp")
|
|
request.is_finished.return_value = False
|
|
outputs = []
|
|
for index in range(15):
|
|
result = talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.full((1, 16), index)}}, request)
|
|
if result is not None:
|
|
outputs.append(result.codes.audio.reshape(16, -1).T)
|
|
tail = talker2code2wav_async_chunk(manager, None, request, is_finished=True)
|
|
outputs.append(tail.codes.audio.reshape(16, -1).T)
|
|
assert [len(part) for part in outputs] == [1, 2, 4, 5, 3]
|
|
torch.testing.assert_close(torch.cat(outputs), torch.arange(15)[:, None].expand(-1, 16))
|
|
assert tail.meta.finished.item()
|
|
|
|
|
|
@torch.inference_mode()
|
|
def test_talker_history_and_generators_follow_requests_across_reordering_and_new_admission():
|
|
model = _small_talker()
|
|
first, second = _request_info("first"), _request_info("second")
|
|
# Unequal spans must select the last hidden row of each request.
|
|
model.make_omni_output(
|
|
torch.tensor([[0.0, 1.0], [0.0, 1.0], [1.0, 0.0], [0.0, 1.0]]),
|
|
model_intermediate_buffer=[first, second],
|
|
request_token_spans=[(0, 3), (3, 4)],
|
|
)
|
|
reordered = model.make_omni_output(
|
|
torch.tensor([[0.0, 1.0], [1.0, 0.0]]),
|
|
model_intermediate_buffer=[second, first],
|
|
request_token_spans=[(0, 1), (1, 2)],
|
|
)
|
|
assert [frame[0, 0].item() for frame in reordered.multimodal_outputs["codes"]["audio"]] == [3, 2]
|
|
assert first["breeze_state"]["history"].tolist() == [[1, 2]]
|
|
assert second["breeze_state"]["history"].tolist() == [[3, 3]]
|
|
generators = model.depth_decoder.generate_frames.call_args.kwargs["generators"]
|
|
assert generators[0] is second["breeze_state"]["generator"]
|
|
assert generators[1] is first["breeze_state"]["generator"]
|
|
|
|
newcomer = _request_info("newcomer")
|
|
admitted = model.make_omni_output(
|
|
torch.tensor([[1.0, 0.0], [1.0, 0.0]]),
|
|
model_intermediate_buffer=[newcomer, first],
|
|
request_token_spans=[(0, 1), (1, 2)],
|
|
)
|
|
assert [frame[0, 0].item() for frame in admitted.multimodal_outputs["codes"]["audio"]] == [1, 1]
|
|
assert newcomer["breeze_state"]["history"].tolist() == [[1]]
|
|
assert first["breeze_state"]["history"].tolist() == [[1, 2, 1]]
|
|
assert second["breeze_state"]["history"].tolist() == [[3, 3]]
|
|
|
|
|
|
@torch.inference_mode()
|
|
def test_talker_eos_does_not_shift_live_state_or_poison_the_next_batch():
|
|
model = _small_talker()
|
|
finished, live = _request_info("finished"), _request_info("live")
|
|
output = model.make_omni_output(
|
|
torch.tensor([[-1.0, -1.0], [0.0, 1.0]]),
|
|
model_intermediate_buffer=[finished, live],
|
|
request_token_spans=[(0, 1), (1, 2)],
|
|
)
|
|
assert output.multimodal_outputs["codes"]["audio"][0].numel() == 0
|
|
assert finished["breeze_state"]["history"].numel() == 0
|
|
assert live["breeze_state"]["current"].tolist() == [[3, 3, 3]]
|
|
assert model.compute_logits(torch.zeros(2, 2)).argmax(-1).tolist() == [7, 0]
|
|
|
|
model.make_omni_output(torch.tensor([[0.0, 1.0]]), model_intermediate_buffer=[live], request_token_spans=[(0, 1)])
|
|
assert live["breeze_state"]["history"].tolist() == [[3, 3]]
|
|
assert model.compute_logits(torch.zeros(1, 2)).argmax(-1).tolist() == [0]
|
|
|
|
|
|
@pytest.mark.parametrize("temperature", [0.0, 0.9])
|
|
@torch.inference_mode()
|
|
def test_talker_mixed_eos_and_cfg_matches_independent_request_state(temperature):
|
|
"""Batch EOS decisions must not mix CFG state or change a request's RNG stream."""
|
|
finished_companion = "finished_cfg" + CFG_UNCOND_SUFFIX
|
|
live_companion = "live_cfg" + CFG_UNCOND_SUFFIX
|
|
order = ["finished_cfg", "live_plain", finished_companion, "finished_plain", live_companion, "live_cfg"]
|
|
hidden_by_id = {
|
|
"finished_cfg": [-1.0, -1.0],
|
|
finished_companion: [-1.0, -1.0],
|
|
"finished_plain": [-1.0, -1.0],
|
|
"live_plain": [0.0, 1.0],
|
|
"live_cfg": [1.0, 0.0],
|
|
live_companion: [1.0, 0.0],
|
|
}
|
|
|
|
def make_infos():
|
|
infos = {request_id: _request_info(request_id) for request_id in order}
|
|
for index, (request_id, info) in enumerate(infos.items()):
|
|
info["breeze_sampling"].update(temperature=temperature, top_k=1)
|
|
if "cfg" not in request_id:
|
|
# Interleave distinct sampler groups while each request still
|
|
# matches its independently decoded state and RNG trajectory.
|
|
info["breeze_sampling"].update(temperature=temperature * 0.5, top_p=0.8)
|
|
info["breeze_state"]["generator"].manual_seed(42 + index)
|
|
info["breeze_state"]["history"] = torch.tensor([[3, 3]])
|
|
if "cfg" in request_id:
|
|
info["breeze_prompt"]["guidance_scale"] = 4.0
|
|
if request_id.endswith(CFG_UNCOND_SUFFIX):
|
|
info["breeze_prompt"]["role"] = "uncond"
|
|
return infos
|
|
|
|
def sample_depth(hidden, first, *, temperature, generators, **kwargs):
|
|
# Consume request-local RNG state in the downstream phase as well as
|
|
# codebook-0 sampling, without requiring depth CUDA graphs on CPU.
|
|
tail = (
|
|
torch.stack([torch.randint(0, 4, (2,), generator=generator) for generator in generators])
|
|
if temperature > 0
|
|
else first[:, None].expand(-1, 2)
|
|
)
|
|
return torch.cat((first[:, None], tail), dim=1)
|
|
|
|
batched, independent = _small_talker(), _small_talker()
|
|
batched.depth_decoder.generate_frames.side_effect = sample_depth
|
|
independent.depth_decoder.generate_frames.side_effect = sample_depth
|
|
batch_infos, independent_infos = make_infos(), make_infos()
|
|
initial_rng = {rid: info["breeze_state"]["generator"].get_state() for rid, info in batch_infos.items()}
|
|
saved_frames = []
|
|
# The second step drops both EOS requests and moves the live CFG companion
|
|
# ahead of its parent, exercising shrinking and reordered batches.
|
|
for step_ids in (order, [live_companion, "live_cfg", "live_plain"]):
|
|
expected_audio, expected_tokens = {}, {}
|
|
for request_id in step_ids:
|
|
if request_id.endswith(CFG_UNCOND_SUFFIX):
|
|
continue
|
|
pair = [request_id]
|
|
if independent_infos[request_id]["breeze_prompt"]["guidance_scale"] != 1.0:
|
|
pair.append(request_id + CFG_UNCOND_SUFFIX)
|
|
result = independent.make_omni_output(
|
|
torch.tensor([hidden_by_id[rid] for rid in pair]),
|
|
model_intermediate_buffer=[independent_infos[rid] for rid in pair],
|
|
request_token_spans=[(i, i + 1) for i in range(len(pair))],
|
|
)
|
|
tokens = independent.compute_logits(torch.zeros(len(pair), 2)).argmax(-1)
|
|
for index, rid in enumerate(pair):
|
|
expected_audio[rid] = result.multimodal_outputs["codes"]["audio"][index].clone()
|
|
expected_tokens[rid] = tokens[index]
|
|
|
|
result = batched.make_omni_output(
|
|
torch.tensor([hidden_by_id[rid] for rid in step_ids]),
|
|
model_intermediate_buffer=[batch_infos[rid] for rid in step_ids],
|
|
request_token_spans=[(i, i + 1) for i in range(len(step_ids))],
|
|
)
|
|
actual_tokens = batched.compute_logits(torch.zeros(len(step_ids), 2)).argmax(-1)
|
|
for index, rid in enumerate(step_ids):
|
|
actual = result.multimodal_outputs["codes"]["audio"][index]
|
|
torch.testing.assert_close(actual, expected_audio[rid], atol=0, rtol=0)
|
|
torch.testing.assert_close(actual_tokens[index], expected_tokens[rid], atol=0, rtol=0)
|
|
saved_frames.append((actual, actual.clone()))
|
|
for rid in order:
|
|
actual, expected = batch_infos[rid]["breeze_state"], independent_infos[rid]["breeze_state"]
|
|
for key in ("history", "current"):
|
|
torch.testing.assert_close(actual[key], expected[key], atol=0, rtol=0)
|
|
torch.testing.assert_close(actual["generator"].get_state(), expected["generator"].get_state())
|
|
|
|
# Unconditional companions never own sampling; an EOS request consumes
|
|
# only codebook-0 sampling, with no downstream depth draws.
|
|
for rid in (finished_companion, live_companion):
|
|
torch.testing.assert_close(batch_infos[rid]["breeze_state"]["generator"].get_state(), initial_rng[rid])
|
|
for rid in ("finished_cfg", "finished_plain"):
|
|
reference_generator = torch.Generator().set_state(initial_rng[rid])
|
|
sampling = batch_infos[rid]["breeze_sampling"]
|
|
sample_logits(
|
|
torch.zeros(1, batched.config.vocab_size),
|
|
sampling["temperature"],
|
|
sampling["top_k"],
|
|
sampling["top_p"],
|
|
reference_generator,
|
|
)
|
|
torch.testing.assert_close(
|
|
batch_infos[rid]["breeze_state"]["generator"].get_state(), reference_generator.get_state()
|
|
)
|
|
|
|
for rid in ("finished_cfg", finished_companion, "finished_plain", live_companion):
|
|
assert batch_infos[rid]["breeze_state"]["history"].tolist() == [[3, 3]]
|
|
assert batch_infos["live_cfg"]["breeze_state"]["history"].tolist() == [[3, 3, 1, 2]]
|
|
assert batch_infos["live_plain"]["breeze_state"]["history"].tolist() == [[3, 3, 3, 3]]
|
|
assert batch_infos["live_cfg"]["breeze_state"]["current"] is batch_infos[live_companion]["breeze_state"]["current"]
|
|
for original, snapshot in saved_frames:
|
|
torch.testing.assert_close(original, snapshot, atol=0, rtol=0)
|
|
|
|
|
|
@torch.inference_mode()
|
|
def test_talker_all_eos_preserves_history_and_skips_depth():
|
|
model = _small_talker()
|
|
infos = [_request_info("first"), _request_info("second")]
|
|
for info in infos:
|
|
info["breeze_state"]["history"] = torch.tensor([[1, 2]])
|
|
output = model.make_omni_output(
|
|
torch.tensor([[-1.0, -1.0], [-2.0, -2.0]]),
|
|
model_intermediate_buffer=infos,
|
|
request_token_spans=[(0, 1), (1, 2)],
|
|
)
|
|
|
|
model.depth_decoder.generate_frames.assert_not_called()
|
|
assert model.compute_logits(torch.zeros(2, 2)).argmax(-1).tolist() == [7, 7]
|
|
assert all(frame.numel() == 0 for frame in output.multimodal_outputs["codes"]["audio"])
|
|
for info in infos:
|
|
assert info["breeze_state"]["history"].tolist() == [[1, 2]]
|
|
assert info["breeze_state"]["current"].tolist() == [[0, 0, 0]]
|