[Performance] LSTM/GRU scan: canonical strides + cuDNN flat-storage clones + thread-local recurrent mode

[vmoens]

Stack from ghstack (oldest at bottom):

• [Doc] Migrate shifted=True callers to legacy/compact + docstring polish #3776
• [Example] Expose compact GAE cat dimension #3775
• [Example] Add Isaac RNN PPO rollout mode flags #3769
• [Example] Isaac Lab RNN PPO with compact memory + knowledge-base notes #3763
• [BugFix] Fix GAE compact path bias on recurrent value nets at internal truncations #3771
• [Refactor] Simplify LSTM/GRUModule recurrent-mode shape normalization #3768
• [Refactor] Keep [B, T] dim in value estimators #3767
• [Feature] Collector.fake_tensordict() / MultiCollector.fake_tensordict() #3764
• [Feature] MultiCollector: track policy version per fresh continue command #3759
• [BugFix] GAE shifted=True: tolerate missing next obs, non-canonical strides, docs #3757
• [Feature] timeit.mark_start/mark_end for non-context-manager timing #3762
• [BugFix] PolicyVersion: int64 dtype + preserve tensordict device #3755
• -> [Performance] LSTM/GRU scan: canonical strides + cuDNN flat-storage clones + thread-local recurrent mode #3754
• [BugFix] Recurrent policy auto-register with policy_factory #3753

Three intertwined fixes to the scan / triton recurrent backends.

• Canonical-stride check. A [1, 4, 5] tensor with strides (5, 5, 1)
  passes is_contiguous() but torch._higher_order_ops.scan and the
  triton kernels read strides directly and reject non-canonical layouts.
  Add _canonical_stride + _canonical_contiguous and re-materialize
  inputs / hidden buffers when strides drift off the C-canonical layout.

• cuDNN flat-storage aliasing. nn.LSTM / nn.GRU with cuDNN flatten
  all per-layer parameters into a single storage; the scan HOP tracer
  walks the FakeTensor graph and rejects the aliased per-layer views as
  inputs. Clone the weight views before closing the scan body. The
  per-layer carry now also clones x_t and the transpose+flatten output
  (the only remaining aliasing edge) so the existing .clone() on the
  full torch.stack(...) carry can drop.

• Thread-local recurrent_mode. _ContextManager was a single mutable
  module-level flag, so spawning a collector worker thread saw the
  parent's recurrent_mode setting. Wrap in _RecurrentModeContextManager
  using contextvars.ContextVar so per-thread state is isolated.

Tests cover: _canonical_contiguous no-op vs re-materialize, scan/triton
output parity under non-canonical strides, scan-under-torch.compile
aliasing pin for both LSTM and GRU, LSTM scan/pad PPO-advantage parity,
non-canonical hidden-buffer regression for scan, and a thread-local
set_recurrent_mode test.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Run pull request jobs on OSDC runners in shadow mode

❌ 3 New Failures, 1 Cancelled Job

As of commit 244f61c with merge base 0a01ee8 ():

NEW FAILURES - The following jobs have failed:

• Build Windows Wheels / pytorch/rl / build-wheel-py3_10-cpu (gh)
  Process completed with exit code 1.
• Build Windows Wheels / pytorch/rl / upload / upload-wheel-py3_10-cpu (gh)
  Unable to download artifact(s): Artifact not found for name: pytorch_rl__3.10_cpu_x64
• Unit-tests on Linux / tests-stable-gpu (3.12, 13.0, 2) / linux-job (gh)
  test/test_collectors.py::TestUpdateParams::test_param_sync_mixed_device[SharedMemWeightSyncScheme-cpu-env_device0-False-MultiAsyncCollector]

CANCELLED JOB - The following job was cancelled. Please retry:

• Unit-tests on Windows / unittests-cpu (3.10, windows.4xlarge, cpu) / windows-job (gh)

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

[vmoens]

what if we were doing that for hidden0_in before the loop? Would the view still be contiguous?

[vmoens]

ditto, should we do it before the loop?