757 lines
30 KiB
Python
757 lines
30 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Unit tests for 310P patch wiring.
|
|
|
|
The tests load patch modules from source with fake Qwen3-TTS dependencies, so
|
|
they validate the patch contract without loading real model or NPU kernels.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def _repo_root() -> Path:
|
|
marker = Path("vllm_omni") / "platforms" / "npu" / "_310p" / "patch"
|
|
for parent in Path(__file__).resolve().parents:
|
|
if (parent / marker).is_dir():
|
|
return parent
|
|
raise FileNotFoundError(f"could not locate repo root containing {marker}")
|
|
|
|
|
|
def _load_source_module(name: str, path: Path):
|
|
spec = importlib.util.spec_from_file_location(name, path)
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError(f"cannot load import spec for {path}")
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def _install_fake_module(monkeypatch: pytest.MonkeyPatch, name: str, **attrs):
|
|
module = types.ModuleType(name)
|
|
for key, value in attrs.items():
|
|
setattr(module, key, value)
|
|
monkeypatch.setitem(sys.modules, name, module)
|
|
return module
|
|
|
|
|
|
def _install_qwen3_tts_patch_fakes(monkeypatch: pytest.MonkeyPatch):
|
|
class FakeAudioResampler:
|
|
def __init__(self, *, target_sr: int):
|
|
self.target_sr = target_sr
|
|
|
|
def resample(self, wav, *, orig_sr: int):
|
|
del orig_sr
|
|
return wav
|
|
|
|
class FakeCodePredictorAttention(torch.nn.Module):
|
|
def __init__(self, *args, **kwargs):
|
|
del args, kwargs
|
|
super().__init__()
|
|
self.register_buffer("_fusion_causal_mask", torch.ones(1), persistent=False)
|
|
|
|
class FakeCodePredictorDecoderLayer(torch.nn.Module):
|
|
pass
|
|
|
|
class FakeCodePredictorBaseModel(torch.nn.Module):
|
|
pass
|
|
|
|
class FakeProjection(torch.nn.Linear):
|
|
def __init__(self):
|
|
super().__init__(4, 4, bias=False)
|
|
self.call_shapes: list[tuple[int, ...]] = []
|
|
|
|
def forward(self, hidden_states):
|
|
self.call_shapes.append(tuple(hidden_states.shape))
|
|
return super().forward(hidden_states)
|
|
|
|
class FakeCodePredictorWrapper(torch.nn.Module):
|
|
def __init__(self, *args, **kwargs):
|
|
del args, kwargs
|
|
super().__init__()
|
|
self.model = torch.nn.Module()
|
|
self.model.codec_embedding = torch.nn.ModuleList([torch.nn.Embedding(8, 4), torch.nn.Embedding(8, 4)])
|
|
self.model.linear = torch.nn.Linear(4, 4, bias=False)
|
|
self.lm_head = torch.nn.ModuleList([torch.nn.Linear(4, 8, bias=False), torch.nn.Linear(4, 8, bias=False)])
|
|
self.small_to_mtp_projection = FakeProjection()
|
|
with torch.no_grad():
|
|
self.small_to_mtp_projection.weight.copy_(torch.diag(torch.tensor([1.0, 2.0, 3.0, 4.0])))
|
|
self._wrapper_config = SimpleNamespace(use_parallel_embedding=False, sampling_mode="per_call")
|
|
self._projected_codec_embed_weight = None
|
|
|
|
def load_weights(self, weights):
|
|
del weights
|
|
return {"loaded"}
|
|
|
|
class FakeCode2WavBase(torch.nn.Module):
|
|
def __init__(self, *args, **kwargs):
|
|
del args, kwargs
|
|
super().__init__()
|
|
|
|
class FakeEncoder:
|
|
def __init__(self):
|
|
self.to_calls: list[dict] = []
|
|
self.last_input_dtype = None
|
|
|
|
def to(self, **kwargs):
|
|
self.to_calls.append(kwargs)
|
|
return self
|
|
|
|
def encode(self, *, input_values, return_dict: bool):
|
|
assert return_dict
|
|
self.last_input_dtype = input_values.dtype
|
|
return SimpleNamespace(audio_codes=torch.arange(12, dtype=torch.long).reshape(1, 3, 4))
|
|
|
|
class FakeFeatureBatch(dict):
|
|
def to(self, target):
|
|
if isinstance(target, torch.dtype):
|
|
for key, value in list(self.items()):
|
|
if torch.is_floating_point(value):
|
|
self[key] = value.to(dtype=target)
|
|
self.dtype = target
|
|
else:
|
|
self.device = torch.device(target)
|
|
for key, value in list(self.items()):
|
|
self[key] = value.to(device=self.device)
|
|
return self
|
|
|
|
class FakeFeatureExtractor:
|
|
sampling_rate = 24000
|
|
|
|
def __call__(self, *, raw_audio, sampling_rate: int, return_tensors: str):
|
|
assert len(raw_audio) == 1
|
|
assert sampling_rate == self.sampling_rate
|
|
assert return_tensors == "pt"
|
|
return FakeFeatureBatch(
|
|
input_values=torch.ones(1, 1, 8, dtype=torch.float32),
|
|
padding_mask=torch.ones(1, 1, 8, dtype=torch.long),
|
|
)
|
|
|
|
class FakeTalkerBase(torch.nn.Module):
|
|
def __init__(self, *, vllm_config, prefix: str = ""):
|
|
del vllm_config, prefix
|
|
super().__init__()
|
|
self._embedding_dtype = torch.bfloat16
|
|
self._prompt_builder = SimpleNamespace(_embedding_dtype=torch.bfloat16)
|
|
self.encoder = FakeEncoder()
|
|
self._encoder_feature_extractor = FakeFeatureExtractor()
|
|
self._encoder_valid_num_quantizers = 2
|
|
self._encoder_downsample_rate = 2
|
|
|
|
def load_weights(self, weights):
|
|
del weights
|
|
self.encoder.to(dtype=torch.bfloat16)
|
|
return {"loaded"}
|
|
|
|
class FakePromptEmbedsBuilder:
|
|
pass
|
|
|
|
fake_qwen3_code_predictor = _install_fake_module(
|
|
monkeypatch,
|
|
"vllm_omni.model_executor.models.common.qwen3_code_predictor",
|
|
CodePredictorAttention=FakeCodePredictorAttention,
|
|
CodePredictorDecoderLayer=FakeCodePredictorDecoderLayer,
|
|
CodePredictorBaseModel=FakeCodePredictorBaseModel,
|
|
CodePredictorWrapper=FakeCodePredictorWrapper,
|
|
_rotate_half=lambda x: x,
|
|
)
|
|
fake_qwen3_tts_code_predictor_vllm = _install_fake_module(
|
|
monkeypatch,
|
|
"vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm",
|
|
Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM=FakeCodePredictorWrapper,
|
|
Qwen3TTSTalkerCodePredictorModelVLLM=FakeCodePredictorBaseModel,
|
|
CodePredictorWrapper=FakeCodePredictorWrapper,
|
|
)
|
|
fake_qwen3_tts_code2wav = _install_fake_module(
|
|
monkeypatch,
|
|
"vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav",
|
|
Qwen3TTSCode2Wav=FakeCode2WavBase,
|
|
)
|
|
fake_prompt_builder = _install_fake_module(
|
|
monkeypatch,
|
|
"vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder",
|
|
Qwen3TTSPromptEmbedsBuilder=FakePromptEmbedsBuilder,
|
|
mel_spectrogram=lambda *_args, **_kwargs: torch.empty(0),
|
|
)
|
|
fake_talker = _install_fake_module(
|
|
monkeypatch,
|
|
"vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker",
|
|
Qwen3TTSTalkerForConditionalGeneration=FakeTalkerBase,
|
|
Qwen3TTSPromptEmbedsBuilder=FakePromptEmbedsBuilder,
|
|
)
|
|
|
|
_install_fake_module(monkeypatch, "vllm")
|
|
_install_fake_module(monkeypatch, "vllm.multimodal")
|
|
_install_fake_module(monkeypatch, "vllm.multimodal.audio", AudioResampler=FakeAudioResampler)
|
|
_install_fake_module(monkeypatch, "torch_npu", npu_format_cast=lambda weight, _fmt: weight)
|
|
_install_fake_module(monkeypatch, "vllm_ascend")
|
|
_install_fake_module(monkeypatch, "vllm_ascend._310p")
|
|
_install_fake_module(monkeypatch, "vllm_ascend._310p.attention")
|
|
_install_fake_module(
|
|
monkeypatch,
|
|
"vllm_ascend._310p.attention.attention_mask",
|
|
AttentionMaskBuilder310=SimpleNamespace(
|
|
gen_causal_additive_mask=lambda max_seq, device: torch.zeros(
|
|
max_seq,
|
|
max_seq,
|
|
device=device,
|
|
)
|
|
),
|
|
)
|
|
_install_fake_module(monkeypatch, "vllm_ascend.sample")
|
|
_install_fake_module(
|
|
monkeypatch,
|
|
"vllm_ascend.sample.sampler",
|
|
apply_top_k_top_p=lambda logits, **_kwargs: logits,
|
|
random_sample=lambda probs, _generators: probs.argmax(dim=-1, keepdim=True),
|
|
)
|
|
_install_fake_module(
|
|
monkeypatch,
|
|
"vllm_ascend.utils",
|
|
ACL_FORMAT_FRACTAL_NZ=29,
|
|
aligned_16=lambda tensor: tensor,
|
|
maybe_trans_nz=lambda weight: weight,
|
|
nd_to_nz_2d=lambda tensor: tensor,
|
|
)
|
|
_install_fake_module(monkeypatch, "vllm_omni")
|
|
_install_fake_module(monkeypatch, "vllm_omni.model_executor")
|
|
_install_fake_module(monkeypatch, "vllm_omni.model_executor.models")
|
|
_install_fake_module(
|
|
monkeypatch,
|
|
"vllm_omni.model_executor.models.common",
|
|
qwen3_code_predictor=fake_qwen3_code_predictor,
|
|
)
|
|
_install_fake_module(
|
|
monkeypatch,
|
|
"vllm_omni.model_executor.models.qwen3_tts",
|
|
prompt_embeds_builder=fake_prompt_builder,
|
|
qwen3_tts_code2wav=fake_qwen3_tts_code2wav,
|
|
qwen3_tts_code_predictor_vllm=fake_qwen3_tts_code_predictor_vllm,
|
|
qwen3_tts_talker=fake_talker,
|
|
)
|
|
return (
|
|
fake_qwen3_code_predictor,
|
|
fake_qwen3_tts_code_predictor_vllm,
|
|
fake_qwen3_tts_code2wav,
|
|
fake_prompt_builder,
|
|
fake_talker,
|
|
)
|
|
|
|
|
|
def _load_qwen3_tts_patch(monkeypatch: pytest.MonkeyPatch):
|
|
fakes = _install_qwen3_tts_patch_fakes(monkeypatch)
|
|
path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "_310p" / "patch" / "qwen3_tts.py"
|
|
module = _load_source_module("vllm_omni_test_310p_qwen3_tts_patch", path)
|
|
return module, fakes
|
|
|
|
|
|
def test_registry_applies_worker_once_and_model_patch_lazily(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
registry_path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "_310p" / "patch" / "__init__.py"
|
|
registry = _load_source_module("vllm_omni_test_310p_patch_registry", registry_path)
|
|
calls = {"worker": 0, "talker": 0, "code2wav": 0}
|
|
|
|
_install_fake_module(
|
|
monkeypatch,
|
|
"vllm_omni.platforms.npu._310p.patch.worker",
|
|
apply_patch=lambda: calls.__setitem__("worker", calls["worker"] + 1),
|
|
)
|
|
_install_fake_module(
|
|
monkeypatch,
|
|
"vllm_omni.platforms.npu._310p.patch.qwen3_tts",
|
|
apply_talker_patches=lambda: calls.__setitem__("talker", calls["talker"] + 1),
|
|
apply_code2wav_patches=lambda: calls.__setitem__("code2wav", calls["code2wav"] + 1),
|
|
)
|
|
|
|
registry.apply_patches()
|
|
registry.apply_patches()
|
|
registry.apply_model_patches(SimpleNamespace(model_arch="OtherModel"))
|
|
registry.apply_model_patches(SimpleNamespace(model_arch="Qwen3TTSTalkerForConditionalGeneration"))
|
|
registry.apply_model_patches(SimpleNamespace(model_arch="Qwen3TTSCode2Wav"))
|
|
|
|
assert calls == {"worker": 1, "talker": 1, "code2wav": 1}
|
|
|
|
|
|
def test_qwen3_tts_patch_registers_target_classes(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
(
|
|
module,
|
|
(
|
|
fake_code_predictor,
|
|
fake_code_predictor_vllm,
|
|
fake_code2wav,
|
|
fake_prompt_builder,
|
|
fake_talker,
|
|
),
|
|
) = _load_qwen3_tts_patch(monkeypatch)
|
|
|
|
original_common_wrapper = fake_code_predictor.CodePredictorWrapper
|
|
original_vllm_wrapper = fake_code_predictor_vllm.CodePredictorWrapper
|
|
|
|
module.apply_talker_patches()
|
|
module.apply_code2wav_patches()
|
|
|
|
assert fake_talker.Qwen3TTSTalkerForConditionalGeneration is module._Qwen3TTSTalker310P
|
|
assert fake_talker.Qwen3TTSPromptEmbedsBuilder is module._Qwen3TTSPromptEmbedsBuilder310P
|
|
assert fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder is module._Qwen3TTSPromptEmbedsBuilder310P
|
|
assert fake_code_predictor.CodePredictorAttention is module._Qwen3CodePredictorAttention310P
|
|
assert fake_code_predictor.CodePredictorDecoderLayer is module._Qwen3CodePredictorDecoderLayer310P
|
|
assert fake_code_predictor.CodePredictorBaseModel is module._Qwen3CodePredictorBaseModel310P
|
|
assert (
|
|
fake_code_predictor_vllm.Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM
|
|
is module._Qwen3TTSTalkerCodePredictor310P
|
|
)
|
|
assert fake_code_predictor.CodePredictorWrapper is original_common_wrapper
|
|
assert fake_code_predictor_vllm.CodePredictorWrapper is original_vllm_wrapper
|
|
assert fake_code2wav.Qwen3TTSCode2Wav is module._Qwen3TTSCode2Wav310P
|
|
|
|
code2wav = module._Qwen3TTSCode2Wav310P(
|
|
vllm_config=SimpleNamespace(device_config=SimpleNamespace(device=torch.device("cpu")))
|
|
)
|
|
|
|
assert code2wav._npu_decoder_runtime_dtype(torch.device("cpu")) is torch.float16
|
|
|
|
|
|
def test_qwen3_tts_code_predictor_forward_uses_projected_embedding_and_sampling(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
module, _ = _load_qwen3_tts_patch(monkeypatch)
|
|
predictor = module._Qwen3TTSTalkerCodePredictor310P(
|
|
vllm_config=object(),
|
|
config=object(),
|
|
talker_config=object(),
|
|
)
|
|
codec_weight = predictor.model.codec_embedding[0].weight.detach().clone()
|
|
projection_weight = predictor.small_to_mtp_projection.weight.detach().clone()
|
|
|
|
assert predictor.load_weights(iter(())) == {"loaded"}
|
|
torch.testing.assert_close(
|
|
predictor._projected_codec_embed_weight[0],
|
|
torch.nn.functional.linear(codec_weight, projection_weight),
|
|
)
|
|
assert predictor.small_to_mtp_projection.call_shapes == [(8, 4), (8, 4)]
|
|
predictor.small_to_mtp_projection.call_shapes.clear()
|
|
predictor._num_groups = 3
|
|
predictor._model_dtype = torch.float32
|
|
predictor._prefix_graphs_enabled = False
|
|
predictor._bucket_pos_ids = {}
|
|
predictor._device_graphs = {}
|
|
predictor._lm_heads_list = [
|
|
lambda hidden: torch.nn.functional.one_hot(
|
|
torch.full((hidden.shape[0],), 2, dtype=torch.long),
|
|
num_classes=8,
|
|
).to(torch.float32),
|
|
lambda hidden: torch.nn.functional.one_hot(
|
|
torch.full((hidden.shape[0],), 3, dtype=torch.long),
|
|
num_classes=8,
|
|
).to(torch.float32),
|
|
]
|
|
predictor._setup_compile = lambda: None
|
|
predictor._padded_bsz = lambda bsz: bsz
|
|
|
|
def ensure_buffers(device, dtype, padded_bsz):
|
|
predictor._proj_buf = torch.zeros(padded_bsz, 4, 4, device=device, dtype=dtype)
|
|
|
|
predictor._ensure_buffers = ensure_buffers
|
|
predictor._compiled_model_fwd = lambda embeds, _positions: torch.zeros_like(embeds)
|
|
|
|
codes = predictor.forward(
|
|
torch.tensor([1]),
|
|
torch.zeros(1, 4),
|
|
torch.zeros(1, 4),
|
|
do_sample=False,
|
|
)
|
|
|
|
assert codes.tolist() == [[1, 2, 3]]
|
|
assert predictor.small_to_mtp_projection.call_shapes == [(1, 2, 4)]
|
|
torch.testing.assert_close(
|
|
predictor._proj_buf[0, 2],
|
|
predictor._projected_codec_embed_weight[0, 2],
|
|
)
|
|
|
|
filter_calls = []
|
|
sample_calls = []
|
|
|
|
def apply_top_k_top_p(logits, *, p, k, top_k):
|
|
filter_calls.append((p, k, top_k))
|
|
return logits
|
|
|
|
def random_sample(_probs, generators):
|
|
sample_calls.append(generators)
|
|
return torch.tensor([[4]])
|
|
|
|
monkeypatch.setattr(module, "apply_top_k_top_p", apply_top_k_top_p)
|
|
monkeypatch.setattr(module, "random_sample", random_sample)
|
|
predictor._num_groups = 2
|
|
predictor._wrapper_config.sampling_mode = "stored"
|
|
predictor._top_k = 2
|
|
predictor._top_p = 0.8
|
|
predictor._lm_heads_list = predictor._lm_heads_list[:1]
|
|
generator = torch.Generator().manual_seed(1234)
|
|
|
|
sampled_codes = predictor.forward(
|
|
torch.tensor([1]),
|
|
torch.zeros(1, 4),
|
|
torch.zeros(1, 4),
|
|
generator=generator,
|
|
)
|
|
|
|
assert sampled_codes.tolist() == [[1, 4]]
|
|
assert len(filter_calls) == 1
|
|
top_p_tensor, top_k_tensor, top_k_hint = filter_calls[0]
|
|
torch.testing.assert_close(top_p_tensor, torch.tensor([0.8]))
|
|
assert top_k_tensor.tolist() == [2]
|
|
assert top_k_hint == 2
|
|
assert sample_calls == [{0: generator}]
|
|
|
|
|
|
def test_310p_attention_forward_runs_fused_qkv_runtime(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""Exercise the mocked 310P fused-QKV path at runtime on CPU.
|
|
|
|
The reviewer explicitly asked for more than source inspection: this test now
|
|
instantiates the patched attention, runs ``prepare_qkv_weights()``, verifies
|
|
the packed weight/bias cache, and drives the mocked 310P flash-attention op.
|
|
"""
|
|
module, _ = _load_qwen3_tts_patch(monkeypatch)
|
|
|
|
trans_calls = []
|
|
rotary_calls = []
|
|
flash_calls = []
|
|
|
|
def maybe_trans_nz(weight: torch.Tensor) -> torch.Tensor:
|
|
trans_calls.append(weight.detach().clone())
|
|
return weight + 1.0
|
|
|
|
def npu_apply_rotary_pos_emb(
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
cos: torch.Tensor,
|
|
sin: torch.Tensor,
|
|
*,
|
|
layout: str = "BSND",
|
|
rotary_mode: str = "half",
|
|
):
|
|
rotary_calls.append((query.detach().clone(), key.detach().clone(), cos.detach().clone(), sin.detach().clone()))
|
|
return query + cos - sin, key + cos - sin
|
|
|
|
def npu_flash_attention(**kwargs) -> None:
|
|
flash_calls.append(
|
|
{key: value.detach().clone() if torch.is_tensor(value) else value for key, value in kwargs.items()}
|
|
)
|
|
kwargs["out"].copy_((kwargs["query"] + kwargs["key"] + kwargs["value"]).to(kwargs["out"].dtype))
|
|
|
|
monkeypatch.setattr(module, "maybe_trans_nz", maybe_trans_nz)
|
|
monkeypatch.setattr(module, "aligned_16", lambda tensor: tensor)
|
|
monkeypatch.setattr(module.torch_npu, "npu_apply_rotary_pos_emb", npu_apply_rotary_pos_emb, raising=False)
|
|
monkeypatch.setattr(module.torch_npu, "_npu_flash_attention", npu_flash_attention, raising=False)
|
|
|
|
config = SimpleNamespace(
|
|
num_attention_heads=2,
|
|
num_key_value_heads=2,
|
|
head_dim=2,
|
|
hidden_size=4,
|
|
num_code_groups=2,
|
|
attention_bias=True,
|
|
rms_norm_eps=1e-6,
|
|
)
|
|
attention = module._Qwen3CodePredictorAttention310P(config)
|
|
attention.num_heads = 2
|
|
attention.num_kv_heads = 2
|
|
attention.head_dim = 2
|
|
attention.scaling = 0.5
|
|
attention._q_size = attention.num_heads * attention.head_dim
|
|
attention._kv_size = attention.num_kv_heads * attention.head_dim
|
|
attention.q_norm = torch.nn.Identity()
|
|
attention.k_norm = torch.nn.Identity()
|
|
attention.o_proj = torch.nn.Identity()
|
|
attention.qkv_proj = torch.nn.Linear(4, attention._q_size + 2 * attention._kv_size, bias=True)
|
|
with torch.no_grad():
|
|
attention.qkv_proj.weight.copy_(torch.arange(48, dtype=torch.float32).reshape(12, 4) / 10)
|
|
attention.qkv_proj.bias.copy_(torch.arange(12, dtype=torch.float32) / 10)
|
|
attention._split_qkv = lambda qkv: qkv.split([attention._q_size, attention._kv_size, attention._kv_size], dim=-1)
|
|
|
|
attention.prepare_qkv_weights()
|
|
|
|
assert len(trans_calls) == 1
|
|
assert trans_calls[0].is_contiguous()
|
|
expected_fused_weight = attention.qkv_proj.weight.detach().clone().contiguous() + 1.0
|
|
torch.testing.assert_close(attention._fused_qkv_weight, expected_fused_weight)
|
|
assert attention._fused_qkv_bias is attention.qkv_proj.bias
|
|
|
|
hidden_states = (torch.arange(12, dtype=torch.float32).reshape(1, 3, 4) - 3) / 10
|
|
cos = torch.full((1, 3, 2), 0.25, dtype=torch.float32)
|
|
sin = torch.full((1, 3, 2), -0.5, dtype=torch.float32)
|
|
attention_mask = torch.arange(9, dtype=torch.float32).reshape(3, 3)
|
|
|
|
output = attention(hidden_states, (cos, sin), attention_mask=attention_mask)
|
|
|
|
expected_qkv = torch.nn.functional.linear(hidden_states, expected_fused_weight, attention.qkv_proj.bias)
|
|
expected_q_raw, expected_k_raw, expected_v_raw = expected_qkv.split(
|
|
[attention._q_size, attention._kv_size, attention._kv_size],
|
|
dim=-1,
|
|
)
|
|
expected_q = expected_q_raw.view(1, 3, attention.num_heads, attention.head_dim).transpose(1, 2)
|
|
expected_k = expected_k_raw.view(1, 3, attention.num_kv_heads, attention.head_dim).transpose(1, 2)
|
|
expected_v = expected_v_raw.view(1, 3, attention.num_kv_heads, attention.head_dim).transpose(1, 2)
|
|
expected_cos = cos.unsqueeze(1)
|
|
expected_sin = sin.unsqueeze(1)
|
|
expected_q = expected_q + expected_cos - expected_sin
|
|
expected_k = expected_k + expected_cos - expected_sin
|
|
expected_q_f = expected_q.transpose(1, 2).reshape(3, attention.num_heads, attention.head_dim)
|
|
expected_k_f = expected_k.transpose(1, 2).reshape(3, attention.num_kv_heads, attention.head_dim)
|
|
expected_v_f = expected_v.transpose(1, 2).reshape(3, attention.num_kv_heads, attention.head_dim)
|
|
expected_flash = (expected_q_f + expected_k_f + expected_v_f).to(torch.float16).to(torch.float32)
|
|
|
|
# One fused aclnn rope call covers q and k in the 4D BSND layout.
|
|
assert len(rotary_calls) == 1
|
|
assert rotary_calls[0][0].shape == (1, 3, attention.num_heads, attention.head_dim)
|
|
assert rotary_calls[0][1].shape == (1, 3, attention.num_kv_heads, attention.head_dim)
|
|
assert rotary_calls[0][2].shape == (1, 3, 1, attention.head_dim)
|
|
assert rotary_calls[0][3].shape == (1, 3, 1, attention.head_dim)
|
|
assert len(flash_calls) == 1
|
|
torch.testing.assert_close(flash_calls[0]["query"], expected_q_f)
|
|
torch.testing.assert_close(flash_calls[0]["key"], expected_k_f)
|
|
torch.testing.assert_close(flash_calls[0]["value"], expected_v_f)
|
|
torch.testing.assert_close(flash_calls[0]["mask"], attention_mask)
|
|
torch.testing.assert_close(flash_calls[0]["seq_len"], torch.tensor([3], dtype=torch.int32))
|
|
assert flash_calls[0]["scale_value"] == pytest.approx(0.5)
|
|
assert flash_calls[0]["num_heads"] == 2
|
|
assert flash_calls[0]["num_kv_heads"] == 2
|
|
torch.testing.assert_close(output, expected_flash.reshape(1, 3, 4))
|
|
|
|
|
|
def test_qwen3_tts_talker_patch_uses_fp16_runtime_dtype(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
module, _ = _load_qwen3_tts_patch(monkeypatch)
|
|
talker = module._Qwen3TTSTalker310P(vllm_config=object())
|
|
|
|
assert talker._embedding_dtype is torch.float16
|
|
assert talker._prompt_builder._embedding_dtype is torch.float16
|
|
assert talker.talker_mtp_graph_safe is False
|
|
assert talker.talker_mtp_accepts_per_row_generators is True
|
|
assert talker.load_weights([]) == {"loaded"}
|
|
assert talker.encoder.to_calls[-1] == {"device": torch.device("cpu"), "dtype": torch.float32}
|
|
|
|
codes = talker._encode_ref_audio_batch([np.zeros(8, dtype=np.float32)], 24000, device=torch.device("cpu"))
|
|
|
|
assert talker.encoder.last_input_dtype is torch.float32
|
|
assert len(codes) == 1
|
|
assert codes[0].dtype is torch.long
|
|
assert codes[0].shape == (4, 2)
|
|
|
|
|
|
def test_qwen3_tts_prompt_patch_sets_cpu_mel_front_end_flag(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
module, _ = _load_qwen3_tts_patch(monkeypatch)
|
|
|
|
# The 310P builder keeps the base implementation and only opts into the
|
|
# CPU mel front-end, because the NPU has no torch.stft.
|
|
assert issubclass(
|
|
module._Qwen3TTSPromptEmbedsBuilder310P,
|
|
module.prompt_embeds_builder.Qwen3TTSPromptEmbedsBuilder,
|
|
)
|
|
assert module._Qwen3TTSPromptEmbedsBuilder310P._mel_spectrogram_on_cpu is True
|
|
|
|
|
|
def test_qwen3_tts_tokenizer_npu_patch_dispatches_fused_ops(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
rotary_calls = []
|
|
rms_calls = []
|
|
|
|
def rotary_mul(hidden_states, cos, sin):
|
|
rotary_calls.append((hidden_states, cos, sin))
|
|
return hidden_states + 1
|
|
|
|
def rms_norm(hidden_states, weight, *, epsilon):
|
|
rms_calls.append((hidden_states, weight, epsilon))
|
|
return hidden_states * weight, None
|
|
|
|
_install_fake_module(
|
|
monkeypatch,
|
|
"torch_npu",
|
|
npu_rotary_mul=rotary_mul,
|
|
npu_rms_norm=rms_norm,
|
|
)
|
|
_install_fake_module(monkeypatch, "vllm")
|
|
_install_fake_module(monkeypatch, "vllm.logger", init_logger=lambda _name: SimpleNamespace(debug=lambda *_: None))
|
|
|
|
rotary_path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "layers" / "rotary_embedding.py"
|
|
rotary_module = _load_source_module("vllm_omni_test_npu_rotary_embedding", rotary_path)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"vllm_omni.platforms.npu.layers.rotary_embedding",
|
|
rotary_module,
|
|
)
|
|
|
|
class FakeRMSNorm(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.weight = torch.nn.Parameter(torch.tensor([2.0, 3.0]))
|
|
self.variance_epsilon = 1e-5
|
|
|
|
def forward(self, hidden_states):
|
|
return hidden_states
|
|
|
|
def original_rope(q, k, cos, sin):
|
|
del cos, sin
|
|
return q, k
|
|
|
|
tokenizer = _install_fake_module(
|
|
monkeypatch,
|
|
"vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2",
|
|
Qwen3TTSTokenizerV2DecoderRMSNorm=FakeRMSNorm,
|
|
apply_rotary_pos_emb=original_rope,
|
|
)
|
|
tokenizer_package = _install_fake_module(
|
|
monkeypatch,
|
|
"vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz",
|
|
modeling_qwen3_tts_tokenizer_v2=tokenizer,
|
|
)
|
|
_install_fake_module(monkeypatch, "vllm_omni")
|
|
_install_fake_module(monkeypatch, "vllm_omni.model_executor")
|
|
_install_fake_module(monkeypatch, "vllm_omni.model_executor.models")
|
|
_install_fake_module(monkeypatch, "vllm_omni.model_executor.models.qwen3_tts")
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz",
|
|
tokenizer_package,
|
|
)
|
|
|
|
path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "models" / "qwen3_tts_tokenizer_v2.py"
|
|
module = _load_source_module("vllm_omni_test_qwen3_tts_tokenizer_npu_patch", path)
|
|
module.apply_qwen3_tts_tokenizer_v2_patch()
|
|
|
|
q = torch.zeros(1, 2, 3, 2)
|
|
k = torch.ones_like(q)
|
|
cos = torch.zeros(1, 3, 2)
|
|
sin = torch.ones_like(cos)
|
|
q_out, k_out = tokenizer.apply_rotary_pos_emb(q, k, cos, sin)
|
|
norm = FakeRMSNorm()
|
|
norm_out = norm(torch.ones(1, 2))
|
|
|
|
assert len(rotary_calls) == 2
|
|
assert rotary_calls[0][1].shape == (1, 3, 1, 2)
|
|
torch.testing.assert_close(q_out, q + 1)
|
|
torch.testing.assert_close(k_out, k + 1)
|
|
assert len(rms_calls) == 1
|
|
assert rms_calls[0][2] == pytest.approx(1e-5)
|
|
torch.testing.assert_close(norm_out, torch.tensor([[2.0, 3.0]]))
|
|
|
|
|
|
def _load_qwen3_tts_code2wav_npu_patch(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
*,
|
|
platform_dtype_override: bool = True,
|
|
):
|
|
linear_weights = []
|
|
conv_weights = []
|
|
|
|
def maybe_trans_nz(weight):
|
|
linear_weights.append(weight)
|
|
return weight
|
|
|
|
def format_cast(weight, fmt):
|
|
conv_weights.append((weight, fmt))
|
|
return weight
|
|
|
|
class FakeDecoder(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear = torch.nn.Linear(4, 4)
|
|
self.conv = torch.nn.Conv1d(4, 4, 3)
|
|
self.deconv = torch.nn.ConvTranspose1d(4, 4, 4)
|
|
self.grouped_conv = torch.nn.Conv1d(4, 4, 3, groups=2)
|
|
self.cache_precompute_calls = 0
|
|
|
|
def precompute_snake_caches(self):
|
|
self.cache_precompute_calls += 1
|
|
|
|
class FakeCode2Wav:
|
|
def __init__(self, *, vllm_config, prefix=""):
|
|
self.vllm_config = vllm_config
|
|
self.prefix = prefix
|
|
self.decoder = FakeDecoder()
|
|
|
|
def load_weights(self, weights):
|
|
assert list(weights) == []
|
|
self.decoder.to(dtype=self.vllm_config.model_config.dtype)
|
|
self.decoder.precompute_snake_caches()
|
|
return {"loaded"}
|
|
|
|
if platform_dtype_override:
|
|
|
|
def _npu_decoder_runtime_dtype(self, _device: torch.device) -> torch.dtype:
|
|
return torch.float16
|
|
|
|
FakeCode2Wav._npu_decoder_runtime_dtype = _npu_decoder_runtime_dtype # type: ignore[attr-defined]
|
|
|
|
logger = SimpleNamespace(info=lambda *_: None, debug=lambda *_: None, warning=lambda *_: None)
|
|
current_platform = SimpleNamespace(is_npu=lambda: False)
|
|
|
|
target = _install_fake_module(
|
|
monkeypatch,
|
|
"vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav",
|
|
Qwen3TTSCode2Wav=FakeCode2Wav,
|
|
)
|
|
_install_fake_module(monkeypatch, "torch_npu", npu_format_cast=format_cast)
|
|
_install_fake_module(monkeypatch, "vllm")
|
|
_install_fake_module(monkeypatch, "vllm.config", VllmConfig=object)
|
|
_install_fake_module(monkeypatch, "vllm.logger", init_logger=lambda _name: logger)
|
|
_install_fake_module(monkeypatch, "vllm_ascend")
|
|
_install_fake_module(monkeypatch, "vllm_ascend.utils", maybe_trans_nz=maybe_trans_nz)
|
|
_install_fake_module(monkeypatch, "vllm_omni")
|
|
_install_fake_module(monkeypatch, "vllm_omni.platforms", current_omni_platform=current_platform)
|
|
# Not A5: keep the FRACTAL_Z conv weight layout for 310P.
|
|
_install_fake_module(monkeypatch, "vllm_omni.platforms.npu", is_a5=lambda: False)
|
|
|
|
path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "models" / "qwen3_tts.py"
|
|
module = _load_source_module("vllm_omni_test_qwen3_tts_code2wav_npu_patch", path)
|
|
module.apply_qwen3_tts_code2wav_patch()
|
|
return module, target, linear_weights, conv_weights
|
|
|
|
|
|
def _make_patched_code2wav(target, *, dtype=torch.bfloat16):
|
|
return target.Qwen3TTSCode2Wav(
|
|
vllm_config=SimpleNamespace(
|
|
device_config=SimpleNamespace(device=torch.device("cpu")),
|
|
model_config=SimpleNamespace(dtype=dtype),
|
|
),
|
|
prefix="stage1",
|
|
)
|
|
|
|
|
|
def test_qwen3_tts_code2wav_npu_patch_prepares_loaded_decoder(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
module, target, linear_weights, conv_weights = _load_qwen3_tts_code2wav_npu_patch(monkeypatch)
|
|
model = _make_patched_code2wav(target, dtype=torch.bfloat16)
|
|
|
|
assert model.load_weights(iter(())) == {"loaded"}
|
|
|
|
assert model.prefix == "stage1"
|
|
assert model.decoder.linear.weight.dtype is torch.float16
|
|
assert [weight.data_ptr() for weight in linear_weights] == [model.decoder.linear.weight.data_ptr()]
|
|
assert {weight.data_ptr() for weight, _ in conv_weights} == {
|
|
model.decoder.conv.weight.data_ptr(),
|
|
model.decoder.deconv.weight.data_ptr(),
|
|
}
|
|
assert all(fmt == module._ACL_FORMAT_FRACTAL_Z for _, fmt in conv_weights)
|
|
assert model.decoder.cache_precompute_calls == 2
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
|
|
def test_qwen3_tts_code2wav_npu_patch_preserves_stage_dtype(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
dtype: torch.dtype,
|
|
) -> None:
|
|
_, target, _, _ = _load_qwen3_tts_code2wav_npu_patch(monkeypatch, platform_dtype_override=False)
|
|
model = _make_patched_code2wav(target, dtype=dtype)
|
|
|
|
assert model.load_weights(iter(())) == {"loaded"}
|
|
assert all(parameter.dtype is dtype for parameter in model.decoder.parameters())
|
|
assert model.decoder.cache_precompute_calls == 1
|