Adding MoE, PEE, GGEMM then SALM, vLLM related changes by tango4j · Pull Request #16038 · NVIDIA-NeMo/Speech

tango4j · 2026-08-06T00:30:05Z

Important

The Update branch button must only be pressed in very rare occasions.
An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.

What does this PR do?

Adds grouped-GEMM Mixture-of-Experts (MoE) and Parallel Expert Encoder (PEE) support to ASR, and integrates PEE with SpeechLM2 SALM training and generation.

Collection: ASR, SpeechLM2

Changelog

Add GGEMMTransformerEncoder for packed execution of heterogeneous transformer experts, including grouped attention and feed-forward paths with serial reference implementations for correctness checks.
Add MoETransformerEncoder with dense and top-k expert routing plus activation-checkpointing support.
Add ParallelExpertEncoder, combining speech MoE, speaker Sortformer, and sound experts behind a Conformer-compatible interface.
Use memory-efficient per-expert execution during PEE training and packed grouped-GEMM execution during evaluation and inference.
Add long-form PEE inference with windowing, context handling, and streaming speaker-cache support.
Integrate PEE bundles into SALM through model.pe_encoder_path, including external speaker-target routing, speaker-target lengths through context-parallel/chunked batches, and generation-only online inference.
Update SALM vLLM audio handling for PEE-compatible speaker inputs and lengths.
Add focused MoE, GGEMM, PEE, SALM, encoder-chunking, and vLLM tests.

Usage

Configure a PEE bundle as the SALM perception encoder:

from omegaconf import OmegaConf

from nemo.collections.speechlm2.models import SALMAutomodel

cfg = OmegaConf.load("path/to/salm_config.yaml")
cfg.model.pe_encoder_path = "/path/to/parallel_expert_encoder.nemo"

model = SALMAutomodel(cfg.model)
model.configure_model()

PEE automatically uses its memory-efficient per-expert path while training and its packed GGEMM path during evaluation and inference.

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines (author to confirm)
Did you write any new necessary tests?
Did you add or update any necessary documentation? (API documentation and usage example included)
Does the PR affect components that are optional to install? No new optional dependencies are introduced.
- Reviewer: Does the PR have correct import guards for all optional libraries? N/A

PR Type:

New Feature
Bugfix
Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

Local validation: 249 tests passed and 55 tests were skipped as expected.
Formatting, import-order, Python compilation, and git diff --check validation passed.
Related to: N/A

Signed-off-by: Taejin Park <tango4j@gmail.com>

copy-pr-bot · 2026-08-06T00:30:08Z

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.

Signed-off-by: Taejin Park <tango4j@gmail.com>

Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> Signed-off-by: Taejin Park <tango4j@gmail.com>

…ith 'import' and 'import from'' Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> Signed-off-by: Taejin Park <tango4j@gmail.com>

tango4j · 2026-08-06T09:34:21Z

/ok to test 025605c

pzelasko · 2026-08-06T12:46:19Z

@@ -0,0 +1,825 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.


This file should live in tests/collections/asr right?

pzelasko · 2026-08-06T15:30:07Z

+        # per-encoder (see ParallelExpertEncoder.set_activation_checkpointing). Kept as a
+        # flag rather than a checkpoint_wrapper around each layer so module structure and
+        # state_dict keys stay identical to a checkpoint saved without it.
+        self.activation_checkpointing = False


Is this change necessary? Our training worked so far with AC and without this workaround.

All the changes in Transformer Encoder are removed. Transformer encoder file is untouched.

…mport Signed-off-by: Taejin Park <tango4j@gmail.com>

Signed-off-by: Taejin Park <tango4j@gmail.com>

tango4j · 2026-08-07T05:34:10Z

/ok to test dab3466

…ith 'import' and 'import from'' Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> Signed-off-by: Taejin Park <tango4j@gmail.com>

…mport Signed-off-by: Taejin Park <tango4j@gmail.com>

Signed-off-by: Taejin Park <tango4j@gmail.com>

pzelasko · 2026-08-11T14:33:10Z

/ok to test 23028a3

github-actions · 2026-08-11T16:02:11Z

[🤖]: Hi @tango4j 👋,

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.

Signed-off-by: Taejin Park <tango4j@gmail.com>

tango4j · 2026-08-11T23:45:29Z

/ok to test aa368b3

github-actions · 2026-08-12T04:40:06Z

[🤖]: Hi @tango4j 👋,

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.

Adding MoE, PEE, GGEMM then SALM, vLLM related changes

e2d47d7

Signed-off-by: Taejin Park <tango4j@gmail.com>

github-actions Bot added the ASR label Aug 6, 2026

github-advanced-security AI found potential problems Aug 6, 2026

View reviewed changes

Fix MoE QA and lint findings

37e7134

Signed-off-by: Taejin Park <tango4j@gmail.com>

tango4j marked this pull request as ready for review August 6, 2026 07:46

Revised docstrings that makes GGEMM generalizable

2ad4c2e

Signed-off-by: Taejin Park <tango4j@gmail.com>

github-advanced-security AI found potential problems Aug 6, 2026

View reviewed changes

Comment thread nemo/collections/asr/parts/utils/sot_speaker_alignment.py Fixed

Comment thread tests/collections/asr/utils/test_sot_speaker_alignment.py Fixed

tango4j and others added 2 commits August 6, 2026 02:33

Potential fix for pull request finding 'CodeQL / Unused global variable'

97f9e1f

Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> Signed-off-by: Taejin Park <tango4j@gmail.com>

