Use Fused Kernels for Depthwise Striding Subsampler: up to 1.5x faster training

nemo/collections/asr/modules/conformer_encoder.py

@@ -83,6 +83,9 @@ class ConformerEncoder(NeuralModule, StreamingEncoder, Exportable, AccessMixin):
             Should be power of 2, 1 (auto-chunking, default), or -1 (no chunking)
         subsampling_conv_channels (int): the size of the convolutions in the subsampling module
             Defaults to -1 which would set it to d_model.
+        use_triton (bool, Optional): use the fused Triton subsampling kernels, for CUDA input with
+            'dw_striding' subsampling. Defaults to None, enabling them whenever Triton is
+            installed. Export always uses the PyTorch path.
         reduction (str, Optional): the method of reduction, choices=['pooling', 'striding']. If no value
             is passed, then no reduction is performed and the models runs with the original 4x subsampling.
         reduction_position (int, Optional): the index of the layer to apply reduction. If -1, apply reduction
@@ -345,6 +348,7 @@ def __init__(
         sync_max_audio_length: bool = True,
         rope_base: float = 10000.0,
         rotary_fraction: float = 1.0,
+        use_triton: bool | None = None,
     ):
         super().__init__()
         d_ff = d_model * ff_expansion_factor
@@ -422,6 +426,7 @@ def __init__(
                     subsampling_conv_chunking_factor=subsampling_conv_chunking_factor,
                     activation=nn.ReLU(True),
                     is_causal=causal_downsampling,
+                    use_triton=use_triton,
                 )
         else:
             self.pre_encode = nn.Linear(feat_in, d_model)

nemo/collections/asr/parts/submodules/subsampling.py

@@ -19,6 +19,9 @@
 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
+from nemo.collections.asr.parts.triton.subsampling import fused_conv_relu_dw
+from nemo.core.utils.optional_libs import TRITON_AVAILABLE
 from nemo.utils import logging
 
 
@@ -126,6 +129,7 @@ def __init__(
         subsampling_conv_chunking_factor=1,
         activation=nn.ReLU(),
         is_causal=False,
+        use_triton: bool | None = None,
     ):
         super(ConvSubsampling, self).__init__()
         self._subsampling = subsampling
