SIGN IN SIGN UP

Simplify stateless RNG split() / fold_in() kernels with no grid-stride loop (#190036)

Follow-up from #189767 to reduce complexity of `split()` / `fold_in()` kernels. On an H100, this generally results in neutral perf or speedups, with some exceptions:

```
  split()

  ┌─────┬────────────┬─────────────┬────────────┬─────────┐
  │ num │     n      │ before (µs) │ after (µs) │ speedup │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 2   │          1 │        6.63 │       6.51 │  1.02x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 2   │      1,024 │        6.49 │       6.58 │  0.99x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 2   │     16,384 │        6.51 │       6.60 │  0.99x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 2   │     65,536 │        6.50 │       6.68 │  0.97x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 2   │    262,144 │        6.49 │       6.64 │  0.98x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 2   │  1,048,576 │       21.55 │      21.19 │  1.02x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 2   │  4,194,304 │      140.07 │     127.49 │  1.10x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 2   │ 16,777,216 │      533.07 │     492.19 │  1.08x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 16  │          1 │        6.50 │       6.43 │  1.01x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 16  │      1,024 │        6.47 │       6.60 │  0.98x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 16  │     16,384 │        6.55 │       6.52 │  1.00x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 16  │     65,536 │        8.89 │       8.14 │  1.09x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 16  │    262,144 │       40.50 │      33.04 │  1.23x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 16  │  1,048,576 │      183.46 │     182.13 │  1.01x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 16  │  4,194,304 │      953.47 │     975.40 │  0.98x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 128 │          1 │        6.38 │       6.51 │  0.98x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 128 │      1,024 │        6.43 │       6.61 │  0.97x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 128 │     16,384 │       14.32 │      13.95 │  1.03x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 128 │     65,536 │       70.55 │      58.48 │  1.21x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 128 │    262,144 │      299.95 │     241.81 │  1.24x  │
  ├─────┼────────────┼─────────────┼────────────┼─────────┤
  │ 128 │  1,048,576 │     1312.20 │    1419.08 │  0.92x  │
  └─────┴────────────┴─────────────┴────────────┴─────────┘

  fold_in()

  ┌─────────┬────────────┬─────────────┬────────────┬─────────┐
  │ variant │     n      │ before (µs) │ after (µs) │ speedup │
  ├─────────┼────────────┼─────────────┼────────────┼─────────┤
  │ int     │          1 │        7.22 │       6.96 │  1.04x  │
  ├─────────┼────────────┼─────────────┼────────────┼─────────┤
  │ int     │      1,024 │        6.93 │       6.91 │  1.00x  │
  ├─────────┼────────────┼─────────────┼────────────┼─────────┤
  │ int     │     16,384 │        7.01 │       6.99 │  1.00x  │
  ├─────────┼────────────┼─────────────┼────────────┼─────────┤
  │ int     │     65,536 │        6.93 │       6.96 │  1.00x  │
  ├─────────┼────────────┼─────────────┼────────────┼─────────┤
  │ int     │    262,144 │        6.93 │       7.02 │  0.99x  │
  ├─────────┼────────────┼─────────────┼────────────┼─────────┤
  │ int     │  1,048,576 │        7.63 │       7.04 │  1.09x  │
  ├─────────┼────────────┼─────────────┼────────────┼─────────┤
  │ int     │  4,194,304 │       70.33 │      66.66 │  1.06x  │
  ├─────────┼────────────┼─────────────┼────────────┼─────────┤
  │ int     │ 16,777,216 │      272.30 │     249.00 │  1.09x  │
  ├─────────┼────────────┼─────────────┼────────────┼─────────┤
  │ tensor  │          1 │        6.98 │       6.58 │  1.06x  │
  ├─────────┼────────────┼─────────────┼────────────┼─────────┤
  │ tensor  │      1,024 │        6.74 │       6.74 │  1.00x  │
  ├─────────┼────────────┼─────────────┼────────────┼─────────┤
  │ tensor  │     16,384 │        6.79 │       6.64 │  1.02x  │
  ├─────────┼────────────┼─────────────┼────────────┼─────────┤
  │ tensor  │     65,536 │        6.68 │       6.61 │  1.01x  │
  ├─────────┼────────────┼─────────────┼────────────┼─────────┤
  │ tensor  │    262,144 │        6.70 │       6.70 │  1.00x  │
  ├─────────┼────────────┼─────────────┼────────────┼─────────┤
  │ tensor  │  1,048,576 │        7.12 │       6.77 │  1.05x  │
  ├─────────┼────────────┼─────────────┼────────────┼─────────┤
  │ tensor  │  4,194,304 │       70.29 │      66.66 │  1.06x  │
  ├─────────┼────────────┼─────────────┼────────────┼─────────┤
  │ tensor  │ 16,777,216 │      272.57 │     248.86 │  1.10x  │
  └─────────┴────────────┴─────────────┴────────────┴─────────┘
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/190036
Approved by: https://github.com/ngimel
ghstack dependencies: #190253
J
Joel Schlosser committed
96ceb08175be9720b7aa6da80ba2b5b52429bd8a
Parent: 8ba9c56
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> on 8/25/2026, 8:59:46 PM