    def _compute_required_workspace(
        self,
        output_shape: Tuple[IntVar, IntVar, IntVar, IntVar],
        q_shape: Tuple[IntVar, IntVar, IntVar, IntVar],
        k_shape: Tuple[IntVar, IntVar, IntVar, IntVar],
    ) -> int:
        """
        Compute workspace size required for attention op.
        """
        is_float32 = self._attrs["inputs"][0]._attrs["dtype"] not in [
            "float16",
            "bfloat16",
        ]

        o_shape = [var._attrs["values"][-1] for var in output_shape]
        # We need a separate buffer of output accumulation
        # - when the intermediate output can't fit into the register file.
        # - when the accumulation type (float) is different from the output type.
        # See https://github.com/NVIDIA/cutlass/blob/209faf7b94ce4ba573d27389fb643962e75d0581/examples/41_fused_multi_head_attention/fmha_grouped.h#L79-L95
        needs_output_accum_buffer = (o_shape[-1] > 128) or not is_float32
        if needs_output_accum_buffer:  # Needs output accumulator buffer
            size_of_accum_element = 4  # Accumulation is always in float
            accu_size = size_of_accum_element * np.prod(o_shape)
        else:
            accu_size = 0

        # The backend which uses kernel_forward.h only needs accumulator buffer
        if not self._attrs["use_grouped_fmha"]:
            return accu_size

        # Number of problems is batch_size * num_heads
        problem_count = q_shape[0].upper_bound() * q_shape[1].upper_bound()

        size_of_int = 4
        size_of_int64 = 8
        # GEMM size is specified by 3 ints: m, n, k
        size_of_gemm_coord = 3 * size_of_int

        # There are two GEMM sizes for each problem, corresponding to 2 matrix
        # multiplications in attention
        problem_sizes_size = 2 * size_of_gemm_coord * problem_count

        # For each problem, need space for leading dimensions of 5 matrices:
        # Q, K, V, O. Leading dimensions are in int64.
        ld_sizes = 4 * size_of_int64 * problem_count

        # For each problem, pointers to 5 matrices: Q, K, V, O, O_accum
        size_of_ptr = 8  # 64-bit arch
        ptrs_sizes = 5 * size_of_ptr * problem_count
        total_size = problem_sizes_size + accu_size + ld_sizes + ptrs_sizes
        return total_size
Compute workspace size required for attention op.
python/aitemplate/compiler/ops/attention/mem_eff_attention.py