Use Fused Kernels for Depthwise Striding Subsampler: up to 1.5x faster training by MahmoudAshraf97 · Pull Request #16114 · NVIDIA-NeMo/Speech

MahmoudAshraf97 · 2026-08-19T15:40:21Z

Important

The Update branch button must only be pressed in very rare occassions.
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 ?

This PR implements triton kernels for the subsampler in ConformerEncoder, profiling showed that the subsampler is heavily bottlenecked by memory bandwidth, so the fused kernels focus on reducing memory bandwidth by removing intermediary tensors

The first kernel implements:

mel input: [B, T, F, 1]
    ⇩
conv0 output: [B, T/2, F/2, 256] 
    ⇩
ReLU: [B, T/2, F/2, 256] 
    ⇩
Masking: [B, T/2, F/2, 256] 
    ⇩
conv_dw1 output: [B, T/4, F/4, 256]

it cuts memory bandwidth by more than 350x and even more in causal mode since it includes an extra padding layer, this fusion is not possible with torch.compile since it does not fuse the activation, the bias, or the masking into the convolution kernels, this kernel is most of the speedup

The second kernel is a general depthwise conv2d kernel, it's faster than torch kernel in forward since it fuses the pre and post masking and the causal padding, it has the most impact on M and L models, almost no impact on XL models because the encoder layers dominate the runtime, it is also faster than torch backward by 2x for the same reason

below are the speedup numbers expressed in relative speedup in E2E encoder step fwd+bwd, it uses the optimization in #16052 in both arms of the comparison

Eager

size	fwd	bwd	e2e
M	1.38x	1.44x	1.42x
L	1.36x	1.36x	1.36x
XL	1.05x	1.20x	1.14x

torch.compile (max-autotune-no-cudagraphs)

size	fwd	bwd	e2e
M	1.28x	1.63x	1.50x
L	1.22x	1.47x	1.39x
XL	1.11x	1.23x	1.19x

Collection: ASR

Usage

Defaults to enabled if triton is installed, otherwise disabled, can also be controlled using the training config:

model:
  encoder:
    use_triton: false

GitHub Actions CI

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

Trusted PRs run automatically through copy-pr-bot. For an untrusted PR, a maintainer can trigger CI by commenting
/ok to test <head-sha>; repeat this after a new push if the PR remains untrusted.

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?
Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

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?

@pzelasko @nithinraok

copy-pr-bot · 2026-08-19T15:40:25Z

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.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: a972c7eda8

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-08-19T15:47:07Z

        sync_max_audio_length: bool = True,
        rope_base: float = 10000.0,
        rotary_fraction: float = 1.0,
+        use_triton: bool | None = None,


Document the Triton default and YAML opt-out

For every existing dw_striding configuration in an environment with Triton installed, this new option silently changes the default execution path, yet none of the user-facing ASR documentation or example YAML configurations explains that behavior or shows use_triton: false. Add the new default, eligibility constraints, and opt-out to the relevant FastConformer documentation/config examples so users can discover and control the behavior.

AGENTS.md reference: AGENTS.md:L58-L63

Useful? React with 👍 / 👎.

this is similar to the use_pytorch_sdpa option, where it only exists in the module docstring and the configs only, not the documentation

yes please add to docs. we should have added use_pytorch_sdpa info to docs as well.

added the documentation

chatgpt-codex-connector · 2026-08-19T15:47:07Z

+                    mel_freq,
+                    pad_start,
+                )
+                relu_out = tl.dot(conv_taps, feats) + conv_bias[:, None]


Respect the caller's disabled-TF32 setting

When an fp32 caller disables TF32 through torch.backends.cudnn.allow_tf32 = False and torch.backends.cuda.matmul.allow_tf32 = False for numerical accuracy, this tl.dot still uses Triton's default TF32 input precision. Because the Triton path is now enabled automatically, such callers silently receive reduced-precision convolution results despite explicitly disabling TF32; select IEEE input precision or fall back to the PyTorch path when TF32 is disabled.

Useful? React with 👍 / 👎.

Most training is in either mixed precision or half precision, where TF32 has no effect at all, regardless leaving this for the maintainers call

Signed-off-by: MahmoudAshraf97 <hassouna97.ma@gmail.com>

nithinraok

Great work, couple of comments.

