Add EasyMagpie vLLM-Omni serving

.github/workflows/cicd-main-speech.yml

@@ -41,6 +41,58 @@ jobs:
       dockerfile: docker/Dockerfile
       runner: ${{ inputs.runner }}
 
+  easymagpie-vllm-serving-changes:
+    runs-on: ubuntu-latest
+    outputs:
+      run: ${{ steps.decision.outputs.run }}
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v6
+        with:
+          fetch-depth: 0
+      - name: Get PR info
+        id: get-pr-info
+        if: startsWith(github.ref, 'refs/heads/pull-request/')
+        uses: nv-gha-runners/get-pr-info@main
+      - name: Get changed files
+        id: changed-files
+        uses: step-security/changed-files@v45.0.1
+        with:
+          sha: ${{ github.sha }}
+          base_sha: ${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').base.sha || github.event.pull_request.base.sha }}
+          files: |
+            .github/workflows/cicd-main-speech.yml
+            tools/easymagpie_vllm_omni/**
+            tests/collections/tts/easymagpie_vllm_omni/serving/**
+            tests/functional_tests/L0_Unit_Tests_GPU_TTS_EasyMagpie_vLLM.sh
+      - name: Decide whether to run serving tests
+        id: decision
+        env:
+          ANY_CHANGED: ${{ steps.changed-files.outputs.any_changed }}
+          EVENT_NAME: ${{ github.event_name }}
+          TESTS_TO_RUN: ${{ inputs.test_to_run }}
+        run: |
+          run="$ANY_CHANGED"
+
+          if [[ "$EVENT_NAME" == "schedule" ]]; then
+            run=false
+          elif [[ "$EVENT_NAME" == "workflow_dispatch" ]] && \
+              jq -e 'index("all") != null or index("L0_Unit_Tests_GPU_TTS_EasyMagpie_vLLM") != null' \
+                <<<"$TESTS_TO_RUN" >/dev/null; then
+            run=true
+          fi
+
+          echo "run=$run" | tee -a "$GITHUB_OUTPUT"
+
+  easymagpie-vllm-serving-build:
+    needs: [easymagpie-vllm-serving-changes]
+    if: needs.easymagpie-vllm-serving-changes.outputs.run == 'true'
+    uses: ./.github/workflows/_build_container.yml
+    with:
+      image-name: nemo_container_easymagpie_vllm
+      dockerfile: tools/easymagpie_vllm_omni/Dockerfile.ci
+      runner: ${{ inputs.runner }}
+
   unit-tests:
     strategy:
       fail-fast: false
@@ -146,6 +198,28 @@ jobs:
           cpu-only: ${{ matrix.cpu-only || false }}
           is_optional: ${{ matrix.is-optional || false }}
 
+  easymagpie-vllm-serving-tests:
+    needs: [easymagpie-vllm-serving-build]
+    runs-on: ${{ inputs.runner }}
+    name: L0_Unit_Tests_GPU_TTS_EasyMagpie_vLLM
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v6
+      - name: Checkout NeMo
+        uses: actions/checkout@v6
+        with:
+          path: ${{ github.run_id }}
+      - name: main
+        uses: ./.github/actions/test-template
+        with:
+          runner: ${{ runner.name }}
+          script: L0_Unit_Tests_GPU_TTS_EasyMagpie_vLLM
+          is_unit_test: true
+          tests_to_run: ${{ inputs.test_to_run }}
+          image: ${{ needs.easymagpie-vllm-serving-build.outputs.image }}
+          test-data-path: ${{ vars.DEFAULT_TEST_DATA_PATH || '/mnt/datadrive/TestData' }}
+          timeout: 15
+
   e2e-tests:
     strategy:
       fail-fast: false

tests/collections/tts/easymagpie_vllm_omni/conversion/conftest.py

@@ -0,0 +1,25 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Path setup for Speech-side EasyMagpie checkpoint conversion tests."""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+EASYMAGPIE_ROOT = Path(__file__).resolve().parents[5] / "tools" / "easymagpie_vllm_omni"
+SCRIPTS_DIR = EASYMAGPIE_ROOT / "scripts"
+
+# Conversion uses Speech and pure PyTorch helpers, never the vLLM runtime.
+sys.path.insert(0, str(EASYMAGPIE_ROOT))
+sys.path.insert(0, str(SCRIPTS_DIR))

tests/collections/tts/easymagpie_vllm_omni/conversion/test_convert_codec.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import convert_codec as converter
+import pytest
+
+
+def _valid_decoder_config() -> dict:
+    return {
+        "_target_": "nemo.collections.tts.modules.audio_codec_modules.ResNetDecoder",
+        "is_causal": True,
+        "activation": "half_snake",
+    }
+
+
+@pytest.mark.parametrize(
+    ("field", "value", "message"),
+    [
+        ("_target_", "some.OtherDecoder", "ResNetDecoder"),
+        ("is_causal", False, "causal"),
+        ("activation", "snake", "half_snake"),
+    ],
+)
+def test_validate_decoder_config_rejects_unsupported_codec(field, value, message):
+    config = _valid_decoder_config()
+    config[field] = value
+
+    with pytest.raises(ValueError, match=message):
+        converter.validate_decoder_config(config)

tests/collections/tts/easymagpie_vllm_omni/conversion/test_convert_to_vllm.py

@@ -0,0 +1,136 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import sys
+import types
+
+import convert_to_vllm as converter  # noqa: E402
+import pytest
+import torch
+from omegaconf import OmegaConf
+
+
+class _FakeEmbeddingModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.anchor = torch.nn.Parameter(torch.zeros(1))
+        self.text_embedding = None
+        self.cfg_unk_token_id = 12
+        self.interruption_token_id = 13
+        self.cfg = types.SimpleNamespace(embedding_dim=2)
+
+    def embed_text_tokens(self, ids, text_lens, disable_cas_embedding):
+        del text_lens, disable_cas_embedding
+        return ids.unsqueeze(-1).expand(-1, -1, self.cfg.embedding_dim).float()
+
+
+def test_precompute_text_embeddings_includes_multiturn_interruption_token():
+    table = converter.precompute_text_embeddings(_FakeEmbeddingModel(), batch_size=8)
+
+    assert table.shape == (14, 2)
+    torch.testing.assert_close(table[-1], torch.tensor([13.0, 13.0]))
+
+
+def test_build_config_exports_multiturn_text_metadata(monkeypatch):
+    class _FakeNemotronHConfig:
+        def __init__(self, **kwargs):
+            self.__dict__.update(kwargs)
+
+    monkeypatch.setitem(
+        sys.modules,
+        "nemo.collections.tts.modules.nemotron_h_decoder",
+        types.SimpleNamespace(NemotronHConfig=_FakeNemotronHConfig),
+    )
+    mode = types.SimpleNamespace(
+        text_input_mode="streaming",
+        streaming_phonemes_delay=3,
+        streaming_speech_delay=5,
+    )
+    model = types.SimpleNamespace(
+        cfg=OmegaConf.create(
+            {
+                "decoder_type": "nemotron_h",
+                "hidden_dim": 4,
+                "embedding_dim": 4,
+                "nemotron_h_config": {"hidden_size": 4},
+                "local_transformer_type": "ar",
+                "use_multiturn_dataset": True,
+            }
+        ),
+        eos_id=101,
+        interruption_token_id=103,
+        num_audio_codebooks=2,
+        codebook_size=32,
+        frame_stacking_factor=1,
+        phoneme_tokenizer=None,
+        mode_name_to_mode={"default": mode},
+        default_inference_mode="default",
+        training_modes=[],
+        task_embedding=None,
+        audio_bos_id=32,
+        audio_eos_id=33,
+        mask_token_id=36,
+    )
+
+    config = converter.build_config(model, vocab_size=104, torch_dtype="float32")
+
+    assert config["text_eos_id"] == 101
+    assert config["text_interruption_id"] == 103
+    assert config["use_multiturn_dataset"] is True
+
+
+def _validation_model():
+    mode = types.SimpleNamespace(
+        text_input_mode="streaming",
+        streaming_phonemes_delay=3,
+        streaming_speech_delay=5,
+    )
+    return types.SimpleNamespace(
+        cfg=OmegaConf.create(
+            {
+                "decoder_type": "nemotron_h",
+                "hidden_dim": 4,
+                "embedding_dim": 4,
+                "nemotron_h_config": {"hidden_size": 4},
+                "local_transformer_type": "ar",
+            }
+        ),
+        mode_name_to_mode={"default": mode},
+        default_inference_mode="default",
+    )
+
+
+@pytest.mark.parametrize(
+    ("field", "value", "message"),
+    [
+        ("decoder_type", "huggingface", "Nemotron-H"),
+        ("local_transformer_type", "none", "local_transformer_type='ar'"),
+        ("hidden_dim", 8, "hidden_dim.*embedding_dim"),
+    ],
+)
+def test_validate_model_config_rejects_unsupported_model(field, value, message):
+    model = _validation_model()
+    model.cfg[field] = value
+
+    with pytest.raises(ValueError, match=message):
+        converter.validate_model_config(model)
+
+
+def test_validate_model_config_rejects_non_streaming_default_mode():
+    model = _validation_model()
+    model.mode_name_to_mode["default"].text_input_mode = "full"
+
+    with pytest.raises(ValueError, match="text_input_mode='streaming'"):
+        converter.validate_model_config(model)

tests/collections/tts/easymagpie_vllm_omni/serving/conftest.py

@@ -0,0 +1,69 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Fixtures and dependency boundary for EasyMagpie vLLM serving tests."""
+from __future__ import annotations
+
+import importlib.util
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+_SERVING_DEPENDENCIES = ("vllm", "vllm_omni")
+_HAS_SERVING_DEPENDENCIES = all(importlib.util.find_spec(name) is not None for name in _SERVING_DEPENDENCIES)
+
+# Ordinary Speech test environments do not install the pinned serving stack.
+# The dedicated serving job provides both dependencies and collects this suite.
+collect_ignore_glob = [] if _HAS_SERVING_DEPENDENCIES else ["test_*.py"]
+
+EASYMAGPIE_ROOT = Path(__file__).resolve().parents[5] / "tools" / "easymagpie_vllm_omni"
+SCRIPTS_DIR = EASYMAGPIE_ROOT / "scripts"
+
+# Load the standalone serving package from this checkout without installing Speech.
+sys.path.insert(0, str(EASYMAGPIE_ROOT))
+sys.path.insert(0, str(SCRIPTS_DIR))
+
+# Equal dimensions exercise the identity projection path.
+_DEFAULT_ARCH: dict = dict(
+    hidden_dim=64,
+    embedding_dim=64,
+    audio_embedding_dim=64,
+    num_audio_codebooks=2,
+    codebook_size=32,
+    frame_stacking_factor=2,
+    local_transformer_n_layers=2,
+    local_transformer_n_heads=4,
+    local_transformer_hidden_dim=64,
+)
+
+
+def build_vllm_config(**arch_overrides):
+    """Build a minimal eager ``VllmConfig`` for the code predictor."""
+    import torch
+    from vllm.config import CompilationMode
+
+    arch = {**_DEFAULT_ARCH, **arch_overrides}
+    hf_config = types.SimpleNamespace(**arch)
+    return types.SimpleNamespace(
+        model_config=types.SimpleNamespace(hf_config=hf_config, dtype=torch.float32),
+        scheduler_config=types.SimpleNamespace(max_num_batched_tokens=128),
+        compilation_config=types.SimpleNamespace(mode=CompilationMode.NONE),
+    )
+
+
+@pytest.fixture
+def vllm_config_factory():
+    """Fixture returning the :func:`build_vllm_config` factory."""
+    return build_vllm_config