    def __init__(
        self,
        causal_type: CausalType,
        epilogue_math_name: str,
        alpha0: float,
        alpha1: float,
        alpha1_divide_by_seq_len: bool = False,
    ) -> None:
        """Initialize classic_b2b_bmm op."""
        super().__init__()
        self._attrs["has_profiler"] = False
        self._attrs["causal_type"] = causal_type
        self._attrs["alpha0"] = alpha0
        self._attrs["alpha1"] = alpha1
        self._attrs["alpha1_divide_by_seq_len"] = alpha1_divide_by_seq_len

        import cutlass_lib

        if epilogue_math_name not in cutlass_lib.library.EpilogueMathName:
            raise RuntimeError(
                "Unsupported epilogue function! Please check "
                "python/aitemplate/utils/mk_cutlass_lib/extra_enum.py for a list of supported epilogue functions."
            )
        self._attrs["epilogue_math_name"] = epilogue_math_name
Initialize classic_b2b_bmm op.
python/aitemplate/compiler/ops/b2b_bmm/b2b_bmm_base.py