[lanluo-nvidia]

Description

#4921 introduced the following test failure:
https://github.com/pytorch/TensorRT/actions/runs/40760002261/job/123120801180

Fixes # (issue)

Type of change

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

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Checklist:

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

[shoumikhin]

The pass fix works: on torch 2.15 the backend now logs "Folded 2" where it found nothing before. Two of the new test cases fail on torch 2.16 though, see inline. Tested on A100.

[shoumikhin]

This lookup returns an empty dict. matmul has no entry in the decomposition table, so the update adds nothing.

On torch 2.15 that does not matter, because AOT tracing already turns matmul into bmm. On the 2.16 nightly that main uses, aot_export_joint_simple keeps aten.matmul. So both joint_export cases fail with 0 != 2 on the bmm count, before the pass runs. The dynamo-lowering job on this PR already fails on exactly these two cases.

This works on both versions:

from torch._export.utils import _get_decomp_for_cia

decompositions[torch.ops.aten.matmul.default] = _get_decomp_for_cia(
    torch.ops.aten.matmul.default
)

With it, all four new cases pass with your fix and fail without it, on torch 2.15 and 2.16.

[lanluo-nvidia]

Agreed—the lookup was empty. Replaced it with _get_decomp_for_cia(aten.matmul.default) so the joint-export tests explicitly exercise the bmm form.”

[shoumikhin]

Same no-op lookup as above. Here it can simply be removed: make_fx already splits matmul into bmm on both versions.

[lanluo-nvidia]

Removed the no-op lookup; make_fx uses the existing decomposition table directly.

[shoumikhin]

This reads as if every torch.compile graph is covered now. One case is still not: scaled_dot_product_attention(..., enable_gqa=True) with one KV head and decompose_attention=True. In the bmm form the KV head is repeated by the matmul's own broadcasting expand, and _unflatten_bmm_operand skips any expand that changes the shape. So nothing folds on torch 2.15 with plain torch.compile, or on 2.16 with use_distributed_mode_trace. The output is still correct. Only the speedup is missing.

The simple fix is to say here that the bmm form is matched for repeat_kv style graphs. Matching a broadcast from a single KV head would also work, but that can be a follow-up. I don't think this needs to block the merge.

[lanluo-nvidia]

Narrowed the wording to repeat_kv style graphs in the guide and API docstrings, and documented the remaining enable_gqa=True broadcasting limitation. Leaving that optimization for a follow-up.”

[shoumikhin]

Thanks, all three points are addressed. Looks good to me.