Tiny upstream pr

[qianfengz]

This PR provide:

1. Synchronize to latest composable kernel commit which added inline-asm implementation of fp32 to bf16 RTN conversion. Using inline-asm RTN conversion is able to improve the performance when BF16+RTN is used
2. Add compiler options for compiling c++ extension on ROCM/HIP, which is able to improve the performance of HIP FMHA BWD on ROCM 6.2

• The following are benchmark results compared with triton when using RTN with those compiling options added on ROCM 6.2

Run reference fwd:
Reference fwd time: 28.90159034729004
Run reference bwd:
Reference bwd time: 48.68329620361328
Run triton fwd:
Triton fwd time: 2.0252671241760254
Run triton bwd:
Triton bwd time: 6.977703094482422
Run CK fwd:
xformers fwd time: 1.8350895643234253
Run CK fwd:
xformers bwd time: 7.089707374572754
(triton_dq - ref_dq).abs().mean()=tensor(0.0002, device='cuda:0', dtype=torch.bfloat16)
(triton_dk - ref_dk).abs().mean()=tensor(0.0001, device='cuda:0', dtype=torch.bfloat16)
(triton_dv - ref_dv).abs().mean()=tensor(0.0004, device='cuda:0', dtype=torch.bfloat16)
(xformer_dq - ref_dq).abs().mean()=tensor(0.0002, device='cuda:0', dtype=torch.bfloat16, grad_fn=<MeanBackward0>)
(xformer_dk - ref_dk).abs().mean()=tensor(0.0002, device='cuda:0', dtype=torch.bfloat16, grad_fn=<MeanBackward0>)
(xformer_dv - ref_dv).abs().mean()=tensor(6.7234e-05, device='cuda:0', dtype=torch.bfloat16, grad_fn=<MeanBackward0>)

• The following are benchmark results compared with triton when using RTN without those compiling options added on ROCM 6.2

Run reference fwd:
Reference fwd time: 28.867050170898438
Run reference bwd:
Reference bwd time: 48.91793441772461
Run triton fwd:
Triton fwd time: 2.056668996810913
Run triton bwd:
Triton bwd time: 6.982858180999756
Run CK fwd:
xformers fwd time: 1.8234171867370605
Run CK fwd:
xformers bwd time: 7.428786754608154
(triton_dq - ref_dq).abs().mean()=tensor(0.0002, device='cuda:0', dtype=torch.bfloat16)
(triton_dk - ref_dk).abs().mean()=tensor(0.0001, device='cuda:0', dtype=torch.bfloat16)
(triton_dv - ref_dv).abs().mean()=tensor(0.0004, device='cuda:0', dtype=torch.bfloat16)
(xformer_dq - ref_dq).abs().mean()=tensor(0.0002, device='cuda:0', dtype=torch.bfloat16, grad_fn=<MeanBackward0>)
(xformer_dk - ref_dk).abs().mean()=tensor(0.0002, device='cuda:0', dtype=torch.bfloat16, grad_fn=<MeanBackward0>)
(xformer_dv - ref_dv).abs().mean()=tensor(8.7738e-05, device='cuda:0', dtype=torch.bfloat16, grad_fn=<MeanBackward0>)

[danthe3rd]

Thanks! can you fix the formatting of setup.py tho? (see linter CI)

[qianfengz]

Any further layout changing is needed ?

[danthe3rd]

Sorry, forgot about that PR :)
Let me merge it