nithinraok · 2026-08-24T18:24:13Z

        sync_max_audio_length: bool = True,
        rope_base: float = 10000.0,
        rotary_fraction: float = 1.0,
+        use_triton: bool | None = None,


yes please add to docs. we should have added use_pytorch_sdpa info to docs as well.

nithinraok · 2026-08-24T18:46:22Z

+            use_triton = TRITON_AVAILABLE
+        self.conv.fuse_triton = (
+            use_triton and subsampling == 'dw_striding' and self._sampling_num >= 2 and isinstance(activation, nn.ReLU)
+        )


but if user requests use_triton and striding, then code currently silently skips use of triton. Could you raise warning instead.

added a warning for all non-eligible configs

nithinraok · 2026-08-24T18:51:01Z

 from torch.nn import LayerNorm

 from nemo.collections.asr.parts.submodules.causal_convs import CausalConv1D, CausalConv2D
+from nemo.collections.asr.parts.triton.depthwise_conv import dw_conv2d


wouldn;t this raise error, when triton is not available?

No, since all triton code in that file are guarded with TRITON_AVAILABLE flag, the alternative is to remove the gating and gate this import only

rewritten the files to follow the same pattern used in the ngram triton implementation, the indentation is gone and it raises on the first forward if use_triton=True was explicitly requested

Signed-off-by: MahmoudAshraf97 <hassouna97.ma@gmail.com>

MahmoudAshraf97 requested a review from a team as a code owner August 19, 2026 15:40

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

nemo-automation-bot Bot added the community-request label Aug 19, 2026

chatgpt-codex-connector Bot reviewed Aug 19, 2026

View reviewed changes

pzelasko requested review from artbataev and nithinraok August 19, 2026 16:10

svcnvidia-nemo-ci added the waiting-on-maintainers Waiting on maintainers to respond label Aug 21, 2026

MahmoudAshraf97 added 10 commits August 22, 2026 04:48

initial commit

cf7e23f

Signed-off-by: MahmoudAshraf97 <hassouna97.ma@gmail.com>

simplify

85c29e0

Signed-off-by: MahmoudAshraf97 <hassouna97.ma@gmail.com>

better variable names and documentation

3c48c6b

Signed-off-by: MahmoudAshraf97 <hassouna97.ma@gmail.com>

address review concerns

9c00939

Signed-off-by: MahmoudAshraf97 <hassouna97.ma@gmail.com>

bump minimum torch version for export guard

b5e7fb3

Signed-off-by: MahmoudAshraf97 <hassouna97.ma@gmail.com>

update lock file

79a18db

Signed-off-by: MahmoudAshraf97 <hassouna97.ma@gmail.com>

disable autotune for tests

b82485d

Signed-off-by: MahmoudAshraf97 <hassouna97.ma@gmail.com>

..

d2b7e0a

Signed-off-by: MahmoudAshraf97 <hassouna97.ma@gmail.com>

reduce autotune range

edc8778

Signed-off-by: MahmoudAshraf97 <hassouna97.ma@gmail.com>

honor deterministic path

bbe6e87

Signed-off-by: MahmoudAshraf97 <hassouna97.ma@gmail.com>

MahmoudAshraf97 force-pushed the fast-pre-encoder branch from a972c7e to bbe6e87 Compare August 22, 2026 08:49

nithinraok requested changes Aug 24, 2026

View reviewed changes

svcnvidia-nemo-ci added waiting-on-customer Waiting on the original author to respond and removed waiting-on-maintainers Waiting on maintainers to respond labels Aug 24, 2026

MahmoudAshraf97 added 2 commits August 24, 2026 17:12

remove kernels indentation and warn on unsupported configs

9ff5c11

Signed-off-by: MahmoudAshraf97 <hassouna97.ma@gmail.com>

add documentation

51d1024

Signed-off-by: MahmoudAshraf97 <hassouna97.ma@gmail.com>

MahmoudAshraf97 requested a review from nithinraok August 24, 2026 21:20

svcnvidia-nemo-ci removed the waiting-on-customer Waiting on the original author to respond label Aug 25, 2026

Uh oh!

Conversation

MahmoudAshraf97 commented Aug 19, 2026

What does this PR do ?

Usage

GitHub Actions CI

Before your PR is "Ready for review"

Who can review?

Uh oh!

copy-pr-bot Bot commented Aug 19, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

chatgpt-codex-connector Bot Aug 19, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Aug 19, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nithinraok left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants