    def __init__(
        self,
        causal,
        dropout=0,
        variable_seq_length_kv=False,
        variable_seq_length_q=False,
        use_grouped_fmha=False,
    ) -> None:
        """Initialize attention module"""
        super().__init__()
        assert dropout == 0
        self._attrs["op"] = "mem_eff_attention"
        self._attrs["has_profiler"] = False
        self._attrs["dropout"] = dropout
        self._attrs["causal"] = causal
        self._attrs["variable_seq_length_kv"] = variable_seq_length_kv
        self._attrs["variable_seq_length_q"] = variable_seq_length_q
        self._attrs["head_size"] = -1
        self._attrs["workspace"] = 0
        self._attrs["use_grouped_fmha"] = use_grouped_fmha
        self.exec_key_template = EXEC_KEY_TEMPLATE
        self.shape_eval_template = SHAPE_FUNC_TEMPLATE
Initialize attention module
python/aitemplate/compiler/ops/attention/mem_eff_attention.py