[ROCm] add the bias all row -inf support for jax unfused-attn#556

wangye805 wants to merge 4 commits intodevROCm/TransformerEngine:devfrom

yewang12/bias_all_row_inf_checkROCm/TransformerEngine:yewang12/bias_all_row_inf_checkCopy head branch name to clipboard