[Perf][Qwen3-TTS] Fuse QKV and gate_up projections in code predictor (#4958)

Signed-off-by: sheenlin <sheenlin@tencent.com>
Co-authored-by: sheenlin <sheenlin@tencent.com>
Co-authored-by: yurain <yurain26@users.noreply.github.com>

tests/model_executor/models/moss_tts/test_moss_fused_load.py

@@ -0,0 +1,120 @@
+# 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."))

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -674,3 +674,291 @@ 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.<i>.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)