drisspg ae02d663cd [FlexAttention] Fix output layout (#135882) ...

We previously only supported the same v_head dim and + qk_head dim. When allowed for different head-dims I accidently kept the same query strides for the output. This PR fixes this bug as well it ensures that we always produce output in the same stride order as the input query.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/135882
Approved by: https://github.com/yanboliang, https://github.com/Chillee

2024-09-13 16:36:05 +00:00

..

__init__.py

[HOO] add hints_wrapper to support passing context hints (#132860)

2024-08-26 18:21:22 +00:00

associative_scan.py

Implementation of scan (#134102)

2024-09-10 04:51:16 +00:00

auto_functionalize.py

Track base of FunctionalTensor in inference mode. (#135141)

2024-09-06 00:10:25 +00:00

cond.py

Revert "[Dynamo] Use custom backend to reenter metadata tf mode when tracing while/cond (#134732)"

2024-09-13 12:52:58 +00:00

effects.py

[aotd] Make effects op registry WeakKeyDictionary (#134470)

2024-08-28 12:12:00 +00:00

executorch_call_delegate.py

[hop] require hops to override __call__. (#134352)

2024-08-28 19:56:40 +00:00

flex_attention.py

[FlexAttention] Fix output layout (#135882)

2024-09-13 16:36:05 +00:00

hints_wrap.py

[HOO] add hints_wrapper to support passing context hints (#132860)

2024-08-26 18:21:22 +00:00

map.py

[hop] preserve metadata in re-tracing hop subgraph by running with interpreter (#135159)

2024-09-05 21:36:56 +00:00

out_dtype.py

Make the __module__ name of HOO to be always "torch.ops.higher_order" (#132775)

2024-08-08 16:55:09 +00:00

run_const_graph.py

[hop] require hops to override __call__. (#134352)

2024-08-28 19:56:40 +00:00

scan.py

Revert "[Dynamo] Use custom backend to reenter metadata tf mode when tracing while/cond (#134732)"

2024-09-13 12:52:58 +00:00

strict_mode.py

[hop] require hops to override __call__. (#134352)

2024-08-28 19:56:40 +00:00

torchbind.py

[hop] require hops to override __call__. (#134352)

2024-08-28 19:56:40 +00:00

triton_kernel_wrap.py

Add option to tweak inductor stride settings for user-defined triton kernels (#135530)

2024-09-11 00:11:17 +00:00

utils.py

Implementation of scan (#134102)

2024-09-10 04:51:16 +00:00

while_loop.py

Revert "[Dynamo] Use custom backend to reenter metadata tf mode when tracing while/cond (#134732)"

2024-09-13 12:52:58 +00:00

wrap.py

[export][reland] Convert autocast to HOO (#132677)

2024-08-05 22:34:52 +00:00