229 lines
7.5 KiB
Python
229 lines
7.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.data import DiffusionOutput
|
|
from vllm_omni.diffusion.models.omnivoice.pipeline_omnivoice import (
|
|
OmniVoicePipeline,
|
|
_PreparedOmniVoiceRequest,
|
|
)
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
|
|
_attention_metadata_from_cu_seqs,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def test_sdpa_fallback_mask_preserves_packed_sequence_boundaries() -> None:
|
|
cu_seqs = torch.tensor([0, 2, 3, 6, 6], dtype=torch.int32)
|
|
|
|
metadata = _attention_metadata_from_cu_seqs(cu_seqs, 6, needs_sdpa_mask=True)
|
|
|
|
assert metadata.attn_mask is not None
|
|
expected = torch.tensor(
|
|
[
|
|
[1, 1, 0, 0, 0, 0],
|
|
[1, 1, 0, 0, 0, 0],
|
|
[0, 0, 1, 0, 0, 0],
|
|
[0, 0, 0, 1, 1, 1],
|
|
[0, 0, 0, 1, 1, 1],
|
|
[0, 0, 0, 1, 1, 1],
|
|
],
|
|
dtype=torch.bool,
|
|
)
|
|
torch.testing.assert_close(metadata.attn_mask[0, 0], expected)
|
|
|
|
|
|
def test_eager_attention_metadata_honors_exact_max_seqlen() -> None:
|
|
cu_seqs = torch.tensor([0, 2, 3, 6, 6], dtype=torch.int32)
|
|
|
|
metadata = _attention_metadata_from_cu_seqs(
|
|
cu_seqs,
|
|
6,
|
|
needs_sdpa_mask=False,
|
|
max_seqlen=3,
|
|
)
|
|
|
|
assert metadata.extra["max_seqlen_q"] == 3
|
|
assert metadata.extra["max_seqlen_k"] == 3
|
|
|
|
|
|
def test_request_batch_prepare_error_does_not_skip_valid_request() -> None:
|
|
"""A malformed request must not prevent another request from completing."""
|
|
prepared = _PreparedOmniVoiceRequest(
|
|
input_ids=torch.zeros(5, 8, dtype=torch.long),
|
|
audio_mask=torch.ones(5, dtype=torch.bool),
|
|
cond_len=3,
|
|
target_len=2,
|
|
seed=42,
|
|
)
|
|
|
|
def prepare_request(prompt, extra):
|
|
del extra
|
|
if prompt == "bad":
|
|
return DiffusionOutput(error="invalid prompt")
|
|
return prepared
|
|
|
|
def collate_requests(requests):
|
|
assert requests == [prepared]
|
|
return prepared.input_ids, prepared.audio_mask, [prepared.cond_len]
|
|
|
|
generator_calls = []
|
|
|
|
def generate_tokens(**kwargs):
|
|
generator_calls.append(kwargs)
|
|
assert kwargs["target_lens"] == [prepared.target_len]
|
|
return torch.ones(1, 8, prepared.target_len, dtype=torch.long)
|
|
|
|
pipeline = SimpleNamespace(
|
|
_prepare_request_input=prepare_request,
|
|
_collate_request_inputs=collate_requests,
|
|
generator=generate_tokens,
|
|
decoder=lambda tokens: tokens.float(),
|
|
num_step=32,
|
|
guidance_scale=7.0,
|
|
t_shift=1.0,
|
|
layer_penalty_factor=0.0,
|
|
position_temperature=0.0,
|
|
class_temperature=0.0,
|
|
)
|
|
sampling = OmniDiffusionSamplingParams(num_inference_steps=32, guidance_scale=7.0)
|
|
batch = DiffusionRequestBatch(
|
|
requests=[
|
|
OmniDiffusionRequest(prompt="bad", sampling_params=sampling, request_id="bad"),
|
|
OmniDiffusionRequest(
|
|
prompt="good",
|
|
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=32, guidance_scale=7.0),
|
|
request_id="good",
|
|
),
|
|
]
|
|
)
|
|
|
|
outputs = OmniVoicePipeline.forward(pipeline, batch)
|
|
|
|
assert len(outputs) == 2
|
|
assert outputs[0].error == "invalid prompt"
|
|
assert outputs[1].error is None
|
|
torch.testing.assert_close(outputs[1].output, torch.ones(1, 8, 2))
|
|
assert generator_calls[0]["guidance_scale"] == 7.0
|
|
|
|
|
|
def test_request_batch_honors_explicit_sampling_overrides() -> None:
|
|
"""Request steps and guidance must override the OmniVoice defaults."""
|
|
prepared = _PreparedOmniVoiceRequest(
|
|
input_ids=torch.zeros(5, 8, dtype=torch.long),
|
|
audio_mask=torch.ones(5, dtype=torch.bool),
|
|
cond_len=3,
|
|
target_len=2,
|
|
seed=42,
|
|
)
|
|
captured_sampling = []
|
|
|
|
def generate_tokens(**kwargs):
|
|
captured_sampling.append((kwargs["num_step"], kwargs["guidance_scale"]))
|
|
return torch.ones(1, 8, prepared.target_len, dtype=torch.long)
|
|
|
|
pipeline = SimpleNamespace(
|
|
_prepare_request_input=lambda prompt, extra: prepared,
|
|
_collate_request_inputs=lambda requests: (prepared.input_ids, prepared.audio_mask, [prepared.cond_len]),
|
|
generator=generate_tokens,
|
|
decoder=lambda tokens: tokens.float(),
|
|
num_step=32,
|
|
guidance_scale=2.0,
|
|
t_shift=1.0,
|
|
layer_penalty_factor=0.0,
|
|
position_temperature=0.0,
|
|
class_temperature=0.0,
|
|
)
|
|
batch = DiffusionRequestBatch(
|
|
requests=[
|
|
OmniDiffusionRequest(
|
|
prompt="hello",
|
|
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=12, guidance_scale=6.5),
|
|
request_id="request",
|
|
)
|
|
]
|
|
)
|
|
|
|
outputs = OmniVoicePipeline.forward(pipeline, batch)
|
|
|
|
assert outputs[0].error is None
|
|
assert captured_sampling == [(12, 6.5)]
|
|
|
|
|
|
def test_request_batch_prepare_error_preserves_surrounding_output_indices() -> None:
|
|
"""A middle prepare error must not shift valid outputs into the wrong slots."""
|
|
prepared_before = _PreparedOmniVoiceRequest(
|
|
input_ids=torch.zeros(3, 8, dtype=torch.long),
|
|
audio_mask=torch.ones(3, dtype=torch.bool),
|
|
cond_len=2,
|
|
target_len=1,
|
|
seed=1,
|
|
)
|
|
prepared_after = _PreparedOmniVoiceRequest(
|
|
input_ids=torch.zeros(5, 8, dtype=torch.long),
|
|
audio_mask=torch.ones(5, dtype=torch.bool),
|
|
cond_len=3,
|
|
target_len=2,
|
|
seed=2,
|
|
)
|
|
|
|
def prepare_request(prompt, extra):
|
|
del extra
|
|
if prompt == "before":
|
|
return prepared_before
|
|
if prompt == "after":
|
|
return prepared_after
|
|
return DiffusionOutput(error="invalid middle request")
|
|
|
|
def collate_requests(requests):
|
|
assert requests == [prepared_before, prepared_after]
|
|
return (
|
|
torch.cat([prepared_before.input_ids, prepared_after.input_ids]),
|
|
torch.cat([prepared_before.audio_mask, prepared_after.audio_mask]),
|
|
[prepared_before.cond_len, prepared_after.cond_len],
|
|
)
|
|
|
|
def generate_tokens(**kwargs):
|
|
assert kwargs["target_lens"] == [1, 2]
|
|
before = torch.full((1, 8, 1), 11, dtype=torch.long)
|
|
after = torch.full((1, 8, 2), 22, dtype=torch.long)
|
|
return torch.cat([before, after], dim=-1)
|
|
|
|
pipeline = SimpleNamespace(
|
|
_prepare_request_input=prepare_request,
|
|
_collate_request_inputs=collate_requests,
|
|
generator=generate_tokens,
|
|
decoder=lambda tokens: tokens.float(),
|
|
num_step=32,
|
|
guidance_scale=2.0,
|
|
t_shift=1.0,
|
|
layer_penalty_factor=0.0,
|
|
position_temperature=0.0,
|
|
class_temperature=0.0,
|
|
)
|
|
batch = DiffusionRequestBatch(
|
|
requests=[
|
|
OmniDiffusionRequest(
|
|
prompt=prompt,
|
|
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=32),
|
|
request_id=prompt,
|
|
)
|
|
for prompt in ("before", "bad", "after")
|
|
]
|
|
)
|
|
|
|
outputs = OmniVoicePipeline.forward(pipeline, batch)
|
|
|
|
assert len(outputs) == 3
|
|
torch.testing.assert_close(outputs[0].output, torch.full((1, 8, 1), 11.0))
|
|
assert outputs[1].error == "invalid middle request"
|
|
torch.testing.assert_close(outputs[2].output, torch.full((1, 8, 2), 22.0))
|