Support flash-attention in forms of custom-call

[ApsarasX]

Usage:

import torch_xla.core.functions as xf

xf.flash_attn(
    query,
    key,
    value,
    *,
    dropout_rate=0.0,
    scale=None,
    is_causal=False,
    alibi_slopes=None,
    deterministic=False,
    return_softmax=False,
)

xf.flash_attn_varlen(
    query,
    key,
    value,
    cu_seqlens_query,
    cu_seqlens_key,
    *,
    max_seqlen_q,
    max_seqlen_k,
    dropout_rate=0.0,
    scale=None,
    is_causal=False,
    alibi_slopes=None,
    deterministic=False,
    return_softmax=False,
)