Add scalling factor to ring attention

[Sirorezka]

Hi, I want to pass softmax_scale factor to ring-attention (to use with cp), but it is currently not added in the patched _ring_compute_attention and not proprely patched in ring_fa3_varlen_func.

Desirrable behvaiour:
_ring_compute_attention should accept softmax_scale attribute
ring_fa3_varlen_func should accept softmax_scale attribute as well.

Thanks!

[mikasenghaas]

cc @S1ro1 @samsja

[S1ro1]

@Sirorezka Yeah we do not support this currently, however open to review a PR adding it if all the CP wrappers we support (FA2, FA3 and FA4) support it. I believe it should be as simple as extending _ring_compute_attention to accept softmax_scale and correctly passing it to the underlying ring_func.

[samsja]

Hi, I want to pass softmax_scale factor to ring-attention (to use with cp), but it is currently not added in the patched _ring_compute_attention and not proprely patched in ring_fa3_varlen_func.

Desirrable behvaiour: _ring_compute_attention should accept softmax_scale attribute ring_fa3_varlen_func should accept softmax_scale attribute as well.

Thanks!

this should work though with ulysse cp which do not use ring. ulysse is also faster in some case

[Sirorezka]

@S1ro1 thx. I am on it. Will create a PR.