Flash attention - NVIDIA V100

[hpc100]

Thank you for your work.

I tried to install the repo on a server equipped with an NVIDIA Tesla V100. However, the instructions using the command:
#flash attention pip install git+https://github.com/Dao-AILab/flash-attention.git don't seem to support the V100 architecture.

I wonder if it is possible to use LitePT without flash attention. If so, which parameters do I need to adapt (e.g., batch size, ...) and what are the potential effects on the final IoU?

Thanks.

[ywyue]

Hi @hpc100, sorry for the late reply. I was on holiday last week.

I would first recommend trying to find a way to install flash-attention v2 on the V100. If that is not possible, you need to make the following modifications:

• reduce enc_patch_size and dec_patch_size to a value that your GPU can afford, as turning off flash attention will increase memory consumption significantly so the original patch size of 1024 will likely not work.
• replace L328-L329 here with:

self.patch_size_max = patch_size
self.patch_size = 0
self.attn_drop = torch.nn.Dropout(attn_drop)

• add following at the beginning of def forward(self, point): here

self.patch_size = min(
     offset2bincount(point.offset).min().tolist(), self.patch_size_max
)

• replace code block in L353-L381 here with:

pos = point.grid_coord[order]
pos = pos.reshape(-1, K, 3)

# encode and reshape qkv: (N', K, 3, H, C') => (3, N', H, K, C')
q, k, v = (
    qkv.reshape(-1, K, 3, H, C // H).permute(2, 0, 3, 1, 4).unbind(dim=0)
)

# rope
# workround to make pointrope cuda float32 happy
q = self.rope(q.float(), pos).to(qkv.dtype)
k = self.rope(k.float(), pos).to(qkv.dtype)

# attn
q = q.float()
k = k.float()
attn = (q * self.scale) @ k.transpose(-2, -1)  # (N', H, K, K)

attn = attn.float()
attn = self.softmax(attn)
attn = self.attn_drop(attn).to(qkv.dtype)
feat = (attn @ v).transpose(1, 2).reshape(-1, C)

I am not sure the potential effects on final IoU as I didn't test it rigorously. If you decide to give it a try, feel free to share your results here.