I found FlashSAC very interesting, and I tried to speed it up further using Triton!
Motivation
Pytorch profiling on A100 identified categorical TD-target construction as launch-bound:
- Existing PyTorch path: 2 compiled regions and 5 CUDA kernel launches
- Proposed Triton path: 1 fused CUDA kernel
The kernel combines minimum-Q selection, Bellman support mapping, clamping, exponentiation, and deterministic categorical projection.
Results
A100 integration benchmark with Torch 2.9 / Triton 3.5:
| Batch size |
PyTorch |
Triton |
Speedup |
Saved per update |
| 512 |
0.315 ms |
0.093 ms |
3.39× |
0.222 ms |
| 2048 |
0.308 ms |
0.097 ms |
3.16× |
0.211 ms |
I also completed 12 end-to-end training runs: 6 paired PyTorch/Triton comparisons across three seeds on HalfCheetah and PickSingleYCB.
- Triton was faster in 4 of 6 comparisons
- Mean wall-clock saving: 3.23%
- Median wall-clock saving: 2.39%
- Observed range: -2.26% to +9.87%
Numerical validation across random, strided, and heavily clamped inputs showed:
- Maximum absolute error below
3e-6
- Probability-mass error below
4.8e-7
- NaN propagation matching the existing PyTorch path
- Deterministic accumulation without atomics
Integration
The fast path would be explicitly enabled with:
--overrides agent.categorical_target_backend=triton
PyTorch remains the default. The original implementation stays unchanged and is automatically used for unsupported inputs, dependency versions, GPUs, or any Triton runtime failure.
The optimization is restricted to the measured FP32, 101-bin, Ampere-or-newer, Torch 2.9 / Triton 3.5 configuration. The Torch 2.5 / Triton 3.1 configuration falls back because it did not measure faster.
The proposed patch would add only minimal option wiring and a private _triton/ directory, without dependency, CI, or lockfile changes.
If this direction looks worthwhile, I would love to receive feedback and test on broader tasks.
I found FlashSAC very interesting, and I tried to speed it up further using Triton!
Motivation
Pytorch profiling on A100 identified categorical TD-target construction as launch-bound:
The kernel combines minimum-Q selection, Bellman support mapping, clamping, exponentiation, and deterministic categorical projection.
Results
A100 integration benchmark with Torch 2.9 / Triton 3.5:
I also completed 12 end-to-end training runs: 6 paired PyTorch/Triton comparisons across three seeds on HalfCheetah and PickSingleYCB.
Numerical validation across random, strided, and heavily clamped inputs showed:
3e-64.8e-7Integration
The fast path would be explicitly enabled with:
PyTorch remains the default. The original implementation stays unchanged and is automatically used for unsupported inputs, dependency versions, GPUs, or any Triton runtime failure.
The optimization is restricted to the measured FP32, 101-bin, Ampere-or-newer, Torch 2.9 / Triton 3.5 configuration. The Torch 2.5 / Triton 3.1 configuration falls back because it did not measure faster.
The proposed patch would add only minimal option wiring and a private
_triton/directory, without dependency, CI, or lockfile changes.If this direction looks worthwhile, I would love to receive feedback and test on broader tasks.