associative_scan lowering

[jloftin-nv]

Description

higher_order.associative_scan (CUDA combine_mode="pointwise") survives export and has no converter, breaking the graph on classical selective-scan SSM layers.

Add a post-lowering pass that matches the Mamba affine combine (a_la_r, a_rb_l+b_r) and, for static scan length, replaces the HOP with a Hillis–Steele scan of existing aten ops. Unrecognized combines and dynamic S are left alone.

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

• New feature (non-breaking change which adds functionality)

Checklist:

• [ X] My code follows the style guidelines of this project (You can use the linters)
• [ X] I have performed a self-review of my own code
• [ X] I have commented my code, particularly in hard-to-understand areas and hacks
• [ X] I have made corresponding changes to the documentation
• [X ] I have added tests to verify my fix or my feature
• [X ] New and existing unit tests pass locally with my changes
• [ X] I have added the relevant labels to my PR in so that relevant reviewers are notified

[micwill755]

Requesting changes on the identity pad (ones_like / zeros_like after decompositions won’t convert) and on the TRT numeric test, which can still pass via PyTorch fallback.

[micwill755]

ones_like and zeros_like are inserted after run_decompositions, and Dynamo has no converter for them. The identity pad can run in PyTorch even if the mul/add scan converts, so tests can pass without the scan in TRT.

You already slice a_head/b_head and a_prefix/b_prefix. Use those directly instead of padding with dummy 1s and 0s:

keep the head (a_head, b_head)
combine the tail (prefix with a/b from step onward)
cat(head, combined_tail)

If you prefer to keep the pad, emit aten.full.default rather than ones_like / zeros_like.

[micwill755]

This can still pass with the PyTorch fallback from the ones_like / zeros_like comment: min_block_size=1 will compile the convertible pieces and leave those ops in PyTorch. I would add an assert to check the lowered graph has no HOP and no ones_like / zeros_like, and TRT-compile a non-power-of-two S.