Add flash attention impl to flex attention

[drisspg]

Stack from ghstack (oldest at bottom):

• -> Add flash attention impl to flex attention #160109
• Add cutedsl template support to compile #160108

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/160109

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 35 New Failures

As of commit 5fe2e74 with merge base 74871d4 ():

NEW FAILURES - The following jobs have failed:

• Check Labels / Check labels (gh)
  RuntimeError: Error checking labels: PR does not have required labels
• inductor / linux-jammy-cpu-py3.9-gcc11-inductor / test (cpu_inductor_torchbench, 1, 2, linux.8xlarge.amx) (gh)
  lennard_jones
• inductor / linux-jammy-cpu-py3.9-gcc11-inductor / test (cpu_inductor_torchbench, 2, 2, linux.8xlarge.amx) (gh)
  vision_maskrcnn
• inductor / linux-jammy-cpu-py3.9-gcc11-inductor / test (dynamic_cpu_inductor_huggingface, 1, 1, linux.8xlarge.amx) (gh)
  YituTechConvBert
• inductor / linux-jammy-cpu-py3.9-gcc11-inductor / test (dynamic_cpu_inductor_timm, 1, 2, linux.8xlarge.amx) (gh)
  mixer_b16_224
• inductor / linux-jammy-cpu-py3.9-gcc11-inductor / test (dynamic_cpu_inductor_timm, 2, 2, linux.8xlarge.amx) (gh)
  xcit_large_24_p8_224
• inductor / linux-jammy-cpu-py3.9-gcc11-inductor / test (dynamic_cpu_inductor_torchbench, 1, 2, linux.8xlarge.amx) (gh)
  lennard_jones
• inductor / linux-jammy-cpu-py3.9-gcc11-inductor / test (dynamic_cpu_inductor_torchbench, 2, 2, linux.8xlarge.amx) (gh)
  vision_maskrcnn
• inductor / linux-jammy-cpu-py3.9-gcc11-inductor / test (inductor_torchbench_cpu_smoketest_perf, 1, 1, linux.24xl.spr-metal) (gh)
  TypeError: unsupported operand type(s) for |: 'type' and 'type'
• inductor / unit-test / linux-jammy-cpu-py3.9-gcc11-inductor / test (inductor_amx, 1, 2, linux.8xlarge.amx) (gh)
  TypeError: unsupported operand type(s) for |: 'type' and 'type'
• inductor / unit-test / linux-jammy-cpu-py3.9-gcc11-inductor / test (inductor_amx, 2, 2, linux.8xlarge.amx) (gh)
  TypeError: unsupported operand type(s) for |: 'type' and 'type'
• inductor / unit-test / linux-jammy-cpu-py3.9-gcc11-inductor / test (inductor_avx2, 1, 2, linux.10xlarge.avx2) (gh)
  TypeError: unsupported operand type(s) for |: 'type' and 'type'
• inductor / unit-test / linux-jammy-cpu-py3.9-gcc11-inductor / test (inductor_avx2, 2, 2, linux.10xlarge.avx2) (gh)
  TypeError: unsupported operand type(s) for |: 'type' and 'type'
• pull / linux-docs / build-docs-python-false (gh)
  Process completed with exit code 2.
• pull / linux-jammy-py3.9-clang12 / test (crossref, 1, 2, linux.2xlarge) (gh)
  RuntimeError: inductor/test_cutlass_backend 1/1 failed!
• pull / linux-jammy-py3.9-clang12 / test (crossref, 2, 2, linux.2xlarge) (gh)
  RuntimeError: inductor/test_fused_attention 1/1 failed!
• pull / linux-jammy-py3.9-clang12 / test (default, 1, 5, linux.4xlarge) (gh)
  RuntimeError: inductor/test_cutlass_backend 1/1 failed!
• pull / linux-jammy-py3.9-clang12 / test (default, 2, 5, linux.4xlarge) (gh)
  RuntimeError: inductor/test_kernel_benchmark 1/1 failed!
• pull / linux-jammy-py3.9-clang12 / test (default, 3, 5, linux.4xlarge) (gh)
  RuntimeError: inductor/test_fused_attention 1/1 failed!
• pull / linux-jammy-py3.9-clang12 / test (default, 4, 5, linux.4xlarge) (gh)
  RuntimeError: inductor/test_combo_kernels 1/1 failed!
• pull / linux-jammy-py3.9-clang12 / test (default, 5, 5, linux.4xlarge) (gh)
  RuntimeError: dynamo/test_functions 1/1 failed!
• pull / linux-jammy-py3.9-clang12 / test (dynamo_wrapped, 1, 3, linux.2xlarge) (gh)
  TypeError: unsupported operand type(s) for |: 'type' and 'type'
• pull / linux-jammy-py3.9-clang12 / test (dynamo_wrapped, 2, 3, linux.2xlarge) (gh)
  TypeError: unsupported operand type(s) for |: 'type' and 'type'
• pull / linux-jammy-py3.9-clang12 / test (dynamo_wrapped, 3, 3, linux.2xlarge) (gh)
  TypeError: unsupported operand type(s) for |: 'type' and 'type'
• pull / linux-jammy-py3.9-clang12 / test (einops, 1, 1, linux.2xlarge) (gh)
  dynamo/test_einops.py::TestEinops::test_layers_version_0_6_1
• pull / linux-jammy-py3.9-gcc11 / test (backwards_compat, 1, 1, linux.2xlarge) (gh)
  test_correct_module_names
• pull / linux-jammy-py3.9-gcc11 / test (default, 1, 5, linux.2xlarge) (gh)
  RuntimeError: inductor/test_cutlass_backend 1/1 failed!
• pull / linux-jammy-py3.9-gcc11 / test (default, 2, 5, linux.2xlarge) (gh)
  RuntimeError: inductor/test_fused_attention 1/1 failed!
• pull / linux-jammy-py3.9-gcc11 / test (default, 3, 5, linux.2xlarge) (gh)
  inductor/test_flex_decoding.py::TestFlexDecodingCPU::test_non_pow_2_headdim_head_dim_121_float32_cpu_float32
• pull / linux-jammy-py3.9-gcc11 / test (default, 4, 5, linux.2xlarge) (gh)
  inductor/test_flex_decoding.py::TestFlexDecodingCPU::test_non_pow_2_headdim_head_dim_94_float32_cpu_float32
• pull / linux-jammy-py3.9-gcc11 / test (default, 5, 5, linux.2xlarge) (gh)
  RuntimeError: inductor/test_aot_inductor_package 1/1 failed!
• pull / linux-jammy-py3.9-gcc11 / test (distributed, 1, 2, linux.2xlarge) (gh)
  RuntimeError: distributed/tensor/test_attention 1/1 failed!
• pull / linux-jammy-py3.9-gcc11 / test (distributed, 2, 2, linux.2xlarge) (gh)
  RuntimeError: distributed/test_dynamo_distributed 1/1 failed!
• pull / linux-jammy-py3.9-gcc11 / test (docs_test, 1, 1, linux.2xlarge) (gh)
  Process completed with exit code 2.
• pull / linux-jammy-py3.9-gcc11 / test (numpy_2_x, 1, 1, linux.2xlarge) (gh)
  RuntimeError: dynamo/test_functions 1/1 failed!

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[drisspg]

Closing, will continue to iterate on: drisspg#2