Potential fix for pull request finding 'CodeQL / Module is imported w…

025605c

…ith 'import' and 'import from'' Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> Signed-off-by: Taejin Park <tango4j@gmail.com>

tango4j assigned pzelasko Aug 6, 2026

copy-pr-bot Bot temporarily deployed to public August 6, 2026 09:35 Inactive

tango4j assigned KunalDhawan Aug 6, 2026

copy-pr-bot Bot temporarily deployed to public August 6, 2026 09:38 Inactive

copy-pr-bot Bot temporarily deployed to test August 6, 2026 09:39 Inactive

copy-pr-bot Bot temporarily deployed to public August 6, 2026 09:39 Inactive

pzelasko reviewed Aug 6, 2026

View reviewed changes

tango4j requested a review from ipmedenn August 6, 2026 23:53

ipmedenn reviewed Aug 6, 2026

View reviewed changes

Comment thread nemo/collections/asr/modules/transformer_encoder.py Outdated

tango4j added 3 commits August 6, 2026 20:23

Reverted Transformer Encoder to original state, fixed SOT alignment i…

f1a3e91

…mport Signed-off-by: Taejin Park <tango4j@gmail.com>

Fixing failing NRT PEE-based jobs for text-only Omni dataset

509b97c

Signed-off-by: Taejin Park <tango4j@gmail.com>

Fix lint errors

dab3466

Signed-off-by: Taejin Park <tango4j@gmail.com>

copy-pr-bot Bot temporarily deployed to public August 7, 2026 05:35 Inactive

copy-pr-bot Bot temporarily deployed to test August 7, 2026 05:36 Inactive

copy-pr-bot Bot temporarily deployed to public August 7, 2026 05:38 Inactive

copy-pr-bot Bot temporarily deployed to public August 7, 2026 05:39 Inactive

tango4j and others added 7 commits August 11, 2026 08:56

Potential fix for pull request finding 'CodeQL / Module is imported w…

296f01f

…ith 'import' and 'import from'' Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> Signed-off-by: Taejin Park <tango4j@gmail.com>

Reverted Transformer Encoder to original state, fixed SOT alignment i…

893d349

…mport Signed-off-by: Taejin Park <tango4j@gmail.com>

Fixing failing NRT PEE-based jobs for text-only Omni dataset

8bf18f0

Signed-off-by: Taejin Park <tango4j@gmail.com>

Fix lint errors

817bbb4

Signed-off-by: Taejin Park <tango4j@gmail.com>

Fix MoE CPU gradient tests

6627b96

Signed-off-by: Taejin Park <tango4j@gmail.com>

Revised all the docstrings and comments to make them concise

f8afc2d

Signed-off-by: Taejin Park <tango4j@gmail.com>

Fixed black and isort

23028a3

Signed-off-by: Taejin Park <tango4j@gmail.com>

pzelasko force-pushed the moe_pee_ggemm branch from 7033beb to 23028a3 Compare August 11, 2026 12:56

copy-pr-bot Bot temporarily deployed to public August 11, 2026 14:34 Inactive

pzelasko previously approved these changes Aug 11, 2026

View reviewed changes

copy-pr-bot Bot temporarily deployed to public August 11, 2026 14:37 Inactive

copy-pr-bot Bot temporarily deployed to public August 11, 2026 14:38 Inactive

copy-pr-bot Bot deployed to test August 11, 2026 14:40 Active

copy-pr-bot Bot temporarily deployed to public August 11, 2026 15:00 Inactive

tango4j added 2 commits August 11, 2026 15:57

Adding safe-guards to freeze non-ASR encoder during training

58b923d

Signed-off-by: Taejin Park <tango4j@gmail.com>

Adding safe-guards to freeze non-ASR encoder during training

d665c44

Signed-off-by: Taejin Park <tango4j@gmail.com>

tango4j dismissed pzelasko’s stale review via d665c44 August 11, 2026 23:08

Adding missed changes on parallel expert encoder

2fb5e1a

Signed-off-by: Taejin Park <tango4j@gmail.com>

github-advanced-security AI found potential problems Aug 11, 2026

View reviewed changes

Comment thread tests/collections/asr/test_parallel_expert_encoder.py Fixed

isort and black issue resolved

aa368b3

Signed-off-by: Taejin Park <tango4j@gmail.com>

copy-pr-bot Bot deployed to public August 11, 2026 23:46 Active

copy-pr-bot Bot deployed to test August 11, 2026 23:47 Active

copy-pr-bot Bot deployed to public August 11, 2026 23:50 Active

copy-pr-bot Bot deployed to public August 11, 2026 23:51 Active

		@@ -0,0 +1,825 @@
		# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.

Uh oh!

Conversation

tango4j commented Aug 6, 2026

What does this PR do?

Changelog

Usage

GitHub Actions CI

Before your PR is "Ready for review"

Who can review?

Additional Information

Uh oh!

copy-pr-bot Bot commented Aug 6, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tango4j commented Aug 6, 2026

Uh oh!

pzelasko Aug 6, 2026

Choose a reason for hiding this comment

Uh oh!

tango4j Aug 7, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

pzelasko Aug 6, 2026

Choose a reason for hiding this comment

Uh oh!

tango4j Aug 7, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tango4j commented Aug 7, 2026

Uh oh!

pzelasko commented Aug 11, 2026

Uh oh!

github-actions Bot commented Aug 11, 2026

Uh oh!

Uh oh!

tango4j commented Aug 11, 2026

Uh oh!

github-actions Bot commented Aug 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants