Revert "[Perf][Qwen3-TTS] Fuse QKV and gate_up projections in code predictor …"

This reverts commit 6fdc4ca62d.

tests/model_executor/models/moss_tts/test_moss_fused_load.py

@@ -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."))

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.<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)

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())

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
+* HF-compatible CPU/CUDA numerics with NPU-only fused norm/RoPE fast paths
-  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
@@ -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
+# quality degradation. The Ascend fused norm/RoPE kernels below are dispatched
-# same algebra but can still change GPU bit-level accumulation. The Ascend
+# by the current device platform.
-# 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.
+        # Separate q/k/v projections matching HF (no fused packing)
-        #
-        # 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_size = self.num_heads * self.head_dim
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=bias)
-        self._kv_size = self.num_kv_heads * self.head_dim
+        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=bias)
-        self.qkv_proj = nn.Linear(
+        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=bias)
-            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)
@@ -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
+        q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape_q)).transpose(1, 2)
-        # ``k/v`` we re-densify the packed tail once so SDPA does not inherit
+        k = self.k_norm(self.k_proj(hidden_states).view(hidden_shape_kv)).transpose(1, 2)
-        # the packed-buffer stride pattern from the fused projection.
+        v = self.v_proj(hidden_states).view(hidden_shape_kv).transpose(1, 2)
-        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
@@ -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:
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
-        # one matmul instead of two, with the same algebraic result as
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
-        # 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:
-        gate_up = self.gate_up_proj(hidden_states)
+        return self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.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)
 
 
 # ===================================================================
@@ -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]] = []

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,

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))
+                tgt = "local_transformer.model." + sub
-                continue
+            elif name.startswith("local_transformer.local_lm_heads."):
-            if 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(

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
+    # Default tensor parallel plan for base model `Qwen3TTSTalkerCodePredictor`
-    # parallel plan for the packed qkv/gate_up projections until the fused
+    base_model_tp_plan = {
-    # loader and ``_split_qkv`` path are taught TP-aware packing/sharding:
+        "layers.*.self_attn.q_proj": "colwise",
-    # a generic colwise split of the plain ``nn.Linear`` would split the
+        "layers.*.self_attn.k_proj": "colwise",
-    # packed [q, k, v] layout across ranks instead of preserving each rank's
+        "layers.*.self_attn.v_proj": "colwise",
-    # local q/k/v groups.
+        "layers.*.self_attn.o_proj": "rowwise",
-    base_model_tp_plan = {}
+        "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"]),

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())
+        self._fused_qkv_weight = maybe_trans_nz(
-        if self.qkv_proj.bias is not None:
+            torch.cat((self.q_proj.weight, self.k_proj.weight, self.v_proj.weight), dim=0).contiguous()
-            self._fused_qkv_bias = self.qkv_proj.bias
+        )
+        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, k, v = qkv.split((self._q_size, self._kv_size, self._kv_size), dim=-1)
-        q = self.q_norm(q_raw.view(hidden_shape_q)).transpose(1, 2)
+        q = self.q_norm(q.view(bsz, seq_len, self.num_heads, self.head_dim)).transpose(1, 2)
-        k = self.k_norm(k_raw.view(hidden_shape_kv)).transpose(1, 2)
+        k = self.k_norm(k.view(bsz, seq_len, self.num_kv_heads, self.head_dim)).transpose(1, 2)
-        v = v_raw.view(hidden_shape_kv).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: