From 6fdc4ca62dc6080bc683552c6ec9e310e9c80867 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9E=97=E9=91=AB?=
 <77673148+l-wave@users.noreply.github.com>
Date: Tue, 4 Aug 2026 19:59:32 +0800
Subject: [PATCH] [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>
---
 .../models/moss_tts/test_moss_fused_load.py   | 120 ++++++++
 .../qwen3_tts/test_code_predictor_dtype.py    | 288 ++++++++++++++++++
 tests/platforms/npu/test_310p_patches.py      |  15 +
 .../models/common/qwen3_code_predictor.py     | 176 ++++++++++-
 .../moss_tts/modeling_moss_tts_local.py       |  18 ++
 .../moss_tts/modeling_moss_tts_talker.py      |  16 +-
 .../qwen3_tts/configuration_qwen3_tts.py      |  17 +-
 .../platforms/npu/_310p/patch/qwen3_tts.py    |  24 +-
 8 files changed, 638 insertions(+), 36 deletions(-)
 create mode 100644 tests/model_executor/models/moss_tts/test_moss_fused_load.py

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
new file mode 100644
index 000000000..60039c693
--- /dev/null
+++ b/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."))
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 5e31c2cbd..01ff0cdaf 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,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)
diff --git a/tests/platforms/npu/test_310p_patches.py b/tests/platforms/npu/test_310p_patches.py
index e138a6d63..91e74349a 100644
--- a/tests/platforms/npu/test_310p_patches.py
+++ b/tests/platforms/npu/test_310p_patches.py
@@ -414,6 +414,21 @@ 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 1bea2a124..ea2ea4557 100644
--- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py
@@ -3,7 +3,8 @@
 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 NPU-only fused norm/RoPE fast paths
+* HF-compatible CPU/CUDA numerics with packed qkv/gate_up projections
+  and 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
@@ -45,8 +46,9 @@ 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. The Ascend fused norm/RoPE kernels below are dispatched
-# by the current device platform.
+# 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.
 #
 # See: https://github.com/vllm-project/vllm-omni/issues/2274
 
@@ -175,11 +177,23 @@ class CodePredictorAttention(nn.Module):
         self.scaling = self.head_dim**-0.5
         self.max_seq = int(config.num_code_groups) + 1
 
-        # Separate q/k/v projections matching HF (no fused packing)
+        # 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.
         bias = getattr(config, "attention_bias", False)
-        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._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.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)
@@ -197,6 +211,15 @@ 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,
@@ -257,9 +280,14 @@ 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)
 
-        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)
+        # 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)
 
         cos, sin = position_embeddings
         # cos/sin are [batch, seq_len, head_dim], need unsqueeze at dim=1 for heads
@@ -295,12 +323,22 @@ class CodePredictorMLP(nn.Module):
 
     def __init__(self, config, *, prefix: str = "") -> None:
         super().__init__()
-        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)
+        # 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.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        return self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
+        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)
 
 
 # ===================================================================
@@ -411,18 +449,128 @@ 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
 
 
@@ -994,7 +1142,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 weights directly (no fused projection remapping needed)."""
+        """Load wrapper weights, delegating ``model.*`` tensors to the inner fused loader."""
         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 2fd0beefd..eb6bfa251 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,6 +23,8 @@ 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
@@ -61,6 +63,22 @@ 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 171e57328..aa7c9b3b0 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,7 +1179,13 @@ 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."):
@@ -1194,8 +1200,9 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
                     continue  # non-persistent buffer, recomputed at runtime
                 if sub.startswith("embed_tokens."):
                     sub = "codec_embedding." + sub[len("embed_tokens.") :]
-                tgt = "local_transformer.model." + sub
-            elif name.startswith("local_transformer.local_lm_heads."):
+                local_weights.append((f"model.{sub}", tensor))
+                continue
+            if name.startswith("local_transformer.local_lm_heads."):
                 tgt = "local_lm_heads." + name[len("local_transformer.local_lm_heads.") :]
             else:
                 tgt = name
@@ -1210,6 +1217,11 @@ 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 28ca4f362..36bf663dd 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,16 +173,13 @@ class Qwen3TTSTalkerCodePredictorConfig(PretrainedConfig):
     model_type = "qwen3_tts_talker_code_predictor"
     keys_to_ignore_at_inference = ["past_key_values"]
 
-    # 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",
-    }
+    # 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 = {}
     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 713b67f81..136109805 100644
--- a/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py
+++ b/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py
@@ -157,11 +157,9 @@ 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(
-            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)
+        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
 
     def forward(
         self,
@@ -170,11 +168,17 @@ 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, 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)
+        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)
 
         cos, sin = position_embeddings
         cos = cos.unsqueeze(1)
@@ -304,7 +308,7 @@ class _Qwen3TTSTalkerCodePredictor310P(
             for layer in self.model.layers:
                 attention = layer.self_attn
                 attention.prepare_qkv_weights()
-                qkv_projections.update((attention.q_proj, attention.k_proj, attention.v_proj))
+                qkv_projections.add(attention.qkv_proj)
 
             for module in self.modules():
                 if isinstance(module, nn.Linear) and module not in qkv_projections: