Can this api support V100 GPU?

[epiphanyer]

I get wrong when running this api on V100 with the following bug:

python: /project/lib/Dialect/TritonGPU/IR/Dialect.cpp:832: llvm::SmallVector mlir::triton::gpu::NvidiaMmaEncodingAttr::getElemsPerThread(llvm::ArrayRef, mlir::Type) const: Assertion `rank == 2 || (rank == 3 && isAmpere()) && "Unexpected rank of mma layout"' failed.

Dose this mean only support for A100? How can I fix it? Thank you.

[lucidrains]

@epiphanyer hey, thank you for your interest

could you try the flex attention variants?

[epiphanyer]

@lucidrains Thank you for your reply.
The flex attention variant can work.
I suspect the nvidia driver on the previous machine was too old. I ran the flex attention successfully on H20 with driver version 550.144.04.
By the way, it seems to work only on sequences that are multiples of 128.