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