HTDemucs 6s — single-graph export for full kernel fusion on TensorRT

Title: HTDemucs 6s — single-graph export for full kernel fusion on TensorRT

Thanks to the TRT team for making inference this straightforward — wanted to share a deployment note that might be useful for your docs.

I got HTDemucs v4 6-stem running at ~5 seconds end-to-end for a 3-minute track on RTX 3090 using TensorRT 10.15.1 SDK. Once the graph was correctly structured, the optimizer did exactly what it's supposed to do — no issues on the TensorRT side at all.

The non-obvious part was the export. HTDemucs has a dual-path architecture where the time-domain and frequency-domain branches cross-communicate through a transformer bottleneck. The standard workaround — externalizing STFT/ISTFT as a second model input — severs that connection and produces numerically incorrect output (pytorch/pytorch #135343). The fix is keeping _spec() inside the trace:

class WaveformOnlyWrapper(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, input_waveform):
        z = self.model._spec(input_waveform)
        return self.model(input_waveform, z)

Three things required simultaneously: legacy tracer (dynamo=False — dynamo can't trace _spec() inside this architecture), opset_version=17 (floor for FFT operator support in ONNX, not a preference), and the wrapper above. With all three in place TensorRT received the complete dual-path graph and achieved full kernel fusion across both branches.

One note on Polygraphy — came across it after the fact. do_constant_folding=True at export time handled what Polygraphy would have done. Wasn't strictly necessary for this model.

Built with TensorRT 10.15.1 SDK (static engine) — the export approach should generalize to TRT-RTX's JIT path as well.

github.com/MansfieldPlumbing/Demucs_v4_TRT
huggingface.co/MansfieldPlumbing/Demucs_v4_TRT

Happy to talk shop if anyone's interested.

The report describes a WaveformOnlyWrapper around _spec(), with legacy tracing (dynamo=False) and ONNX opset 17, but names no TensorRT-RTX source file or test. Start by locating the export entry point that handles this model and determine whether reproducing the complete dual-path graph and full kernel fusion is within this repository's scope.