feat(torch-extras): Update `flash-attention` to 2.0.2; build 1.x and 2.x editions

[Eta0]

This change updates flash-attention in the torch-extras container to version 2.0.2.

The .github/workflows/torch-extras.yml workflow now publishes flash-attention 1.x and 2.x variants of the torch-extras container, distinguished by the suffixes -flash_attn<VERSION>. Right now the variants are:

• -flash-attn2.0.2
• -flash-attn1.0.9
  • Also updated from version 1.0.7 in this PR.

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5719570265 Image: ghcr.io/coreweave/ml-containers/torch:es-flash-attn-2-41a6446-base-cuda12.1.1-torch2.0.1-vision0.15.2-audio2.0.2

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5719570265 Image: ghcr.io/coreweave/ml-containers/torch:es-flash-attn-2-41a6446-base-cuda12.0.1-torch2.0.1-vision0.15.2-audio2.0.2

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5719570265 Image: ghcr.io/coreweave/ml-containers/torch:es-flash-attn-2-41a6446-base-cuda11.8.0-torch2.0.1-vision0.15.2-audio2.0.2

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5719570264 Image: ghcr.io/coreweave/ml-containers/torch-extras:es-flash-attn-2-41a6446-nccl-cuda11.8.0-nccl2.16.2-1-torch2.0.1-vision0.15.2-audio2.0.2-flash_attn1.0.9

[Eta0]

Did this successfully build earlier?

@wbrown On a testing container, yes. The code is valid. On our CI, it's only finishing most of the time. Seems to be resource-constrained. Requested some more powerful runners from @todie (related: #25) in lieu of dropping the concurrency even lower.

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5719570264 Image: ghcr.io/coreweave/ml-containers/torch-extras:es-flash-attn-2-41a6446-nccl-cuda11.8.0-nccl2.16.2-1-torch2.0.1-vision0.15.2-audio2.0.2-flash_attn2.0.2

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5719570264 Image: ghcr.io/coreweave/ml-containers/torch:es-flash-attn-2-41a6446-nccl-cuda12.0.1-nccl2.18.3-1-torch2.0.1-vision0.15.2-audio2.0.2

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5719570264 Image: ghcr.io/coreweave/ml-containers/torch:es-flash-attn-2-41a6446-nccl-cuda12.1.1-nccl2.18.3-1-torch2.0.1-vision0.15.2-audio2.0.2

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5719570264 Image: ghcr.io/coreweave/ml-containers/torch-extras:es-flash-attn-2-41a6446-nccl-cuda12.1.1-nccl2.18.3-1-torch2.0.1-vision0.15.2-audio2.0.2-flash_attn2.0.2

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5719570264 Image: ghcr.io/coreweave/ml-containers/torch-extras:es-flash-attn-2-41a6446-nccl-cuda12.1.1-nccl2.18.3-1-torch2.0.1-vision0.15.2-audio2.0.2-flash_attn1.0.9

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5721928378 Image: ghcr.io/coreweave/ml-containers/torch-extras:es-flash-attn-2-750fc34-nccl-cuda12.0.1-nccl2.18.3-1-torch2.0.1-vision0.15.2-audio2.0.2-flash_attn2.0.2

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5721928378 Image: ghcr.io/coreweave/ml-containers/torch-extras:es-flash-attn-2-750fc34-nccl-cuda12.0.1-nccl2.18.3-1-torch2.0.1-vision0.15.2-audio2.0.2-flash_attn1.0.9

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5721944894 Image: ghcr.io/coreweave/ml-containers/torch-extras:es-flash-attn-2-750fc34-nccl-cuda11.8.0-nccl2.16.2-1-torch2.0.1-vision0.15.2-audio2.0.2-flash_attn2.0.2

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5721944894 Image: ghcr.io/coreweave/ml-containers/torch-extras:es-flash-attn-2-750fc34-nccl-cuda11.8.0-nccl2.16.2-1-torch2.0.1-vision0.15.2-audio2.0.2-flash_attn1.0.9

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5722249778 Image: ghcr.io/coreweave/ml-containers/torch-extras:es-flash-attn-2-5fbc6bb-nccl-cuda12.1.1-nccl2.18.3-1-torch2.0.1-vision0.15.2-audio2.0.2-flash_attn1.0.9

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5722249778 Image: ghcr.io/coreweave/ml-containers/torch-extras:es-flash-attn-2-5fbc6bb-nccl-cuda12.1.1-nccl2.18.3-1-torch2.0.1-vision0.15.2-audio2.0.2-flash_attn2.0.2

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5722319221 Image: ghcr.io/coreweave/ml-containers/torch-extras:es-flash-attn-2-5fbc6bb-base-cuda12.1.1-torch2.0.1-vision0.15.2-audio2.0.2-flash_attn2.0.2

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5722319221 Image: ghcr.io/coreweave/ml-containers/torch-extras:es-flash-attn-2-5fbc6bb-base-cuda12.1.1-torch2.0.1-vision0.15.2-audio2.0.2-flash_attn1.0.9

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5722324333 Image: ghcr.io/coreweave/ml-containers/torch-extras:es-flash-attn-2-5fbc6bb-base-cuda12.0.1-torch2.0.1-vision0.15.2-audio2.0.2-flash_attn2.0.2

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5722324333 Image: ghcr.io/coreweave/ml-containers/torch-extras:es-flash-attn-2-5fbc6bb-base-cuda12.0.1-torch2.0.1-vision0.15.2-audio2.0.2-flash_attn1.0.9

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5722326374 Image: ghcr.io/coreweave/ml-containers/torch-extras:es-flash-attn-2-5fbc6bb-base-cuda11.8.0-torch2.0.1-vision0.15.2-audio2.0.2-flash_attn2.0.2

[github-actions[bot]]

@Eta0 Build complete, success: https://github.com/coreweave/ml-containers/actions/runs/5722326374 Image: ghcr.io/coreweave/ml-containers/torch-extras:es-flash-attn-2-5fbc6bb-base-cuda11.8.0-torch2.0.1-vision0.15.2-audio2.0.2-flash_attn1.0.9