def get_attention_flops(self, msl: int) -> int:
        """Computes the attention flops for the batch.

        Args:
            msl (int): The batch sequence length.

        Returns:
            attn_flops (int): The attention flops.
        """
        return (
            self.model.config.n_layers * 2 * 2 *
            (self.model.config.d_model * (msl**2))
        )
Computes the attention flops for the batch.

        Args:
            msl (int): The batch sequence length.

        Returns:
            attn_flops (int): The attention flops.
        
llmfoundry/models/mpt/modeling_mpt.py