trunk/9455d3d34559799d09196341ca545d946610eb2e: Teach min-cut about online softmax prep (#185783)

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

Approved by: https://github.com/shunting314