[Feature][Performance] Triton backend for GRU / LSTM with intermediate resets

[vmoens]

Summary

Adds recurrent_backend="triton" to GRUModule and LSTMModule. The low-level kernels fuse the time loop for one layer and one direction into a single CUDA launch and apply the is_init reset cheaply inside the loop, avoiding the pack_padded_sequence / pad_packed_sequence round trip that dominates the cuDNN pad backend whenever any reset occurs mid-rollout. Sits alongside the existing pad and scan backends.

• New file: torchrl/modules/tensordict_module/_rnn_triton.py — forward and backward kernels, both K-tiled along the recurrent contraction.
• Optional dep, guarded by _has_triton. Clear RuntimeError if the backend is selected without triton installed.
• recurrent_compute_dtype selects fp32 (TF32 on Ampere/Hopper, default — matches cuDNN behavior) or bf16.
• Hidden size is internally padded to the next power of two (Python-side; no in-kernel masking).
• num_layers > 1 is handled by stacking per-layer Triton calls.
• Dropout is applied between per-layer Triton calls, matching PyTorch's recurrent dropout semantics.
• LSTM backward now propagates gradients flowing through intermediate cell states (dc_t).

Limitations (prototype)

• The low-level kernels operate on one layer and one direction at a time; multilayer execution launches one Triton call per layer.
• Bidirectional and LSTM proj_size are not native Triton kernel paths yet. These configurations currently use a pad-compatible correctness path rather than getting the Triton speedup.
• The autograd wrapper saves per-layer gate activations explicitly, so activation memory scales linearly with layer count and can exceed cuDNN's opaque reserve_space footprint.
• Shape changes can trigger separate Triton compilation/autotune. Fixed training shapes are the intended fast path.
• Backward uses tl.atomic_add for reset-state gradients, so deterministic mode is not implemented yet.

Numerical correctness

Forward + backward match the pad backend within TF32 noise (~5e-4 rel err for fp32, ~5e-3 for bf16) across multiple shape / reset_prob combinations. The test coverage now includes multilayer GRU/LSTM, dropout train/eval behavior, bidirectional/projection correctness fallback, and LSTM losses that backpropagate through ("next", "hidden1").

Performance (H100, B=2048, T=256, H=128, fp32, fwd+bwd)

reset_prob	cuDNN (pad)	scan_compile	triton
0%	33.8 ms	70.9 ms	11.7 ms
1%	309.1 ms	71.1 ms	11.7 ms
10%	310.6 ms	71.2 ms	11.7 ms

Triton is reset-invariant and 5–25x faster than cuDNN at non-zero reset probabilities. cuDNN still wins at reset=0% for LSTM at H>=256 (its persistent-LSTM path is hard to beat without TMA + warp specialisation). The gap is small for GRU.

Test plan

• test/test_tensordictmodules.py::TestLSTMModule::test_lstm_module_triton_backend_matches_pad — forward agreement vs pad at H in {16, 64} x {fp32, bf16}, with several mid-rollout resets.
• test/test_tensordictmodules.py::TestLSTMModule::test_lstm_module_triton_backward — gradient agreement vs pad at H=64.
• test/test_tensordictmodules.py::TestLSTMModule::test_lstm_module_triton_extended_forward_matches_pad — multilayer, dropout train/eval, bidirectional fallback, projection fallback, and combined projected bidirectional LSTM forward parity.
• test/test_tensordictmodules.py::TestLSTMModule::test_lstm_module_triton_extended_backward_matches_pad — multilayer, bidirectional fallback, projection fallback, and intermediate cell-state gradient parity.
• Same coverage for GRUModule where applicable.
• All tests skip cleanly when triton is unavailable or CUDA is missing.
• benchmarks/bench_gru_reset_backends.py extended with triton_td, num_layers, dropout sweeps, and optional torch.compile flags.

Follow-ups (intentionally out of scope here)

• Native Triton bidirectional and LSTM proj_size kernels.
• torch.compile first-class integration via torch.library.custom_op (right now the kernel runs inside compiled graphs but as an opaque node).
• vmap support for _GRUFn.apply / _LSTMFn.apply.
• CPU fallback or explicit CPU tensor error.
• In-kernel masking for non-power-of-two H to eliminate the wrapper pad/unpad copies.
• Deterministic backward mode for reset-state gradient accumulation.
• Docs/tutorial/SOTA integration once the prototype scope is accepted.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3738

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 2 Unrelated Failures

As of commit 14e2a61 with merge base 1f1f8bf ():

NEW FAILURE - The following job has failed:

• Libs Tests on Linux / unittests-gym (3.10, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 6f36efb526a3efdee028e21a100438acc3859a74ea80a9f6ecafd117a694f494 /exec failed with exit code 1

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• Unit-tests on Linux / tests-optdeps (3.12, 13.0) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

BROKEN TRUNK - The following job failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Unit-tests on Linux / tests-olddeps (3.10, 11.8) / linux-job (gh) (trunk failure)
  test/test_tensordictmodules.py::TestGRUModule::test_gru_parallel_within[False-False-False]

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vmoens]

Remaining work before this is release-ready

Tracking what would need to land before flipping the triton backend from prototype to first-class:

Correctness gaps

• num_layers > 1 — supported by stacking per-layer Triton calls.
• Dropout — applied between per-layer Triton calls, matching PyTorch's layer-dropout semantics.
• Bidirectional and proj_size — not supported by the Triton kernels yet; current correctness path falls back to pad-compatible execution.
• LSTM intermediate dc_t — the backward kernel now propagates gradients flowing through intermediate cell states.

Production-grade integration

• torch.compile first-class support — wrap as @torch.library.custom_op + register_fake + register_autograd. Today the kernel runs in compiled graphs but as an opaque node (no fusion with surrounding ops).
• vmap — _GRUFn.apply / _LSTMFn.apply aren't vmap-aware.
• CPU fallback — currently CUDA-only; need either a vanilla torch implementation or a hard error on CPU tensors.
• In-kernel H masking — eliminate the wrapper-side pad/unpad copies.
• Determinism — backward uses tl.atomic_add for dh0 / dc0, which can give tiny run-to-run gradient differences. Either expose a deterministic flag that switches to a single-pass reduction or document the non-determinism.

Coverage

• Sota-implementations CI — add a recurrent SAC/PPO config that exercises recurrent_backend="triton" end-to-end (per CLAUDE.md §11).
• Docs — extend docs/source/reference/modules.rst with a paragraph on the new backend.
• Tutorial — short page contrasting pad / scan / triton on the resets-during-RL workload, with the benchmark table (per CLAUDE.md §9).
• Wider numerical test coverage — current tests cover H ∈ {16, 64}. Extend to H ∈ {128, 256} and add a property-style test on randomized is_init masks.

[vmoens]

Follow-up PR:

• Native Triton bidirectional / proj_size
• torch.library.custom_op compile integration
• vmap
• in-kernel H masking
• deterministic backward mode
• full docs/tutorial/SOTA coverage