Ye Tao 49ab40481e [dynamo][decomp][CPU] Preserve native CPU flash-attention output layout (#199160) ...

problem stated in https://github.com/pytorch/pytorch/issues/199155

Native CPU flash attention allocates its output with empty_like(query), but its Python decomposition previously always packed the result in (L, N, H, E) order. Export can elide contiguous() based on the native layout, leaving an invalid downstream view when a batch-first attention projection is decomposed.

Restore the query’s physical dimension order using the existing symbolic stride-order helper and a contiguous clone in physical order. Skip conversion when output strides are already known to match the query.

Exact native stride preservation for singleton and empty dimensions remains a TODO requiring additional handling and tests.

Add native/FakeTensor stride and gradient coverage, static/dynamic MHA and BERT projection exports, and a regression test for a data-dependent symbolic batch size created by boolean indexing. The new regression verifies that export, decomposition and the downstream view succeed and produce results matching eager execution

Pull Request resolved: https://github.com/pytorch/pytorch/pull/199160
Approved by: https://github.com/aorenste

2026-10-28 09:51:27 +00:00

..

__init__.py

Make the softshrink decompositions match eager bitwise (#200802)

2026-10-19 05:45:33 +00:00

decompositions_for_jvp.py

[Feature Refactoring][decomp] Make log_sigmoid_forward's buffer shape device-agnostic instead of backend-enumerated (#197530)

2026-10-15 01:51:56 +00:00

decompositions_for_rng.py

[Inductor] Add decomposition for prims.uniform (#187887)

2026-07-10 19:16:35 +00:00

decompositions.py

[dynamo][decomp][CPU] Preserve native CPU flash-attention output layout (#199160)

2026-10-28 09:51:27 +00:00