def __init__(
        self,
        batch_size,
        max_seq_len,
        dropout=0,
        causal=False,
        dtype="float16",
    ):
        """Initialize attention module, create a tensor for seqlen"""
        super().__init__()
        self.cu_length = Parameter(shape=[batch_size + 1], dtype="int32")
        self.op = flash_attention(
            batch_size=batch_size,
            dropout=dropout,
            max_seq_len=max_seq_len,
            causal=causal,
        )
Initialize attention module, create a tensor for seqlen
python/aitemplate/frontend/nn/attention.py