diff --git a/tests/model_executor/models/moss_tts/test_moss_fused_load.py b/tests/model_executor/models/moss_tts/test_moss_fused_load.py deleted file mode 100644 index 60039c693..000000000 --- a/tests/model_executor/models/moss_tts/test_moss_fused_load.py +++ /dev/null @@ -1,120 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Regression tests for MOSS-TTS-Realtime fused code-predictor loading. - -The MOSS-TTS-Realtime depth transformer shares ``CodePredictorBaseModel``, whose -q/k/v and gate/up projections are fused. Before the fix, the talker loaded those -shards with a per-tensor ``default_weight_loader``, which silently dropped every -q/k/v/gate/up shard and left the fused params randomly initialized -- no -exception, just a corrupted model. - -These tests assert the property that failure mode violates: after -``load_weights``, every fused parameter is both reported loaded and numerically -equal to the concatenation of the shards that fed it. -""" - -import pytest -import torch - -from vllm_omni.model_executor.models.moss_tts.configuration_moss_tts import ( - MossTTSLocalTransformerConfig, -) -from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local import ( - MossTTSRealtimeLocalTransformer, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def _tiny_cfg(): - return MossTTSLocalTransformerConfig( - head_dim=8, - hidden_size=32, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - intermediate_size=64, - max_position_embeddings=8, - ) - - -def _hf_shards(cfg, *, drop: set[str] | None = None): - """Synthetic HF-style checkpoint tensors for the depth transformer.""" - drop = drop or set() - h, inter = cfg.hidden_size, cfg.intermediate_size - q_out = cfg.num_attention_heads * cfg.head_dim - kv_out = cfg.num_key_value_heads * cfg.head_dim - out = {} - for layer in range(cfg.num_hidden_layers): - p = f"model.layers.{layer}" - for nm, rows in (("q_proj", q_out), ("k_proj", kv_out), ("v_proj", kv_out)): - key = f"{p}.self_attn.{nm}.weight" - if nm not in drop: - out[key] = torch.randn(rows, h) - out[f"{p}.self_attn.o_proj.weight"] = torch.randn(h, q_out) - for nm in ("gate_proj", "up_proj"): - key = f"{p}.mlp.{nm}.weight" - if nm not in drop: - out[key] = torch.randn(inter, h) - out[f"{p}.mlp.down_proj.weight"] = torch.randn(h, inter) - out[f"{p}.input_layernorm.weight"] = torch.randn(h) - out[f"{p}.post_attention_layernorm.weight"] = torch.randn(h) - for nm in ("q_norm", "k_norm"): - out[f"{p}.self_attn.{nm}.weight"] = torch.randn(cfg.head_dim) - out["model.norm.weight"] = torch.randn(h) - return out - - -def test_every_fused_param_is_loaded_and_correctly_packed(): - cfg = _tiny_cfg() - lt = MossTTSRealtimeLocalTransformer(cfg) - shards = _hf_shards(cfg) - - loaded = lt.load_weights(iter([(k, v) for k, v in shards.items()])) - - params = dict(lt.named_parameters()) - fused = [n for n in params if ".self_attn.qkv_proj." in n or ".mlp.gate_up_proj." in n] - assert fused, "model declares no fused params -- test would be vacuous" - - # 1. Reported as loaded (this is what the old path silently failed). - unreported = [n for n in fused if n not in loaded] - assert not unreported, f"fused params not reported loaded: {unreported}" - - # 2. Numerically equal to the concatenation of their shards, in order. - for layer in range(cfg.num_hidden_layers): - p = f"model.layers.{layer}" - want_qkv = torch.cat([shards[f"{p}.self_attn.{n}.weight"] for n in ("q_proj", "k_proj", "v_proj")], dim=0) - got_qkv = params[f"{p}.self_attn.qkv_proj.weight"].detach() - assert torch.equal(got_qkv, want_qkv), f"layer {layer}: qkv packing mismatch" - - want_gu = torch.cat([shards[f"{p}.mlp.{n}.weight"] for n in ("gate_proj", "up_proj")], dim=0) - got_gu = params[f"{p}.mlp.gate_up_proj.weight"].detach() - assert torch.equal(got_gu, want_gu), f"layer {layer}: gate_up packing mismatch" - - -def test_incomplete_shard_group_raises_instead_of_corrupting(): - """A checkpoint missing v_proj must fail loudly, not load a random qkv.""" - cfg = _tiny_cfg() - lt = MossTTSRealtimeLocalTransformer(cfg) - shards = _hf_shards(cfg, drop={"v_proj"}) - - with pytest.raises(RuntimeError, match="incomplete fused shards"): - lt.load_weights(iter([(k, v) for k, v in shards.items()])) - - -def test_absent_fused_group_raises(): - """No q/k/v at all must also raise, via the missing-fused-param guard.""" - cfg = _tiny_cfg() - lt = MossTTSRealtimeLocalTransformer(cfg) - shards = _hf_shards(cfg, drop={"q_proj", "k_proj", "v_proj"}) - - with pytest.raises(RuntimeError, match="missing fused parameters"): - lt.load_weights(iter([(k, v) for k, v in shards.items()])) - - -def test_returned_names_keep_the_model_prefix(): - """Returned names must be ``model.*`` for talker re-prefixing.""" - cfg = _tiny_cfg() - lt = MossTTSRealtimeLocalTransformer(cfg) - loaded = lt.load_weights(iter([(k, v) for k, v in _hf_shards(cfg).items()])) - assert loaded, "nothing reported loaded" - assert all(n.startswith("model.") for n in loaded), sorted(n for n in loaded if not n.startswith("model.")) diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py index 01ff0cdaf..5e31c2cbd 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py @@ -674,291 +674,3 @@ class TestGumbelMaxSampling: empirical = torch.bincount(samples, minlength=vocab).float() / n_draws # L1 distance under 0.03 with 20k draws is comfortable for vocab=16. assert (empirical - target).abs().sum().item() < 0.05 - - -class TestCodePredictorFusedProjections: - """Reference tests for the fused qkv_proj / gate_up_proj paths. - - These CPU tests confirm that the packed projections implement the same - linear algebra as the unfused reference under this setup. GPU executions - may still diverge at the bit level when kernel choices change. - """ - - def test_fused_qkv_matches_separate(self, loaded_target_classes) -> None: - torch.manual_seed(0) - hidden = 32 - num_heads, num_kv_heads, head_dim = 4, 2, 8 - q_size = num_heads * head_dim - kv_size = num_kv_heads * head_dim - x = torch.randn(2, 5, hidden) - - q = torch.nn.Linear(hidden, q_size, bias=False) - k = torch.nn.Linear(hidden, kv_size, bias=False) - v = torch.nn.Linear(hidden, kv_size, bias=False) - fused = torch.nn.Linear(hidden, q_size + 2 * kv_size, bias=False) - with torch.no_grad(): - fused.weight.copy_(torch.cat([q.weight, k.weight, v.weight], dim=0)) - - ref_q, ref_k, ref_v = q(x), k(x), v(x) - out_q, out_k, out_v = fused(x).split([q_size, kv_size, kv_size], dim=-1) - assert torch.equal(ref_q, out_q) - assert torch.equal(ref_k, out_k) - assert torch.equal(ref_v, out_v) - - def test_attention_split_redensifies_kv_slices(self, loaded_target_classes) -> None: - cp_config, _ = _make_tiny_config(loaded_target_classes) - common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] - attn = common_mod.CodePredictorAttention(cp_config) - - qkv = torch.randn(2, 5, attn._q_size + 2 * attn._kv_size) - q_ref, k_ref, v_ref = qkv.split([attn._q_size, attn._kv_size, attn._kv_size], dim=-1) - q_raw, k_raw, v_raw = attn._split_qkv(qkv) - - assert torch.equal(q_raw, q_ref) - assert torch.equal(k_raw, k_ref) - assert torch.equal(v_raw, v_ref) - assert k_raw.is_contiguous() - assert v_raw.is_contiguous() - - def test_fused_gate_up_matches_separate(self, loaded_target_classes) -> None: - torch.manual_seed(1) - hidden, intermediate = 32, 64 - x = torch.randn(2, 5, hidden) - - g = torch.nn.Linear(hidden, intermediate, bias=False) - u = torch.nn.Linear(hidden, intermediate, bias=False) - fused = torch.nn.Linear(hidden, 2 * intermediate, bias=False) - with torch.no_grad(): - fused.weight.copy_(torch.cat([g.weight, u.weight], dim=0)) - - ref_g, ref_u = g(x), u(x) - out_g, out_u = fused(x).split([intermediate, intermediate], dim=-1) - assert torch.equal(ref_g, out_g) - assert torch.equal(ref_u, out_u) - - def test_load_weights_repacks_hf_shards(self, mocker: MockerFixture, loaded_target_classes) -> None: - """``load_weights`` must transparently re-pack HF q/k/v and gate/up - shards into the fused ``qkv_proj`` and ``gate_up_proj`` weights.""" - _, _, _, code_predictor_model, _ = loaded_target_classes - cp_config, _ = _make_tiny_config(loaded_target_classes) - model = code_predictor_model(cp_config, embedding_dim=cp_config.hidden_size) - - # The test harness mocks ``default_weight_loader`` as a no-op; install a - # real copier so this test verifies both the returned names and the - # packed parameter values/order. - common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] - mocker.patch.object( - common_mod, - "default_weight_loader", - lambda param, weight: param.data.copy_(weight), - ) - - # Build a synthetic "HF checkpoint" that exposes separate q/k/v - # and gate/up shards, mirroring what HuggingFace ships. - param_dict = dict(model.named_parameters(remove_duplicate=False)) - weights: list[tuple[str, torch.Tensor]] = [] - expected_fused: dict[str, torch.Tensor] = {} - torch.manual_seed(42) - for name, param in param_dict.items(): - if name.endswith(".self_attn.qkv_proj.weight"): - prefix = name[: -len(".qkv_proj.weight")] - num_heads = cp_config.num_attention_heads - num_kv_heads = cp_config.num_key_value_heads - head_dim = cp_config.head_dim - q_size = num_heads * head_dim - kv_size = num_kv_heads * head_dim - hidden = cp_config.hidden_size - q_w = torch.randn(q_size, hidden) - k_w = torch.randn(kv_size, hidden) - v_w = torch.randn(kv_size, hidden) - weights.append((f"{prefix}.q_proj.weight", q_w)) - weights.append((f"{prefix}.k_proj.weight", k_w)) - weights.append((f"{prefix}.v_proj.weight", v_w)) - expected_fused[name] = torch.cat([q_w, k_w, v_w], dim=0) - elif name.endswith(".mlp.gate_up_proj.weight"): - prefix = name[: -len(".gate_up_proj.weight")] - inter = cp_config.intermediate_size - hidden = cp_config.hidden_size - gate_w = torch.randn(inter, hidden) - up_w = torch.randn(inter, hidden) - weights.append((f"{prefix}.gate_proj.weight", gate_w)) - weights.append((f"{prefix}.up_proj.weight", up_w)) - expected_fused[name] = torch.cat([gate_w, up_w], dim=0) - else: - # Pass-through parameters (norms, embeddings, o_proj, down_proj). - weights.append((name, torch.randn_like(param))) - - loaded = model.load_weights(weights) - after = dict(model.named_parameters(remove_duplicate=False)) - assert expected_fused, "test config must contain at least one fused projection layer" - for fname, expected in expected_fused.items(): - assert fname in loaded, f"{fname} was not assembled by load_weights" - assert torch.equal(after[fname], expected), f"{fname} does not match the packed shard order" - - def test_load_weights_repacks_bias_shards(self, mocker: MockerFixture, loaded_target_classes) -> None: - """When ``attention_bias=True``, the q/k/v ``.bias`` shards must be - concatenated into ``qkv_proj.bias`` instead of being dropped (which - would leave the fused bias randomly initialized).""" - _, _, _, code_predictor_model, _ = loaded_target_classes - cp_config, _ = _make_tiny_config(loaded_target_classes) - cp_config.attention_bias = True - model = code_predictor_model(cp_config, embedding_dim=cp_config.hidden_size) - - # The test harness mocks ``default_weight_loader`` as a no-op; install a - # real copier so we can assert the packed values land on the parameter. - common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] - mocker.patch.object( - common_mod, - "default_weight_loader", - lambda param, weight: param.data.copy_(weight), - ) - - param_dict = dict(model.named_parameters(remove_duplicate=False)) - weights: list[tuple[str, torch.Tensor]] = [] - expected_bias: dict[str, torch.Tensor] = {} - torch.manual_seed(11) - for name, param in param_dict.items(): - if name.endswith(".self_attn.qkv_proj.weight"): - prefix = name[: -len(".qkv_proj.weight")] - q_size = cp_config.num_attention_heads * cp_config.head_dim - kv_size = cp_config.num_key_value_heads * cp_config.head_dim - hidden = cp_config.hidden_size - q_w, k_w, v_w = (torch.randn(s, hidden) for s in (q_size, kv_size, kv_size)) - q_b, k_b, v_b = (torch.randn(s) for s in (q_size, kv_size, kv_size)) - weights.append((f"{prefix}.q_proj.weight", q_w)) - weights.append((f"{prefix}.k_proj.weight", k_w)) - weights.append((f"{prefix}.v_proj.weight", v_w)) - weights.append((f"{prefix}.q_proj.bias", q_b)) - weights.append((f"{prefix}.k_proj.bias", k_b)) - weights.append((f"{prefix}.v_proj.bias", v_b)) - expected_bias[f"{prefix}.qkv_proj.bias"] = torch.cat([q_b, k_b, v_b], dim=0) - elif name.endswith(".mlp.gate_up_proj.weight"): - prefix = name[: -len(".gate_up_proj.weight")] - inter, hidden = cp_config.intermediate_size, cp_config.hidden_size - weights.append((f"{prefix}.gate_proj.weight", torch.randn(inter, hidden))) - weights.append((f"{prefix}.up_proj.weight", torch.randn(inter, hidden))) - elif name.endswith(".qkv_proj.bias"): - # Fused target is assembled from the q/k/v bias shards above; a - # real HF checkpoint never ships a qkv_proj.bias tensor directly. - continue - else: - weights.append((name, torch.randn_like(param))) - - loaded = model.load_weights(weights) - after = dict(model.named_parameters()) - assert expected_bias, "test config must contain at least one qkv layer" - for bname, exp in expected_bias.items(): - assert bname in loaded, f"{bname} bias was not assembled by load_weights" - assert torch.equal(after[bname], exp), f"{bname} does not match the packed q/k/v bias" - - def test_wrapper_load_weights_handles_model_prefixed_shards( - self, - mocker: MockerFixture, - loaded_target_classes, - ) -> None: - _, _, code_predictor_wrapper, _, _ = loaded_target_classes - cp_config, talker_config = _make_tiny_config(loaded_target_classes) - vllm_config = _make_vllm_config(mocker) - predictor = code_predictor_wrapper( - vllm_config=vllm_config, - config=cp_config, - talker_config=talker_config, - ) - - common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] - mocker.patch.object( - common_mod, - "default_weight_loader", - lambda param, weight: param.data.copy_(weight), - ) - - param_dict = dict(predictor.named_parameters(remove_duplicate=False)) - weights: list[tuple[str, torch.Tensor]] = [] - expected_fused: dict[str, torch.Tensor] = {} - torch.manual_seed(7) - for name, param in param_dict.items(): - if name.endswith(".self_attn.qkv_proj.weight"): - prefix = name[: -len(".qkv_proj.weight")] - q_size = cp_config.num_attention_heads * cp_config.head_dim - kv_size = cp_config.num_key_value_heads * cp_config.head_dim - hidden = cp_config.hidden_size - q_w = torch.randn(q_size, hidden) - k_w = torch.randn(kv_size, hidden) - v_w = torch.randn(kv_size, hidden) - weights.append((f"{prefix}.q_proj.weight", q_w)) - weights.append((f"{prefix}.k_proj.weight", k_w)) - weights.append((f"{prefix}.v_proj.weight", v_w)) - expected_fused[name] = torch.cat([q_w, k_w, v_w], dim=0) - elif name.endswith(".mlp.gate_up_proj.weight"): - prefix = name[: -len(".gate_up_proj.weight")] - inter = cp_config.intermediate_size - hidden = cp_config.hidden_size - gate_w = torch.randn(inter, hidden) - up_w = torch.randn(inter, hidden) - weights.append((f"{prefix}.gate_proj.weight", gate_w)) - weights.append((f"{prefix}.up_proj.weight", up_w)) - expected_fused[name] = torch.cat([gate_w, up_w], dim=0) - else: - weights.append((name, torch.randn_like(param))) - - loaded = predictor.load_weights(weights) - after = dict(predictor.named_parameters(remove_duplicate=False)) - assert expected_fused, "test config must contain at least one fused projection layer" - for fname, expected in expected_fused.items(): - assert fname in loaded, f"{fname} was not loaded through the outer wrapper" - assert torch.equal(after[fname], expected), f"{fname} does not match the packed shard order" - - def test_load_weights_raises_on_incomplete_shards(self, loaded_target_classes) -> None: - """Missing one of q/k/v (or gate/up) must surface a clear error. - - ``CodePredictorBaseModel`` is the inner transformer (its parameter - paths look like ``layers..self_attn.qkv_proj.weight`` -- without - the ``model.`` prefix that the outer ``CodePredictorWrapper`` strips - before forwarding). The test mimics that contract. - """ - _, _, _, code_predictor_model, _ = loaded_target_classes - cp_config, _ = _make_tiny_config(loaded_target_classes) - model = code_predictor_model(cp_config, embedding_dim=cp_config.hidden_size) - - # Find any actual self_attn layer prefix, to make the test robust - # to renaming of the inner model's parameter scheme. - attn_prefix = next( - ( - name[: -len(".qkv_proj.weight")] - for name in dict(model.named_parameters()).keys() - if name.endswith(".self_attn.qkv_proj.weight") - ), - None, - ) - assert attn_prefix is not None, "test config must contain at least one self_attn layer" - - num_heads = cp_config.num_attention_heads - num_kv_heads = cp_config.num_key_value_heads - head_dim = cp_config.head_dim - hidden = cp_config.hidden_size - - # Provide only q + k (no v). The fused param exists in the model, - # so load_weights must refuse to silently leave it uninitialized. - bad_weights: list[tuple[str, torch.Tensor]] = [ - (f"{attn_prefix}.q_proj.weight", torch.randn(num_heads * head_dim, hidden)), - (f"{attn_prefix}.k_proj.weight", torch.randn(num_kv_heads * head_dim, hidden)), - # v_proj missing on purpose. - ] - with pytest.raises(RuntimeError, match="incomplete fused shards"): - model.load_weights(bad_weights) - - def test_load_weights_raises_when_fused_shards_absent(self, loaded_target_classes) -> None: - """Entirely absent q/k/v or gate/up groups must not leave fused params - randomly initialized.""" - _, _, _, code_predictor_model, _ = loaded_target_classes - cp_config, _ = _make_tiny_config(loaded_target_classes) - model = code_predictor_model(cp_config, embedding_dim=cp_config.hidden_size) - - weights = [ - (name, torch.randn_like(param)) - for name, param in model.named_parameters(remove_duplicate=False) - if not (name.endswith(".qkv_proj.weight") or name.endswith(".gate_up_proj.weight")) - ] - - with pytest.raises(RuntimeError, match="missing fused parameters"): - model.load_weights(weights) diff --git a/tests/platforms/npu/test_310p_patches.py b/tests/platforms/npu/test_310p_patches.py index 91e74349a..e138a6d63 100644 --- a/tests/platforms/npu/test_310p_patches.py +++ b/tests/platforms/npu/test_310p_patches.py @@ -414,21 +414,6 @@ def test_qwen3_tts_code_predictor_forward_uses_projected_embedding_and_sampling( assert sample_calls == [{0: generator}] -def test_310p_attention_forward_uses_fused_qkv(monkeypatch: pytest.MonkeyPatch) -> None: - """The fused code predictor dropped q_proj/k_proj/v_proj in favor of a single - ``qkv_proj`` (+ ``_split_qkv``). The 310P attention override must follow - suit, otherwise its NPU flash-attention path raises AttributeError on 310P. - The NPU path itself cannot run in CPU CI, so guard the contract statically.""" - import inspect - - module, _ = _load_qwen3_tts_patch(monkeypatch) - source = inspect.getsource(module._Qwen3CodePredictorAttention310P.forward) - assert "qkv_proj" in source and "_split_qkv" in source - assert "self.q_proj" not in source - assert "self.k_proj" not in source - assert "self.v_proj" not in source - - 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()) diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index ea2ea4557..1bea2a124 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -3,8 +3,7 @@ Shared by Qwen3-Omni and Qwen3-TTS talker models. * SDPA attention (F.scaled_dot_product_attention) with native GQA support -* HF-compatible CPU/CUDA numerics with packed qkv/gate_up projections - and NPU-only fused norm/RoPE fast paths +* HF-compatible CPU/CUDA numerics with NPU-only fused norm/RoPE fast paths * Per-call embedding buffer to avoid cross-request aliasing * Pre-allocated position_ids (read-only, safe to persist) * torch.compile (epilogue_fusion=False) on inner transformer by default @@ -46,9 +45,8 @@ if current_omni_platform.is_npu(): # HuggingFace reference on CPU/CUDA. vLLM's fused kernels (RMSNorm, # QKVParallel, get_rope) introduce small precision differences that compound # across the autoregressive steps of the code predictor, causing severe audio -# quality degradation. Packing qkv/gate_up into plain ``nn.Linear`` keeps the -# same algebra but can still change GPU bit-level accumulation. The Ascend -# fused norm/RoPE kernels below are dispatched by the current device platform. +# quality degradation. The Ascend fused norm/RoPE kernels below are dispatched +# by the current device platform. # # See: https://github.com/vllm-project/vllm-omni/issues/2274 @@ -177,23 +175,11 @@ class CodePredictorAttention(nn.Module): self.scaling = self.head_dim**-0.5 self.max_seq = int(config.num_code_groups) + 1 - # Fused QKV projection. - # - # The HF reference uses three separate ``nn.Linear`` modules for - # q/k/v. Packing those weights row-wise into one ``nn.Linear`` keeps - # the same matmul at the algebra level; see the module-level - # HF-numerics note for the GPU bit-level drift caveat. The goal here - # is narrower: reduce eager host launch overhead while keeping the - # downstream q_norm / k_norm / SDPA math and HF checkpoint loading - # contract intact. + # Separate q/k/v projections matching HF (no fused packing) bias = getattr(config, "attention_bias", False) - self._q_size = self.num_heads * self.head_dim - self._kv_size = self.num_kv_heads * self.head_dim - self.qkv_proj = nn.Linear( - self.hidden_size, - self._q_size + 2 * self._kv_size, - bias=bias, - ) + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=bias) + self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=bias) + self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=bias) self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) self.q_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps) @@ -211,15 +197,6 @@ class CodePredictorAttention(nn.Module): ) self.register_buffer("_fusion_causal_mask", fusion_mask, persistent=False) - def _split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - q_raw = qkv[..., : self._q_size] - kv_raw = qkv[..., self._q_size :] - # ``split`` would leave the non-leading k/v slices strided like the - # packed qkv buffer. Materialize them once here so SDPA sees the same - # dense k/v layout as the unfused projection path. - k_raw, v_raw = (t.contiguous() for t in kv_raw.split(self._kv_size, dim=-1)) - return q_raw, k_raw, v_raw - def _forward_npu_attention( self, q: torch.Tensor, @@ -280,14 +257,9 @@ class CodePredictorAttention(nn.Module): hidden_shape_q = (bsz, seq_len, self.num_heads, self.head_dim) hidden_shape_kv = (bsz, seq_len, self.num_kv_heads, self.head_dim) - # Single fused matmul. ``q`` is immediately consumed by RMSNorm; for - # ``k/v`` we re-densify the packed tail once so SDPA does not inherit - # the packed-buffer stride pattern from the fused projection. - qkv = self.qkv_proj(hidden_states) - q_raw, k_raw, v_raw = self._split_qkv(qkv) - q = self.q_norm(q_raw.view(hidden_shape_q)).transpose(1, 2) - k = self.k_norm(k_raw.view(hidden_shape_kv)).transpose(1, 2) - v = v_raw.view(hidden_shape_kv).transpose(1, 2) + q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape_q)).transpose(1, 2) + k = self.k_norm(self.k_proj(hidden_states).view(hidden_shape_kv)).transpose(1, 2) + v = self.v_proj(hidden_states).view(hidden_shape_kv).transpose(1, 2) cos, sin = position_embeddings # cos/sin are [batch, seq_len, head_dim], need unsqueeze at dim=1 for heads @@ -323,22 +295,12 @@ class CodePredictorMLP(nn.Module): def __init__(self, config, *, prefix: str = "") -> None: super().__init__() - # Fused gate_up projection. Same packing idea as ``qkv_proj`` above: - # one matmul instead of two, with the same algebraic result as - # separate gate/up linears. See the module-level HF-numerics note for - # the GPU bit-level drift caveat. - self._intermediate_size = int(config.intermediate_size) - self.gate_up_proj = nn.Linear( - config.hidden_size, - 2 * config.intermediate_size, - bias=False, - ) + self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - gate_up = self.gate_up_proj(hidden_states) - gate, up = gate_up.split([self._intermediate_size, self._intermediate_size], dim=-1) - return self.down_proj(F.silu(gate) * up) + return self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)) # =================================================================== @@ -449,128 +411,18 @@ class CodePredictorBaseModel(nn.Module): hidden_states = self.norm(hidden_states) return hidden_states.to(input_dtype) - # Mapping from HF checkpoint shard name to the (fused parameter, - # slice index) pair used to assemble the fused tensor. See - # ``load_weights`` below. - _FUSED_QKV_SHARDS: tuple[tuple[str, str, int], ...] = ( - ("q_proj", "qkv_proj", 0), - ("k_proj", "qkv_proj", 1), - ("v_proj", "qkv_proj", 2), - ) - _FUSED_GATE_UP_SHARDS: tuple[tuple[str, str, int], ...] = ( - ("gate_proj", "gate_up_proj", 0), - ("up_proj", "gate_up_proj", 1), - ) - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights, transparently re-packing HF q/k/v and gate/up shards - into the fused ``qkv_proj`` and ``gate_up_proj`` parameters. - - HF checkpoints ship q_proj / k_proj / v_proj (and gate_proj / - up_proj) as separate tensors. Our model defines a single - ``qkv_proj`` (``gate_up_proj``) whose weight is the row-wise - concatenation of those shards in q/k/v (gate/up) order. We - buffer the shards per layer prefix and write the fused parameter - in one go once all shards have arrived. - - Both ``.weight`` and ``.bias`` shards are packed this way, so a - code-predictor config with ``attention_bias=True`` is handled - correctly: the q/k/v bias vectors are concatenated into - ``qkv_proj.bias`` instead of being dropped (which would leave the - fused bias randomly initialized). ``gate_up_proj`` has no bias, - so only its ``.weight`` shards are ever routed. - """ params_dict = dict(self.named_parameters(remove_duplicate=False)) loaded_params: set[str] = set() - - # layer_prefix -> {fused_name: {kind: {slice_idx: shard_tensor}}} - # where ``kind`` is "weight" or "bias". - pending: dict[str, dict[str, dict[str, dict[int, torch.Tensor]]]] = {} - - def _try_finalize_fused( - layer_prefix: str, - fused_name: str, - kind: str, - num_shards: int, - ) -> None: - bucket = pending.get(layer_prefix, {}).get(fused_name, {}).get(kind) - if bucket is None or len(bucket) < num_shards: - return - target_param_name = f"{layer_prefix}.{fused_name}.{kind}" - target_param = params_dict.get(target_param_name) - if target_param is None: - return - ordered = [bucket[i] for i in range(num_shards)] - fused_weight = torch.cat(ordered, dim=0) - weight_loader = getattr(target_param, "weight_loader", default_weight_loader) - weight_loader(target_param, fused_weight) - loaded_params.add(target_param_name) - del pending[layer_prefix][fused_name][kind] - if not pending[layer_prefix][fused_name]: - del pending[layer_prefix][fused_name] - if not pending[layer_prefix]: - del pending[layer_prefix] - - def _route_shard( - name: str, - loaded_weight: torch.Tensor, - shard_table: tuple[tuple[str, str, int], ...], - parent_attr: str, - ) -> bool: - """Stash a HF ``.weight``/``.bias`` shard if its fused counterpart - exists in this module; return True iff the shard was consumed.""" - for shard_name, fused_name, slice_idx in shard_table: - for kind in ("weight", "bias"): - suffix = f".{parent_attr}.{shard_name}.{kind}" - if not name.endswith(suffix): - continue - layer_prefix = name[: -len(f".{shard_name}.{kind}")] # ...self_attn or ...mlp - fused_param_name = f"{layer_prefix}.{fused_name}.{kind}" - if fused_param_name not in params_dict: - # e.g. a gate/up ``.bias`` shard when the fused - # projection has no bias: nothing to pack into. - return False - pending.setdefault(layer_prefix, {}).setdefault(fused_name, {}).setdefault(kind, {})[slice_idx] = ( - loaded_weight - ) - _try_finalize_fused(layer_prefix, fused_name, kind, num_shards=len(shard_table)) - return True - return False - for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue - if _route_shard(name, loaded_weight, self._FUSED_QKV_SHARDS, parent_attr="self_attn"): - continue - if _route_shard(name, loaded_weight, self._FUSED_GATE_UP_SHARDS, parent_attr="mlp"): - continue param = params_dict.get(name) if param is None: continue weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) - - if pending: - unfinished = { - lp: {fn: list(kinds.keys()) for fn, kinds in buckets.items()} for lp, buckets in pending.items() - } - raise RuntimeError( - f"CodePredictor load_weights: incomplete fused shards for layers {unfinished}; " - f"check that the checkpoint contains all q/k/v and gate/up tensors." - ) - - missing_fused = sorted( - name - for name in params_dict - if (".self_attn.qkv_proj." in name or ".mlp.gate_up_proj." in name) and name not in loaded_params - ) - if missing_fused: - raise RuntimeError( - f"CodePredictor load_weights: missing fused parameters {missing_fused}; " - f"check that the checkpoint contains the q/k/v and gate/up shards, " - f"or the already-fused qkv_proj/gate_up_proj tensors." - ) return loaded_params @@ -1142,7 +994,7 @@ class CodePredictorWrapper(nn.Module): logger.info("Prepared NPU code predictor weights: linear=%d", linear_count) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load wrapper weights, delegating ``model.*`` tensors to the inner fused loader.""" + """Load weights directly (no fused projection remapping needed).""" loaded: set[str] = set() model_weights: list[tuple[str, torch.Tensor]] = [] other_weights: list[tuple[str, torch.Tensor]] = [] diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py index eb6bfa251..2fd0beefd 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py @@ -23,8 +23,6 @@ previous KV-cache loop -- causal attention, only the last position is read. from __future__ import annotations -from collections.abc import Iterable - import torch import torch.nn as nn import torch.nn.functional as F @@ -63,22 +61,6 @@ class MossTTSRealtimeLocalTransformer(nn.Module): prefix="model", ) - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Delegate ``model.*`` tensors to ``CodePredictorBaseModel.load_weights``. - - The shared body fuses ``q/k/v_proj`` into ``qkv_proj`` and - ``gate/up_proj`` into ``gate_up_proj``; its loader re-packs the HF - shards and raises on incomplete or missing fused parameters. Going - through it (instead of ``default_weight_loader`` per tensor) keeps - those guarantees for the MossTTSRealtime checkpoint too. - """ - model_weights: list[tuple[str, torch.Tensor]] = [] - for name, tensor in weights: - if name.startswith("model."): - model_weights.append((name[len("model.") :], tensor)) - loaded = self.model.load_weights(iter(model_weights)) - return {f"model.{n}" for n in loaded} - @torch.no_grad() def generate_frame( self, diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py index aa7c9b3b0..171e57328 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py @@ -1179,13 +1179,7 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): # ``gate_proj``/``up_proj`` → ``gate_up_proj``. ``default_weight_loader`` # alone leaves those fused params un-initialised, which silently turns # the backbone into a slightly-corrupted model that never emits EOS. - # The local depth transformer (shared ``CodePredictorBaseModel``) has - # the exact same hazard: its ``qkv_proj``/``gate_up_proj`` are fused, - # so its shards must go through ``CodePredictorBaseModel.load_weights`` - # (via ``MossTTSRealtimeLocalTransformer.load_weights``), which re-packs - # them and raises on incomplete/missing fused parameters. backbone_weights: list[tuple[str, torch.Tensor]] = [] - local_weights: list[tuple[str, torch.Tensor]] = [] skipped: list[str] = [] for name, tensor in weights: if name.startswith("language_model."): @@ -1200,9 +1194,8 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): continue # non-persistent buffer, recomputed at runtime if sub.startswith("embed_tokens."): sub = "codec_embedding." + sub[len("embed_tokens.") :] - local_weights.append((f"model.{sub}", tensor)) - continue - if name.startswith("local_transformer.local_lm_heads."): + tgt = "local_transformer.model." + sub + elif name.startswith("local_transformer.local_lm_heads."): tgt = "local_lm_heads." + name[len("local_transformer.local_lm_heads.") :] else: tgt = name @@ -1217,11 +1210,6 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): for n in backbone_loaded: loaded.add(f"model.{n}") - # Delegate depth-transformer weights to the shared fused loader. - local_loaded = self.local_transformer.load_weights(iter(local_weights)) - for n in local_loaded: - loaded.add(f"local_transformer.{n}") - # Build stacked embedding cache for vectorised decode ops. # embed_tokens[0] is the text embedding; [1..n_vq] are audio codebooks. self._stacked_audio_emb_w = torch.stack( diff --git a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py index 36bf663dd..28ca4f362 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py +++ b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py @@ -173,13 +173,16 @@ class Qwen3TTSTalkerCodePredictorConfig(PretrainedConfig): model_type = "qwen3_tts_talker_code_predictor" keys_to_ignore_at_inference = ["past_key_values"] - # The code predictor currently runs TP=1. Do not advertise a tensor - # parallel plan for the packed qkv/gate_up projections until the fused - # loader and ``_split_qkv`` path are taught TP-aware packing/sharding: - # a generic colwise split of the plain ``nn.Linear`` would split the - # packed [q, k, v] layout across ranks instead of preserving each rank's - # local q/k/v groups. - base_model_tp_plan = {} + # Default tensor parallel plan for base model `Qwen3TTSTalkerCodePredictor` + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py b/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py index 136109805..713b67f81 100644 --- a/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py +++ b/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py @@ -157,9 +157,11 @@ class _Qwen3CodePredictorAttention310P(qwen3_code_predictor.CodePredictorAttenti def prepare_qkv_weights(self) -> None: # Pack QKV once so each graph replay uses one matmul and consumes the # weight directly in the 310P matmul layout. - self._fused_qkv_weight = maybe_trans_nz(self.qkv_proj.weight.contiguous()) - if self.qkv_proj.bias is not None: - self._fused_qkv_bias = self.qkv_proj.bias + self._fused_qkv_weight = maybe_trans_nz( + torch.cat((self.q_proj.weight, self.k_proj.weight, self.v_proj.weight), dim=0).contiguous() + ) + if self.q_proj.bias is not None: + self._fused_qkv_bias = torch.cat((self.q_proj.bias, self.k_proj.bias, self.v_proj.bias), dim=0) def forward( self, @@ -168,17 +170,11 @@ class _Qwen3CodePredictorAttention310P(qwen3_code_predictor.CodePredictorAttenti attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: bsz, seq_len, _ = hidden_states.shape - hidden_shape_q = (bsz, seq_len, self.num_heads, self.head_dim) - hidden_shape_kv = (bsz, seq_len, self.num_kv_heads, self.head_dim) - - # ``prepare_qkv_weights`` packs qkv_proj into the 310P matmul layout; - # split it back out the same way before re-applying the 310P - # flash-attention path. qkv = F.linear(hidden_states, self._fused_qkv_weight, self._fused_qkv_bias) - q_raw, k_raw, v_raw = self._split_qkv(qkv) - q = self.q_norm(q_raw.view(hidden_shape_q)).transpose(1, 2) - k = self.k_norm(k_raw.view(hidden_shape_kv)).transpose(1, 2) - v = v_raw.view(hidden_shape_kv).transpose(1, 2) + q, k, v = qkv.split((self._q_size, self._kv_size, self._kv_size), dim=-1) + q = self.q_norm(q.view(bsz, seq_len, self.num_heads, self.head_dim)).transpose(1, 2) + k = self.k_norm(k.view(bsz, seq_len, self.num_kv_heads, self.head_dim)).transpose(1, 2) + v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) cos, sin = position_embeddings cos = cos.unsqueeze(1) @@ -308,7 +304,7 @@ class _Qwen3TTSTalkerCodePredictor310P( for layer in self.model.layers: attention = layer.self_attn attention.prepare_qkv_weights() - qkv_projections.add(attention.qkv_proj) + qkv_projections.update((attention.q_proj, attention.k_proj, attention.v_proj)) for module in self.modules(): if isinstance(module, nn.Linear) and module not in qkv_projections: