    def _test_mem_eff_attention(
        self,
        batch_size=16,
        nheads=16,
        seqlen=1024,
        n=1024,
        dropout_p=0.0,
        causal=False,
        dtype="float16",
        device="cuda",
        test_name="mem_eff_attention",
        rebuild=True,
        benchmark_ait=False,
        benchmark_pt=False,
        copy_op=False,
        use_perm=True,
        variable_seq_length_kv=False,
        variable_seq_length_q=False,
        skip_pt=False,
        use_grouped_fmha=False,
        atol=1e-3,
        rtol=1e-3,
    ):
        """
        Use skip_pt to avoid CUDA OOM when benchmarking with problem sizes
        which are too large for the PT implementation.
        """
        # Can't skip PT computation if we are benchmarking it
        assert not (benchmark_pt and skip_pt)

        torch_dtype = string_to_torch_dtype(dtype)
        d = n // nheads

        with torch.no_grad():
            x = torch.randn(
                batch_size,
                seqlen,
                n,
                device="cuda",
                dtype=torch_dtype,
            )
            Wqkv = torch.nn.Linear(
                nheads * d,
                3 * nheads * d,
                device=device,
                dtype=torch_dtype,
            )

            if variable_seq_length_kv:
                lengths_kv = torch.randint(0, seqlen + 1, size=(batch_size, 1))
                lengths_kv = lengths_kv.to(device="cuda")
            else:
                lengths_kv = torch.tensor(
                    [seqlen] * batch_size, dtype=torch.int, device="cuda"
                ).reshape(-1, 1)

            if variable_seq_length_q:
                lengths_q = torch.randint(0, seqlen + 1, size=(batch_size, 1))
                lengths_q = lengths_q.to(device="cuda")
            else:
                lengths_q = torch.tensor(
                    [seqlen] * batch_size, dtype=torch.int, device="cuda"
                ).reshape(-1, 1)

            seq_range = torch.arange(seqlen, device="cuda")
            attention_mask_bool_kv = (
                seq_range.unsqueeze(0).expand((batch_size, seqlen)) < lengths_kv
            ).unsqueeze(1)

            attention_mask_bool_q = (
                seq_range.unsqueeze(0).expand((batch_size, seqlen)) < lengths_q
            ).unsqueeze(2)

            x_unpad, indices, cu_seqlens, max_seqlen_in_batch = unpad_input(
                x, attention_mask_bool_kv
            )
            qkv_unpad = rearrange(
                Wqkv(x_unpad), "nnz (t h d) -> nnz t h d", t=3, h=nheads
            )
            qkv = rearrange(Wqkv(x), "b s (t h d) -> b s t h d", t=3, h=nheads)
            q, k, v = torch.split(qkv, 1, dim=2)
            if not skip_pt:
                y_pt = attention_ref(
                    qkv, attention_mask_bool_kv, dropout_p, causal=causal
                )

            total, _, num_heads, head_size = qkv_unpad.shape

            Q = Tensor(
                shape=[batch_size, num_heads, seqlen, head_size],
                dtype=dtype,
                name="q",
                is_input=True,
            )
            K = Tensor(
                shape=[batch_size, num_heads, seqlen, head_size],
                dtype=dtype,
                name="k",
                is_input=True,
            )
            V = Tensor(
                shape=[batch_size, num_heads, seqlen, head_size],
                dtype=dtype,
                name="v",
                is_input=True,
            )

            if variable_seq_length_kv:
                L_kv = Tensor(
                    shape=[batch_size, 1],
                    dtype="int",
                    name="lengths_kv",
                    is_input=True,
                )
            if variable_seq_length_q:
                L_q = Tensor(
                    shape=[batch_size, 1],
                    dtype="int",
                    name="lengths_q",
                    is_input=True,
                )

            mem_eff_attention_op = ops.mem_eff_attention(
                causal=causal,
                variable_seq_length_kv=variable_seq_length_kv,
                variable_seq_length_q=variable_seq_length_q,
                use_grouped_fmha=use_grouped_fmha,
            )
            if copy_op:
                mem_eff_attention_op = ops.mem_eff_attention(
                    **mem_eff_attention_op._get_op_attributes()
                )

            Y = mem_eff_attention_op(
                Q,
                K,
                V,
                L_kv if variable_seq_length_kv else None,
                L_q if variable_seq_length_q else None,
            )

            Y._attrs["is_output"] = True
            Y._attrs["name"] = "output"

            if rebuild:
                target = detect_target()
                module = compile_model(Y, target, "./tmp", test_name)
            else:
                module = Model(os.path.join("./tmp", test_name, "test.so"))

            q = torch.permute(q, (0, 3, 2, 1, 4)).reshape(
                batch_size, num_heads, seqlen, head_size
            )
            k = torch.permute(k, (0, 3, 2, 1, 4)).reshape(
                batch_size, num_heads, seqlen, head_size
            )
            v = torch.permute(v, (0, 3, 2, 1, 4)).reshape(
                batch_size, num_heads, seqlen, head_size
            )

            inputs = {
                "q": q.to(torch_dtype).contiguous(),
                "k": k.to(torch_dtype).contiguous(),
                "v": v.to(torch_dtype).contiguous(),
            }
            if variable_seq_length_kv:
                inputs["lengths_kv"] = lengths_kv.to(torch.int).contiguous()
            if variable_seq_length_q:
                inputs["lengths_q"] = lengths_q.to(torch.int).contiguous()

            y = torch.empty(
                [batch_size, seqlen, num_heads, head_size],
                dtype=torch_dtype,
                device="cuda",
            )
            module.run_with_tensors(inputs, [y])

            ret = {}

            if benchmark_ait or benchmark_pt:
                print(
                    f"batch_size = {batch_size}, nheads = {nheads}, seqlen = {seqlen}, n = {n}, causal = {causal}, dtype = {dtype}"
                )
                print(
                    f"variable_seq_length_kv = {variable_seq_length_kv}, variable_seq_length_q = {variable_seq_length_q}"
                )

            if benchmark_ait:
                # Warm up.
                for _ in range(5):
                    module.run_with_tensors(inputs, [y])
                # Benchmark AIT
                time_per_iter_ms, time_std, _ = module.benchmark_with_tensors(
                    inputs,
                    [y],
                    count=100,
                )
                print(
                    f"AIT benchmark eff-mem-attn time per iter: {time_per_iter_ms:.2f}ms"
                )

                ret["ait_time_per_iter_ms"] = time_per_iter_ms

            # y ~ [batch_size, seqlen, num_heads, head_size]
            # attention_mask_bool_q ~ [batch_size, seqlen, 1]
            if variable_seq_length_q:
                y = y.masked_fill_(~attention_mask_bool_q.unsqueeze(2), 0.0)
                if not skip_pt:
                    y_pt = y_pt.masked_fill_(~attention_mask_bool_q.unsqueeze(2), 0.0)

            if not skip_pt:
                torch.testing.assert_close(
                    y, y_pt.to(torch_dtype), atol=atol, rtol=rtol, equal_nan=True
                )

            if benchmark_pt:
                from aitemplate.testing.benchmark_pt import benchmark_torch_function

                func = attention_ref
                args = (
                    qkv.to(torch_dtype).cuda(),
                    attention_mask_bool_kv.cuda(),
                    dropout_p,
                    False,
                    False,
                )
                duration = benchmark_torch_function(100, func, *args)
                print(
                    f"PT benchmark eff-mem-attn time per iter: {duration:.2f}ms, BS: {batch_size}, QPS: {batch_size / duration:.2f}"
                )

                ret["pt_time_per_iter_ms"] = duration

        return ret
Use skip_pt to avoid CUDA OOM when benchmarking with problem sizes
which are too large for the PT implementation.
tests/unittest/ops/test_attention.py