Environment Configuration

[youngerlearner]

I encountered an environment issue, and I would like to know the supported CUDA version, torch version, and the corresponding flash-attn version in the runtime environment.

[shenyunhang]

CUDA: 12.6
torch: 2.6.0a0+df5bbc09d1.nv24.11
flash_attn: 2.6.3
`

We provide our docker image shenyunhang/pytorch:24.11-py3_2024-1224 at here: https://hub.docker.com/repository/docker/shenyunhang/pytorch/tags

It is based on nvcr.io/nvidia/pytorch:24.11-py3 in https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/tags.

We also install flashattn_hopper-3.0.0b1 to better support the devices with compute capability >= 9.0.