flash_attention_forward

lightning_ir.flash.flash_bert.flash_attention_forward(self: BertSelfAttention, hidden_states: Tensor, attention_mask: FloatTensor | None, *args, **kwargs) → Tuple[Tensor]