Searched defs:alibi_slopes_ (Results 1 – 1 of 1) sorted by relevance
329 void set_params_alibi(Flash_fwd_params ¶ms, std::optional<at::Tensor> &alibi_slopes_, int batch… in set_params_alibi()354 std::optional<at::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads in mha_fwd()551 std::optional<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads in mha_varlen_fwd()816 std::optional<at::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads in mha_bwd()1029 std::optional<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads in mha_varlen_bwd()1263 std::optional<at::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads in mha_fwd_kvcache()