@@ -418,6 +422,15 @@ def __init__(
 
         self.conv = MaskedConvSequential(*layers)
 
+        # The kernels implement `dw_striding`'s layout, [conv, act] + (sampling_num - 1) x
+        # [dw, pw, act], with ReLU baked in; a factor of 2 stops after [conv, act], leaving no
+        # depthwise to fuse.
+        if use_triton is None:
+            use_triton = TRITON_AVAILABLE
+        self.conv.fuse_triton = (
+            use_triton and subsampling == 'dw_striding' and self._sampling_num >= 2 and isinstance(activation, nn.ReLU)
+        )
+
     def get_sampling_frames(self):
         return [1, self.subsampling_factor]
 
@@ -695,10 +708,33 @@ def calculate_conv_output_size(input_size: torch.Tensor, kernel_size: int, strid
 
 
 class MaskedConvSequential(nn.Sequential):
+    # Set by ConvSubsampling; off by default, so every other subsampling type stays on PyTorch.
+    fuse_triton = False
+
     def forward(self, x, lengths):
         # Convert input (batch, time, features) to conv format
         x = x.unsqueeze(1)  # (batch, 1, time, features)
-        current_lengths = lengths.clone().float()
+        current_lengths = lengths
+
+        # Tracing and export cannot capture a Triton launch, the fused kernel returns no input
+        # gradient, and its weight gradients accumulate through atomics, so their summation order
+        # varies between runs.
+        if (
+            self.fuse_triton
+            and x.is_cuda
+            and not x.requires_grad
+            and not torch.are_deterministic_algorithms_enabled()
+            and not (torch.jit.is_tracing() or torch.compiler.is_exporting())
+        ):
+            x, current_lengths, mask = self._forward_fused(x, current_lengths)
+        else:
+            x, current_lengths, mask = self._forward_torch(x, current_lengths)
+
+        # Final masking
+        x = apply_channel_mask(x, mask)
+        return x, current_lengths.long()
+
+    def _forward_torch(self, x, current_lengths):
         mask = self._create_mask(x, current_lengths.long())
 
         # Process through each layer with mask propagation
@@ -711,21 +747,99 @@ def forward(self, x, lengths):
 
             # Update lengths for stride operations with proper padding
             if hasattr(layer, 'stride') and layer.stride != (1, 1):
-                if hasattr(layer, "_left_padding"):
-                    padding = (layer._left_padding, layer._right_padding)  # CausalConv2D
-                else:
-                    padding = layer.padding
                 current_lengths = calculate_conv_output_size(
-                    current_lengths, layer.kernel_size[0], layer.stride[0], padding
+                    current_lengths, layer.kernel_size[0], layer.stride[0], _layer_padding(layer)
                 )
                 mask = self._create_mask(x, current_lengths.long())
 
-        # Final masking
-        x = apply_channel_mask(x, mask)
-        return x, current_lengths.long()
+        return x, current_lengths, mask
+
+    def _forward_fused(self, x, current_lengths):
+        """The `dw_striding` stack, with conv0 and the depthwise layers as Triton kernels.
+
+        The stack is `[conv, act] + (sampling_num - 1) x [dw, pw, act]`. One kernel covers the
+        leading `conv, act, dw`; the loop over `self[3:]` runs each depthwise as a kernel, each
+        pointwise as a linear, and every other layer as itself. Lengths change only at the
+        depthwise layers.
+
+        Tensors are channels-last throughout, `(batch, time, freq, channels)`, and one permute at
+        the end returns the `(batch, channels, time, freq)` the caller expects.
+
+        The kernels read zeros beyond their input lengths and write zeros beyond their output
+        lengths. Only the trailing pointwise and activation touch the padded tail, which
+        `apply_channel_mask` clears at the end of `forward`.
+        """
+        conv0, _, first_depthwise, first_pointwise, activation = self[:5]
+        # conv -> ReLU -> depthwise in one kernel; the intermediate never reaches memory.
+        x, current_lengths = fused_conv_relu_dw(
+            x,
+            conv0.weight,
+            conv0.bias,
+            first_depthwise.weight,
+            first_depthwise.bias,
+            *_layer_padding(conv0),
+            current_lengths,
+        )
+
+        body = self[5:]
+        x = _pointwise_block(x, first_pointwise, activation)
+        for i in range(0, len(body), 3):
+            depthwise, pointwise, activation = body[i : i + 3]
+            # The kernel masks its own output, so it needs the post-stride lengths.
+            next_lengths = calculate_conv_output_size(
+                current_lengths, depthwise.kernel_size[0], depthwise.stride[0], _layer_padding(depthwise)
+            )
+            x = dw_conv2d(
+                x,
+                depthwise.weight,
+                depthwise.bias,
+                depthwise.stride,
+                *_layer_padding(depthwise),
+                current_lengths,
+                next_lengths,
+            )
+            current_lengths = next_lengths
+            x = _pointwise_block(x, pointwise, activation)
+
+        x = x.permute(0, 3, 1, 2)
+        return x, current_lengths, self._create_mask(x, current_lengths.long())
 
     def _create_mask(self, tensor, lengths):
         """Create mask matching tensor dimensions."""
         batch_size, channels, time, features = tensor.shape
         time_mask = torch.arange(time, device=tensor.device).expand(batch_size, time) < lengths.unsqueeze(1)
         return time_mask.unsqueeze(-1).expand(batch_size, time, features).to(tensor.dtype)
+
+
+def _layer_padding(layer):
+    """The (start, end) padding of a convolution.
+
+    nn.Conv2d's `.padding` is (pad_h, pad_w), one value per axis and symmetric within it, so the
+    height value is both edges. CausalConv2D keeps its two edges on private attributes.
+    """
+    if hasattr(layer, "_left_padding"):
+        return layer._left_padding, layer._right_padding
+    return layer.padding[0], layer.padding[0]
+
+
+def _is_depthwise(layer):
+    """A depthwise convolution: one group per channel."""
+    return isinstance(layer, nn.Conv2d) and layer.groups > 1
+
+
+def _is_pointwise(layer):
+    """A 1x1 convolution over all channels, which is a contraction over the channel axis alone."""
+    return isinstance(layer, nn.Conv2d) and layer.groups == 1 and layer.kernel_size == (1, 1)
+
+
+def _pointwise_block(x, conv, activation):
+    # kernel_size=1 convs are pointwise, i.e. linear, but nn.Conv2d dispatches to much slower
+    # cuBLAS kernels. flatten(1) on the weight is a free view, so checkpoints are unchanged.
+    # F.linear on an N-D input returns a view of its 2D result. An in-place activation on a
+    # view copies the whole tensor in backward, so x is flattened and the activation runs on
+    # the 2D result itself.
+    # TODO: remove the shape manipulation once https://github.com/pytorch/pytorch/pull/194077
+    # is in the minimum required PyTorch version.
+    b, t, f, c = x.shape
+    x = nn.functional.linear(x.view(-1, c), conv.weight.flatten(1), conv.bias)
+    return activation(x).view(b, t, f, -1)

nemo/collections/asr/parts/triton/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  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.