The online softmax preparation pattern computes softmax max and sum as a
multi-output prim. When this prim appears in an AOTAutograd joint graph, the
min-cut partitioner previously did not know that it was a recomputable/fusible
operation like the decomposed amax/sum pattern it replaces. That skewed the cut
and could make the forward save a full softmax activation instead of the much
smaller max/sum tuple outputs.
Add prims.prepare_softmax_online to the partitioner's recomputable op set. The
existing post-grad placement is still the normal path, but this makes the
partitioner robust if the online softmax preparation is introduced before
partitioning.
I considered relying only on the post-grad pattern placement that originally
avoided the regression, but that leaves the partitioner's cost model wrong for
this prim and makes future joint-graph uses fragile.
Benchmark Results:
A focused AOTAutograd min-cut reproducer that simulates joint-graph online
softmax placement saved 212,992 activation bytes before the fix and 149,504
after the fix, a 29.8% reduction. The old cut saved an extra full
[2, 128, 128] fp16 softmax activation; the new cut saves the small
[2, 128, 1] fp32 max/sum tuple outputs instead.
Test Plan:
- python test/inductor/test_online_softmax.py -k test_prepare_softmax_after_partitioning
- python test/inductor/test_online_softmax.py -k test_prepare_softmax
- python ../state/148141/benchmark_partition_saved_bytes.py
- lintrunner -a torch/_functorch/partitioners.py test/inductor/test_online_softmax.py
- git diff --check --cached
Fixes #148141
Generated by my agent
Pull Request resolved: #185783
Approved by: https://github.com/shunting314