Add EasyMagpie vLLM-Omni serving

[vklimkov-nvidia]

Summary

• Adds EasyMagpieTTS streaming inference on vLLM-Omni for NemotronTTS, using a Nemotron-H autoregressive talker with a per-codebook local transformer over a 25 fps spectral codec.
• Implements two registered deployment topologies:
  • a talker-only pipeline for acoustic-token generation
  • a two-stage pipeline that connects the talker to a native stateful codec and produces 22.05 kHz waveform chunks
• Supports offline synthesis, OpenAI-compatible whole-text HTTP serving through POST /v1/audio/speech, and incremental text/audio streaming through WS /v1/audio/speech/stream.
• Includes checkpoint conversion, vLLM plugin registration, deployment configs, benchmark tools, notebooks, and focused unit tests.

Why

EasyMagpie generates stacked acoustic codes autoregressively and then decodes them into audio. Running both stages in vLLM-Omni provides batched, high-throughput inference while retaining progressive audio output for streaming and voice-agent use cases.

The conversion tool turns the training-time EasyMagpie and causal codec .nemo checkpoints into a self-contained model directory. It converts the model weights, precomputes the text-embedding lookup, bundles the native codec, saves the tokenizer, and optionally stores speaker embeddings. After conversion, serving requires vLLM/vLLM-Omni and this plugin package, but does not require NeMo at runtime.

The native codec keeps its causal history in vLLM-managed state pages, allowing each invocation to process only newly generated acoustic frames and emit asynchronous PCM chunks without replaying the full context.

Validation

• Repository isort and Black checks pass.
• pytest -q examples/tts/easymagpie_vllm_omni/tests: 70 passed, 1 skipped.
• A real EasyMagpie talker and 25 fps codec conversion completed successfully on an RTX A6000.
• Generated artifacts were validated for a consistent safetensors index, tokenizer and architecture metadata, with 956 talker tensors and 76 codec tensors.
• Conversion imports for add_special_tokens, NemotronHConfig, ModelLoadConfig, and load_easy_magpie_model resolve from the Speech checkout.
• The 128-request RTX A6000 reference at concurrency 32 measured:
  • model whole-text input: 67.75x RTF
  • HTTP whole-text service: 48.43x RTF
  • WebSocket incremental service: 44.12x RTF
  • zero playback underruns in both service modes

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[blisc]

/ok to test 0734cee

[blisc]

Overall the PR looks good from my end. The immediate concerns are

• How we package this component? It looks like it is currently constructed to be completely standalone from the rest of NeMo Speech which can be done.
• On packaging, despite that decision, we need tests to run which means that for testing, the requirements need to be added to NeMo's pyproject.toml and tests need to be moved to the NeMo tests folder, and we need to make sure that they run.
• The current implementation is slightly hard-coded which is fine for initial version, but we need much more strict config validation functions that explain what parameters will break the current implementation, and how we can update the implementation to allow future parameters

[blisc]

Let's make a better naming convention. Not sure that easymagpie_talker should be used for the pipeline without the codec. @paarthneekhara, what should we refer to the backbone+LT as?

[vklimkov-nvidia]

@paarthneekhara could you please advise on the canonical name for the backbone + LT stage? I’d prefer to use the established EasyMagpie terminology rather than introduce a new name here.

[paarthneekhara]

Sorry, missed these. How about easymagpie_codecLM or just easymagpie_LM ?

[blisc]

Let's go with EasyMagpieTTS which decomposes into

• EasyMagpie_LM
  • EasyMagpie_LM_Backbone
  • EasyMagpie_LM_LT
• SpectralCodec-BWE-22kHz

[blisc]

Perhaps instead of keeping a nemo-free implementation here, we ought to have a nn.Module base class that both will inherit.

[vklimkov-nvidia]

The serving decoder has a different runtime contract: NeMo uses NeuralModule, length masks, and weight normalization, while vLLM uses packed batches and explicit cache state. Sharing a base from nemo.collections would also make the isolated serving package depend on Speech. That dependency could make sense if we identify enough reusable components, but currently the serving components require substantially different implementations, so a shared base would add coupling without removing much duplication.

[blisc]

How do we ensure that changes to the torch path won't break this initial vllm implementation? Or if it breaks, how can we get notified that vllm changes are needed?
I want to avoid the situation where torch changes doesn't get noticed until the end.
Should we add a test between the two paths in the repo? Or should we start enforcing doing evals with vllm?
It might be out of scope for this PR, but I want to start some sort of discussion on this.
CC @rlangman

[rfejgin]

+1
I agree: I only looked at the codec part, but it looks like the codec implementation is duplicated from torch into the VLLM case. From the comments above, it sounds like there is no way around that. But yes, we should have tests that catch drift between the two, which will inevitably happen as time passes. There are a few things to discuss:

1. How to test that the two implementations produce the same result: we'll need to about which scenarios to check; define some tolerances; think about random seeding; and minimize autoregressiveness in the tests so that unavoidable numeric differences don't accumulate.
2. Does every commit into the torch path require a corresponding change in the vllm implementation? If we don't do that, the tests will break. But then again, it is a bit of a high bar to require porting every (possibly experimental) change to VLLM. So it may be that these tests are only run on demand, not be part of CI.
3. Given the difficulty and error-proneness of keeping two impelmentations in sync over the long term: (a) it's probably worthwhile to give more thought about whether there's any way to rearchitect our code such that less duplication between the torch and vllm exists. (b) failing that, maybe there are ways to automate the porting to vllm since we'll need to do it regularly. The more tests exists, the more we can hand that task off to an agent with confidence.

[github-actions]

[🤖]: Hi @vklimkov-nvidia 👋,

We wanted to let you know that a CICD pipeline for this PR just finished successfully.

So it might be time to merge this PR or get some approvals.

[vklimkov-nvidia]

@blisc I moved the tests under tests/collections/tts: conversion tests now run in the regular Speech suite. The serving tests cannot run there as-is because vLLM 0.24 pins PyTorch 2.11 and would constrain Speech’s bring-your-own-PyTorch environment, so they now run in a path-gated, isolated GPU job. Could you please /ok to test 09a9e5ae33a84502403301b16dce2e9da88ee1cc so the new job runs?

[vklimkov-nvidia]

/ok to test 87779cb

[github-actions]

[🤖]: Hi @vklimkov-nvidia 👋,

We wanted to let you know that a CICD pipeline for this PR just finished successfully.

So it might be time to merge this PR or get some approvals.