anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Perf][Qwen3-TTS]Cached Incremental Decode (#5202)

Browse Source 
Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
This commit is contained in:
boatmanandYueqian Lin authored and GitHub committed 2026-08-11 05:24:04 +00:00
1 parent 12a5f6fbf3
commit 4f9b65074d
11 files changed
+3528 -735
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py
+254 -236
View File
@@ -24,7 +24,7 @@ class _FakeDecoder(nn.Module):
super().__init__()
self.total_upsample = total_upsample
self.decode_calls: list[dict[str, int]] = []
self.batched_decode_calls: list[dict[str, int]] = []
self.batched_decode_calls: list[dict[str, object]] = []
self.decode_codes: list[torch.Tensor] = []
self.cudagraph_calls: list[dict[str, int | torch.device]] = []
@@ -35,6 +35,7 @@ class _FakeDecoder(nn.Module):
self,
codes: torch.Tensor,
*,
caches=None,
chunk_size: int = 300,
left_context_size: int = 25,
) -> torch.Tensor:
@@ -58,6 +59,7 @@ class _FakeDecoder(nn.Module):
codes: torch.Tensor,
lengths: list[int],
*,
caches=None,
chunk_size: int = 300,
left_context_size: int = 25,
max_batch_size: int = 0,
@@ -69,14 +71,28 @@ class _FakeDecoder(nn.Module):
"max_batch_size": max_batch_size,
"codes_shape": tuple(codes.shape),
"lengths": tuple(lengths),
"caches": caches,
}
)
if caches is not None:
outputs = []
for row, (length, cache) in enumerate(zip(lengths, caches, strict=True)):
output = self.chunked_decode(
codes[row : row + 1, :, :length],
caches=cache,
chunk_size=chunk_size,
left_context_size=left_context_size,
)
outputs.append(output[0])
return outputs
batch = codes.shape[0]
frames = codes.shape[-1]
wav_len = frames * self.total_upsample + 6
wav = torch.arange(wav_len, dtype=torch.float32).view(1, 1, -1)
offsets = torch.arange(batch, dtype=torch.float32).view(batch, 1, 1) * 1000
return wav.expand(batch, 1, wav_len) + offsets
outputs = wav.expand(batch, 1, wav_len) + offsets
return [outputs[row, :, : lengths[row] * self.total_upsample] for row in range(batch)]
def enable_cudagraph(self, **kwargs):
self.cudagraph_calls.append(kwargs)
@@ -157,7 +173,7 @@ def _load_weights_noop(model: Qwen3TTSCode2Wav) -> set[str]:
return model.load_weights(iter(()))
def test_forward_trims_context_on_exact_frame_boundaries():
def test_forward_uses_decoder_audio_contract_on_exact_frame_boundaries():
model = _make_model()
out = model.forward(
@@ -170,7 +186,7 @@ def test_forward_trims_context_on_exact_frame_boundaries():
torch.testing.assert_close(audio, expected)
def test_forward_trims_trailing_padding_without_context():
def test_forward_uses_decoder_audio_contract_without_context():
model = _make_model()
out = model.forward(
@@ -183,24 +199,40 @@ def test_forward_trims_trailing_padding_without_context():
torch.testing.assert_close(audio, expected)
def test_forward_reuses_cached_ref_context_for_followup_chunk():
@pytest.mark.parametrize("skipped_length", [0, 1], ids=["empty", "malformed"])
def test_stateless_batch_uses_original_request_index_for_left_context(skipped_length):
model = _make_model()
first_codes = torch.tensor(
[
9,
8,
1,
2,
9,
8,
3,
4,
skipped_request = torch.zeros(skipped_length, dtype=torch.long)
valid_request = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long)
out = model.forward(
input_ids=torch.cat([skipped_request, valid_request]),
seq_token_counts=[skipped_length, valid_request.numel()],
runtime_additional_information=[
{
"codes": {"audio": torch.zeros(skipped_length, dtype=torch.float32)},
"meta": {"left_context_size": 0},
},
{
"codes": {"audio": torch.zeros(8, dtype=torch.float32)},
"meta": {"left_context_size": 2},
},
],
dtype=torch.long,
)
model.forward(
input_ids=first_codes,
wav_output = out.multimodal_outputs["model_outputs"]
torch.testing.assert_close(wav_output[0], torch.empty(0))
torch.testing.assert_close(wav_output[1], torch.arange(8, 16, dtype=torch.float32))
def test_stateless_forward_trims_reference_prefix_from_decoder_output():
model = _make_model()
def _incremental_decode(codes, lengths, *, caches, **_kwargs):
return torch.arange(16, dtype=torch.float32).view(1, 1, -1)
model.decoder.batched_chunked_decode = _incremental_decode
out = model.forward(
input_ids=torch.arange(8, dtype=torch.long),
runtime_additional_information=[
{
"meta": {
@@ -213,104 +245,48 @@ def test_forward_reuses_cached_ref_context_for_followup_chunk():
],
)
followup_codes = torch.tensor([5, 6, 7, 15, 16, 17], dtype=torch.long)
audio = out.multimodal_outputs["model_outputs"][0]
torch.testing.assert_close(audio, torch.arange(8, 16, dtype=torch.float32))
def test_request_aware_forward_passes_only_delta_frames_to_decoder():
model = _make_model(async_chunk=True)
model.forward(
input_ids=followup_codes,
runtime_additional_information=[
{
"meta": {
"left_context_size": 3,
"ref_context_size": 2,
"ref_context_request_id": "rid",
"ref_context_included": False,
"finished": torch.tensor(True),
}
input_ids=torch.arange(6, dtype=torch.long),
runtime_additional_information=[{"meta": {"request_id": "rid", "left_context_size": 0}}],
)
model.forward(
input_ids=torch.arange(4, dtype=torch.long),
runtime_additional_information=[{"meta": {"request_id": "rid", "left_context_size": 0}}],
)
assert [tuple(codes.shape) for codes in model.decoder.decode_codes[-2:]] == [
(1, _NUM_QUANTIZERS, 3),
(1, _NUM_QUANTIZERS, 2),
]
def test_four_frame_first_chunk_uses_configured_initial_size_without_ramp():
model = _make_model(
async_chunk=True,
stage_connector_config={
"extra": {
"codec_chunk_frames": 25,
"initial_codec_chunk_frames": 4,
}
],
},
)
_load_weights_noop(model)
model.forward(
input_ids=torch.arange(8, dtype=torch.long),
runtime_additional_information=[{"meta": {"request_id": "four-frame-request", "left_context_size": 0}}],
)
torch.testing.assert_close(
model.decoder.decode_codes[-1],
torch.tensor(
[
[
[9, 8, 5, 6, 7],
[9, 8, 15, 16, 17],
]
],
dtype=torch.long,
),
)
assert "rid" not in model._ref_context_cache
def test_forward_fails_fast_when_ref_context_cache_is_missing():
model = _make_model()
followup_codes = torch.tensor([5, 6, 7, 15, 16, 17], dtype=torch.long)
with pytest.raises(ValueError, match="Missing Qwen3-TTS ref context cache"):
model.forward(
input_ids=followup_codes,
runtime_additional_information=[
{
"meta": {
"left_context_size": 3,
"ref_context_size": 2,
"ref_context_request_id": "rid",
"ref_context_included": False,
}
}
],
)
def test_ref_context_cache_evicts_lru_entries_when_requests_abort():
model = _make_model()
model._ref_context_cache_max_entries = 2
first_codes = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long)
for rid in ("a", "b", "c"):
model.forward(
input_ids=first_codes,
runtime_additional_information=[
{
"meta": {
"left_context_size": 2,
"ref_context_size": 2,
"ref_context_request_id": rid,
"ref_context_included": True,
}
}
],
)
assert list(model._ref_context_cache) == ["b", "c"]
assert model._ref_context_cache_bytes == sum(model._tensor_nbytes(t) for t in model._ref_context_cache.values())
def test_ref_context_cache_evicts_to_byte_cap():
model = _make_model()
model._ref_context_cache_max_entries = 100
model._ref_context_cache_max_bytes = 33
first_codes = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long)
for rid in ("a", "b"):
model.forward(
input_ids=first_codes,
runtime_additional_information=[
{
"meta": {
"left_context_size": 2,
"ref_context_size": 2,
"ref_context_request_id": rid,
"ref_context_included": True,
}
}
],
)
assert list(model._ref_context_cache) == ["b"]
assert model._ref_context_cache_bytes == 32
assert model.decoder._initial_codec_chunk_frames == 4
assert model.decoder._incremental_chunk_ramp == []
assert tuple(model.decoder.decode_codes[-1].shape) == (1, _NUM_QUANTIZERS, 4)
def test_connector_codec_chunking_does_not_override_decode_chunking():
@@ -332,7 +308,7 @@ def test_connector_codec_chunking_does_not_override_decode_chunking():
model.forward(
input_ids=torch.arange(12, dtype=torch.long),
runtime_additional_information=[{"meta": {"left_context_size": 0}}],
runtime_additional_information=[{"meta": {"request_id": "rid", "left_context_size": 0}}],
)
assert model.decoder.decode_calls[-1] == {
@@ -391,6 +367,66 @@ def test_forward_emits_zero_samples_for_empty_codec_payload():
assert out.multimodal_outputs["model_outputs"][0].numel() == 0
def test_dummy_runtime_information_provides_finished_request_state():
model = _make_model(async_chunk=True)
runtime_info = model.get_dummy_runtime_additional_information(2)
assert len(runtime_info) == 2
assert [info["meta"]["request_id"] for info in runtime_info] == [
"__qwen3_tts_dummy_run__0",
"__qwen3_tts_dummy_run__1",
]
assert all(info["meta"]["finished"] for info in runtime_info)
def test_dummy_request_state_is_marked_for_graph_stats_suppression():
model = _make_model(async_chunk=True)
runtime_info = model.get_dummy_runtime_additional_information(1)
model.forward(
input_ids=torch.arange(4, dtype=torch.long),
runtime_additional_information=runtime_info,
)
cache = model.decoder.batched_decode_calls[0]["caches"][0]
assert cache["_is_dummy_run"] is True
assert model._decoder_state_cache == {}
def test_decoder_state_cache_capacity_warns_without_evicting_or_failing():
model = _make_model(async_chunk=True)
assert model._decoder_state_cache_warn_entries == 512
model._decoder_state_cache.update({f"active-{index}": {"prefix_frames": 0} for index in range(512)})
with patch("vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav.logger.warning_once") as warning:
model.forward(
input_ids=torch.arange(4, dtype=torch.long),
runtime_additional_information=[{"meta": {"request_id": "new-request"}}],
)
warning.assert_called_once()
assert len(model._decoder_state_cache) == 513
assert "new-request" in model._decoder_state_cache
def test_decoder_state_cache_uses_scheduler_request_id_for_finished_cleanup():
model = _make_model(async_chunk=True)
model.forward(
input_ids=torch.arange(4, dtype=torch.long),
request_ids=["internal-request-id"],
runtime_additional_information=[{"meta": {"request_id": "external-request-id"}}],
)
assert "internal-request-id" in model._decoder_state_cache
assert "external-request-id" not in model._decoder_state_cache
model.on_requests_finished(["internal-request-id"])
assert model._decoder_state_cache == {}
def test_forward_batches_equal_length_requests_in_one_decoder_call():
model = _make_model()
@@ -403,11 +439,15 @@ def test_forward_batches_equal_length_requests_in_one_decoder_call():
],
)
assert model.decoder.decode_calls == [
assert model.decoder.decode_calls == []
assert model.decoder.batched_decode_calls == [
{
"chunk_size": 300,
"left_context_size": 25,
"max_batch_size": 0,
"codes_shape": (2, _NUM_QUANTIZERS, 3),
"lengths": (3, 3),
"caches": None,
}
]
audios = out.multimodal_outputs["model_outputs"]
@@ -415,10 +455,81 @@ def test_forward_batches_equal_length_requests_in_one_decoder_call():
torch.testing.assert_close(audios[1], torch.arange(1004, 1012, dtype=torch.float32))
def test_forward_uses_variable_length_chunk_batching_for_long_bucket_group():
def test_forward_batches_request_aware_decoder_states():
model = _make_model(async_chunk=True)
calls = []
def _batched_chunked_decode(codes, lengths, caches, **kwargs):
calls.append((codes, lengths, caches))
return [torch.arange(16, dtype=torch.float32).view(1, -1) + row * 100 for row in range(2)]
model.decoder.batched_chunked_decode = _batched_chunked_decode
per_request = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long)
out = model.forward(
input_ids=torch.cat([per_request, per_request]),
seq_token_counts=[8, 8],
runtime_additional_information=[
{
"meta": {
"left_context_size": 2,
"ref_context_size": 2,
"ref_context_request_id": request_id,
"request_id": request_id,
"ref_context_included": True,
}
}
for request_id in ("a", "b")
],
)
assert len(calls) == 1
assert tuple(calls[0][0].shape) == (2, 2, 4)
assert calls[0][1] == [4, 4]
assert [cache["prefix_frames"] for cache in calls[0][2]] == [2, 2]
audios = out.multimodal_outputs["model_outputs"]
torch.testing.assert_close(audios[0], torch.arange(16, dtype=torch.float32))
torch.testing.assert_close(audios[1], torch.arange(100, 116, dtype=torch.float32))
def test_forward_allows_followup_delta_without_reference_prefix():
model = _make_model(async_chunk=True)
first_codes = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long)
model.forward(
input_ids=first_codes,
runtime_additional_information=[
{
"meta": {
"request_id": "rid",
"ref_context_request_id": "rid",
"ref_context_size": 2,
"ref_context_included": True,
}
}
],
)
model.forward(
input_ids=torch.tensor([5, 6, 7, 15, 16, 17], dtype=torch.long),
runtime_additional_information=[
{
"meta": {
"request_id": "rid",
"ref_context_request_id": "rid",
"ref_context_size": 2,
"ref_context_included": False,
}
}
],
)
assert [tuple(codes.shape) for codes in model.decoder.decode_codes[-2:]] == [
(1, _NUM_QUANTIZERS, 4),
(1, _NUM_QUANTIZERS, 3),
]
def test_forward_passes_variable_length_padded_batch_to_decoder():
model = _make_model()
model._decode_batch_bucket_frames = [4]
model._decode_variable_chunk_batch_min_frames = 1
out = model.forward(
input_ids=torch.arange(12, dtype=torch.long),
@@ -437,6 +548,7 @@ def test_forward_uses_variable_length_chunk_batching_for_long_bucket_group():
"max_batch_size": 0,
"codes_shape": (2, _NUM_QUANTIZERS, 4),
"lengths": (2, 4),
"caches": None,
}
]
audios = out.multimodal_outputs["model_outputs"]
@@ -444,108 +556,6 @@ def test_forward_uses_variable_length_chunk_batching_for_long_bucket_group():
torch.testing.assert_close(audios[1], torch.arange(1008, 1016, dtype=torch.float32))
def test_forward_bucket_batches_different_length_requests_and_trims_rows():
model = _make_model()
model._decode_batch_bucket_frames = [4]
out = model.forward(
input_ids=torch.arange(18, dtype=torch.long),
seq_token_counts=[4, 6, 8],
runtime_additional_information=[
{"meta": {"left_context_size": 0}},
{"meta": {"left_context_size": 1}},
{"meta": {"left_context_size": 2}},
],
)
assert model.decoder.decode_calls == [
{
"chunk_size": 300,
"left_context_size": 25,
"codes_shape": (3, _NUM_QUANTIZERS, 4),
}
]
assert model.decoder.batched_decode_calls == []
audios = out.multimodal_outputs["model_outputs"]
torch.testing.assert_close(audios[0], torch.arange(8, dtype=torch.float32))
torch.testing.assert_close(audios[1], torch.arange(1004, 1012, dtype=torch.float32))
torch.testing.assert_close(audios[2], torch.arange(2008, 2016, dtype=torch.float32))
def test_forward_bucket_pads_only_to_group_max_frame_length():
model = _make_model()
model._decode_batch_bucket_frames = [8]
out = model.forward(
input_ids=torch.arange(10, dtype=torch.long),
seq_token_counts=[4, 6],
runtime_additional_information=[
{"meta": {"left_context_size": 0}},
{"meta": {"left_context_size": 0}},
],
)
assert model.decoder.decode_calls == [
{
"chunk_size": 300,
"left_context_size": 25,
"codes_shape": (2, _NUM_QUANTIZERS, 3),
}
]
assert model.decoder.batched_decode_calls == []
audios = out.multimodal_outputs["model_outputs"]
torch.testing.assert_close(audios[0], torch.arange(8, dtype=torch.float32))
torch.testing.assert_close(audios[1], torch.arange(1000, 1012, dtype=torch.float32))
def test_forward_splits_bucket_groups_by_configured_max_batch_size():
model = _make_model()
model._decode_batch_bucket_frames = [4]
model._decode_batch_max_size = 2
model.forward(
input_ids=torch.arange(18, dtype=torch.long),
seq_token_counts=[4, 6, 8],
runtime_additional_information=[
{"meta": {"left_context_size": 0}},
{"meta": {"left_context_size": 0}},
{"meta": {"left_context_size": 0}},
],
)
assert model.decoder.decode_calls == [
{
"chunk_size": 300,
"left_context_size": 25,
"codes_shape": (2, _NUM_QUANTIZERS, 3),
},
{
"chunk_size": 300,
"left_context_size": 25,
"codes_shape": (1, _NUM_QUANTIZERS, 4),
},
]
assert model.decoder.batched_decode_calls == []
def test_forward_does_not_pad_singleton_bucket_group():
model = _make_model()
model._decode_batch_bucket_frames = [4]
model.forward(
input_ids=torch.arange(4, dtype=torch.long),
runtime_additional_information=[{"meta": {"left_context_size": 0}}],
)
assert model.decoder.decode_calls == [
{
"chunk_size": 300,
"left_context_size": 25,
"codes_shape": (1, _NUM_QUANTIZERS, 2),
}
]
def test_decode_chunking_override_is_passed_to_cudagraph():
model = _make_model(
async_chunk=True,
@@ -563,27 +573,26 @@ def test_decode_chunking_override_is_passed_to_cudagraph():
_load_weights_noop(model)
assert model.decoder.cudagraph_calls[-1] == {
"capture_sizes": None,
"capture_batch_sizes": None,
"extra_capture_shapes": None,
"compile_shapes": None,
"stateless_capture_sizes": None,
"device": torch.device("cuda"),
"codec_chunk_frames": 25,
"codec_left_context_frames": 72,
"initial_codec_chunk_frames": 1,
"codec_chunk_ramp": None,
"async_chunk": True,
"decode_chunk_size": 400,
"decode_left_context": 17,
}
def test_cudagraph_capture_shapes_can_be_configured():
def test_cudagraph_batch_config_is_preserved():
model = _make_model(
async_chunk=True,
device=torch.device("cuda"),
stage_connector_config={
"extra": {
"decode_cudagraph_capture_sizes": "97,325",
"decode_cudagraph_batch_sizes": [1, 2, 4, 8],
"decode_cudagraph_extra_capture_shapes": ["3:325", [5, 325]],
}
},
)
@@ -591,26 +600,39 @@ def test_cudagraph_capture_shapes_can_be_configured():
_load_weights_noop(model)
call = model.decoder.cudagraph_calls[-1]
assert call["capture_sizes"] == [97, 325]
assert call["capture_batch_sizes"] == [1, 2, 4, 8]
assert call["extra_capture_shapes"] == [(3, 325), (5, 325)]
def test_decode_compile_shapes_can_be_configured():
def test_stateless_cudagraph_capture_sizes_are_passed_from_config():
model = _make_model(
async_chunk=False,
device=torch.device("cuda"),
stage_connector_config={
"extra": {
"decode_cudagraph_capture_sizes": [97, 400],
}
},
)
_load_weights_noop(model)
assert model.decoder.cudagraph_calls[-1]["stateless_capture_sizes"] == [97, 400]
def test_async_cudagraph_ignores_stateless_capture_sizes():
model = _make_model(
async_chunk=True,
device=torch.device("cuda"),
stage_connector_config={
"extra": {
"decode_compile_shapes": ["1:325", [1, 73]],
"decode_cudagraph_capture_sizes": [97, 400],
}
},
)
_load_weights_noop(model)
call = model.decoder.cudagraph_calls[-1]
assert call["compile_shapes"] == [(1, 73), (1, 325)]
assert model.decoder.cudagraph_calls[-1]["stateless_capture_sizes"] is None
def test_decode_tf32_can_be_configured():
@@ -642,23 +664,19 @@ def test_decode_tf32_can_be_configured():
torch.set_float32_matmul_precision(old_matmul_precision)
def test_decode_batch_bucket_frames_can_be_configured():
def test_decode_batch_max_size_can_be_configured():
model = _make_model(
async_chunk=True,
stage_connector_config={
"extra": {
"decode_batch_bucket_frames": "73,169",
"decode_batch_max_size": 10,
"decode_variable_chunk_batch_min_frames": 512,
}
},
)
_load_weights_noop(model)
assert model._decode_batch_bucket_frames == [73, 169]
assert model._decode_batch_max_size == 10
assert model._decode_variable_chunk_batch_min_frames == 512
def test_invalid_decode_batch_max_size_is_rejected():
tests/model_executor/models/qwen3_tts/test_qwen3_tts_incremental_decode.py
+1253
View File
File diff suppressed because it is too large. Load diff
tests/model_executor/models/qwen3_tts/test_qwen3_tts_tokenizer_mask_cache.py
-125
View File
@@ -1,125 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import inspect
import pytest
import torch
from transformers.masking_utils import create_sliding_window_causal_mask
from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import (
Qwen3TTSTokenizerV2DecoderConfig,
)
from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
Qwen3TTSTokenizerV2DecoderTransformerModel,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _make_config(attention_implementation: str) -> Qwen3TTSTokenizerV2DecoderConfig:
config = Qwen3TTSTokenizerV2DecoderConfig(
hidden_size=32,
latent_dim=16,
max_position_embeddings=512,
num_attention_heads=4,
num_key_value_heads=4,
intermediate_size=64,
num_hidden_layers=1,
num_quantizers=2,
decoder_dim=32,
upsample_rates=(2,),
upsampling_ratios=(2,),
)
config._attn_implementation = attention_implementation
return config
def _make_model(
attention_implementation: str,
dtype: torch.dtype = torch.float32,
) -> Qwen3TTSTokenizerV2DecoderTransformerModel:
model = Qwen3TTSTokenizerV2DecoderTransformerModel._from_config(_make_config(attention_implementation))
return model.to(dtype=dtype).eval()
def _create_reference_mask(
model: Qwen3TTSTokenizerV2DecoderTransformerModel,
inputs_embeds: torch.Tensor,
) -> torch.Tensor | None:
sequence_length = inputs_embeds.shape[1]
mask_kwargs = {
"config": model.config,
"attention_mask": None,
"past_key_values": None,
"position_ids": None,
}
sig = inspect.signature(create_sliding_window_causal_mask)
if "input_embeds" in sig.parameters:
mask_kwargs["input_embeds"] = inputs_embeds[:1]
mask_kwargs["cache_position"] = torch.arange(sequence_length)
else:
mask_kwargs["inputs_embeds"] = inputs_embeds[:1]
return create_sliding_window_causal_mask(**mask_kwargs)
@pytest.mark.parametrize("sequence_length", [17, 72, 73, 169, 325])
@pytest.mark.parametrize(
("attention_implementation", "dtype"),
[("sdpa", torch.float32), ("eager", torch.float32), ("eager", torch.float16)],
)
def test_cached_sliding_mask_matches_transformers(
sequence_length: int,
attention_implementation: str,
dtype: torch.dtype,
):
model = _make_model(attention_implementation, dtype)
inputs_embeds = torch.empty(2, sequence_length, model.config.hidden_size, dtype=dtype)
actual = model._get_sliding_attention_mask(inputs_embeds)
expected = _create_reference_mask(model, inputs_embeds)
if expected is None:
assert actual is None
assert model._get_sliding_attention_mask(inputs_embeds) is None
return
assert actual is not None
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
assert actual.shape == (1, 1, sequence_length, sequence_length)
assert actual.is_contiguous()
assert model._get_sliding_attention_mask(inputs_embeds) is actual
def test_cached_eager_mask_tracks_compute_dtype():
model = _make_model("eager")
fp32_inputs = torch.empty(1, 73, model.config.hidden_size)
fp32_mask = model._get_sliding_attention_mask(fp32_inputs)
assert fp32_mask is not None
assert fp32_mask.dtype == torch.float32
model.half()
fp16_inputs = fp32_inputs.half()
fp16_mask = model._get_sliding_attention_mask(fp16_inputs)
assert fp16_mask is not None
assert fp16_mask.dtype == torch.float16
assert torch.isfinite(fp16_mask).all()
assert fp16_mask.min() == torch.finfo(torch.float16).min
@pytest.mark.parametrize(
("kwargs", "message"),
[
({"use_cache": True}, "use_cache=True"),
({"past_key_values": object()}, "past_key_values"),
({"cache_position": torch.arange(1, 18)}, "zero-based cache_position"),
({"position_ids": torch.arange(1, 18).unsqueeze(0)}, "zero-based position_ids"),
],
)
def test_cached_mask_rejects_unsupported_cache_and_positions(kwargs: dict, message: str):
model = _make_model("sdpa")
inputs_embeds = torch.randn(1, 17, model.config.latent_dim)
with pytest.raises(ValueError, match=message):
model(inputs_embeds=inputs_embeds, **kwargs)