Custom Attention Mask

[PaulCouairon]

Hi,

I am trying to use a custom attention mask (for example masking out scores below a certain threshold). The flex-attention backend seems to be suited for this but I am not fully sure to understand how to pass a score_mod function as an argument.

In addition, is #231 compatible with flex-attention backend?

Thank you in advance!

[alihassanijr]

Thanks for your interest.

We don't have any APIs for further modifying the Flex implementation yet. The main reason is that we're trying to focus on features that can be implemented with our own CUDA backends first, and supporting arbitrary masks (#124) is something I'd like to work on down the line, but hard to say when that will be.

We also note that due to some outstanding bugs with Flex, we strongly recommend using our CUDA backends if at all possible. We typically discourage from using Flex with torch.compile (which is what most users are after anyway) just to preserve correctness, but that is just terrible in terms of performance.

Regarding support for different head dim in value, it is not supported in our Flex backend yet, but I think it should just be a matter of allowing it if Flex already supports that.

If it is possible for you to share more details about your use case and workload, perhaps I can see if there's any alternatives I could suggest.

Our Flex mask is constructed here, and the call to create_block_mask is here, so adding in a custom score_mod should mostly be a matter of passing it down.

I can make a branch for that if you like, but in terms of making it a permanent feature I'd have to think about that more.

[alihassanijr]

for example masking out scores below a certain threshold

If it's just a single floating point that clips the dot products, I think it might be easy enough to add to our CUDA backends as well -- but again whether it ends up being a permanent feature is a matter of how common / important the use case is.

[PaulCouairon]

Yes, it’s really just a single floating-point threshold that clips the dot products, similar to the masking described in ProxyCLIP. I understand your point about whether it’s worth making permanent.

[alihassanijr]

I think that should be easy enough to implement in all backends.
I'll look into this.

[alihassanijr]

Just to clarify though, we're interested in masking the dot products, not attention weights, correct?

Something like:

# dot_product_min_value: float

S = q @ k.T
S = S.clip(min=dot_product_min_value)
A = S.softmax(dim=-1)
output = A @ v

[PaulCouairon]

Yes exactly!

[alihassanijr]

@PaulCouairon could you share what hardware you're using?
I can quickly add the feature to the appropriate backend for your use case first.

[PaulCouairon]

I am using Ampere architectures (A6000, A100).

[alihassanijr]

Forgot about the scale; should it clip then scale, or scale then clip?

# dot_product_min_value: float

S = q @ k.T
# S *= attention_scale
S = S.clip(min=dot_product_min_value, max=dot_product_max_value)
# S *= attention_scale
A = S.softmax(dim=-1)
output = A @ v

[alihassanijr]

I assumed clip then scale for now, but it's easy enough to change.

You can build it off of this branch:

https://github.com/alihassanijr/NATTEN-Torch/tree/feat/dot-product-clipping

You can optionally pass dot_product_min: float and dot_product_max: float to na{1d,2d,3d} to clip dot products.

PR to track: #252

NOTE: we don't have gradchecks in NATTEN anymore, and since this feature doesn't exist in SDPA or other standard attention APIs / kernels (as far as I know), I strongly recommend verifying this against some reference implementation in pure PyTorch.
You don't need to implement neighborhood attention -- you can just set kernel_size to your input size, and not do dilation, stride, or causal masking, and then compare against a pure PyTorch implementation of attention with dot product clipping.

I'll try to do that myself too, but I might not have time to work on it this week.

[PaulCouairon]

Thank you very much for your help! I'll try this.

[PaulCouairon]

Thank you again for your support! I built the library and ran a quick test. It seems that this operation clamps the dot product to the defined minimum or maximum values. I apologize for the confusion, I had been thinking of masking instead, i.e., setting the value to -inf if it falls below the threshold, something similar to:

S = S.masked_fill(S < min_value, float("-inf"))

[alihassanijr]

I'll take a look this weekend.