# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
# pyre-unsafe
# pyre-ignore-all-errors[29]

import hashlib
import math
from typing import Callable, List, Optional, Tuple, Type

import pytest
import torch
from mslk.attention import fmha
from mslk.attention.fmha.common import AttentionFwOpBase
from mslk.attention.fmha.merge_training import (
    memory_efficient_attention_partial_autograd,
    merge_attentions_autograd,
    Partial,
)

from .utils import assert_allclose, disable_on_rocm

compute_capability = (0, 0)
if torch.cuda.is_available():
    compute_capability = torch.cuda.get_device_capability("cuda")
sm80_or_better_only = pytest.mark.skipif(
    torch.version.cuda is not None and compute_capability < (8, 0),
    reason="requires sm80+",
)


# This temporary working is necessary because the MTIA test collection might not happen
# on the same device as the device the tests are actually executed on. If test collection
# is done on a device without MTIA, the supported masks will contain masks that MTIA support
# and the corresponding tests will get collected. But when it comes time to actually run the
# tests, the mask won't be supported because it is run on an actual MTIA device.
def get_supported_attn_bias_types(op):
    supported_attn_bias_types = op.SUPPORTED_ATTN_BIAS_TYPES

    try:
        import mtia.host_runtime.torch_mtia.dynamic_library  # noqa  # type: ignore

        supported_attn_bias_types = [
            b
            for b in supported_attn_bias_types
            if not issubclass(
                b,
                (
                    fmha.attn_bias.PagedBlockDiagonalGappyKeysMask,
                    fmha.attn_bias.PagedBlockDiagonalPaddedKeysMask,
                ),
            )
        ]
    except (ImportError, OSError):
        pass

    return supported_attn_bias_types


@disable_on_rocm
@sm80_or_better_only
@pytest.mark.parametrize(
    "op",
    [
        fmha.triton_splitk.FwOp,
        fmha.flash.FwOp,
        fmha.flash3.FwOp,
        None,
    ],
    ids=lambda op: "None" if op is None else op.NAME,
)
@pytest.mark.parametrize("G,H", [(1, 11), (7, 1), (1, 1), (7, 11), (None, 11)])
@pytest.mark.parametrize(
    "write_lse", (False, True), ids=lambda x: "write_lse" if x else ""
)
@pytest.mark.parametrize(
    "stack_inputs", (False, True), ids=lambda x: "stack_inputs" if x else ""
)
def test_merge_attentions_nobias(
    write_lse: bool,
    stack_inputs: bool,
    op: Type[AttentionFwOpBase],
    G: Optional[int],
    H: int,
):
    """
    Merging the same attention twice shouldn't change anything.
    This also tests the shape of the lse output of each permitted op.
    """
    if op is fmha.flash3.FwOp and not op.is_available():
        pytest.skip("Flash3 not available")
    B, Mq, K = 13, 3, 192
    if op is fmha.triton_splitk.FwOp:
        K = 128
    case_name = str((write_lse, G, H, stack_inputs)).encode("ascii")
    # patternlint-disable-next-line poor-choice-of-hash-function
    many_keys = hashlib.md5(case_name).digest()[0] % 2
    M = [5, 100000][many_keys]
    if op is None or torch.bfloat16 in op.SUPPORTED_DTYPES:
        dtype = torch.bfloat16
    else:
        dtype = next(iter(op.SUPPORTED_DTYPES))
    if dtype == torch.float8_e4m3fn:
        pytest.skip("float8 not supported")
    if G is None:
        q = 3 * torch.rand(B, Mq, H, K, dtype=dtype, device="cuda")
        k = (3 * torch.rand(B, M, 1, K, dtype=dtype, device="cuda")).expand(B, M, H, K)
        v = (3 * torch.rand(B, M, 1, K, dtype=dtype, device="cuda")).expand(B, M, H, K)
    else:
        q = 3 * torch.rand(B, Mq, G, H, K, dtype=dtype, device="cuda")
        k = (3 * torch.rand(B, M, G, 1, K, dtype=dtype, device="cuda")).expand(
            B, M, G, H, K
        )
        v = (3 * torch.rand(B, M, G, 1, K, dtype=dtype, device="cuda")).expand(
            B, M, G, H, K
        )
    out1, lse1 = fmha.memory_efficient_attention_partial(q, k, v, op=op)
    assert out1.shape == q.shape
    M_ceil = lse1.shape[-1]
    assert M_ceil >= Mq
    assert lse1.shape == (B, H, M_ceil) if G is None else (B, G, H, M_ceil)
    lse1 = lse1[..., :Mq]

    attn_chunks = [out1, out1]
    lse_chunks = [lse1, lse1]
    attn_chunks_ = torch.stack(attn_chunks) if stack_inputs else attn_chunks
    lse_chunks_ = torch.stack(lse_chunks) if stack_inputs else lse_chunks
    out, lse = fmha.merge_attentions(attn_chunks_, lse_chunks_, write_lse=write_lse)  # type: ignore
    assert out.shape == out1.shape
    assert_allclose(out1, out, rtol=1e-3, atol=1e-3, msg="out")
    if write_lse:
        assert lse is not None
        assert lse.shape[:-1] == lse1.shape[:-1]
        assert_allclose(
            lse1[..., :Mq] + math.log(2), lse[..., :Mq], rtol=1e-3, atol=1e-3, msg="lse"
        )
    else:
        assert lse is None


@disable_on_rocm
@sm80_or_better_only
@pytest.mark.parametrize(
    "dtype,op",
    [
        (torch.bfloat16, fmha.triton_splitk.FwOp_S1),
        # Cutlass's LSE is not consistent
        # (torch.float32, fmha.cutlass.FwOp),
        (torch.bfloat16, fmha.flash.FwOp),
    ],
    ids=lambda o: f"{o.NAME}" if hasattr(o, "NAME") else str(o),
)
@pytest.mark.parametrize("num_queries", [1])
@pytest.mark.parametrize("bmghk", [True, False], ids=lambda x: "bmghk" if x else "")
def test_partial_paged(
    dtype: torch.dtype, op: Type[AttentionFwOpBase], num_queries: int, bmghk: bool
):
    B = 128
    N_H_L = 8
    D_H = 128
    page_size = 256
    G = 2 if bmghk else 1
    block_tables = torch.zeros((B, 1), dtype=torch.int32, device="cuda")
    torch.manual_seed(1)
    output_dtype = torch.float32 if op.SUPPORTS_OUTPUT_DTYPE else None

    B_T = num_queries * B

    q = torch.randn((1, B_T, G, N_H_L, D_H), dtype=dtype, device="cuda")
    k = torch.randn((1, page_size, G, 1, D_H), dtype=dtype, device="cuda")
    v = torch.randn_like(k)
    k = k.expand(1, page_size, G, N_H_L, D_H)
    v = v.expand(1, page_size, G, N_H_L, D_H)
    if not bmghk:
        q = q[:, :, 0]
        k = k[:, :, 0]
        v = v[:, :, 0]

    attn_bias = (
        fmha.attn_bias.PagedBlockDiagonalCausalWithOffsetPaddedKeysMask.from_seqlens(
            q_seqlen=[num_queries] * B,
            kv_seqlen=[1] + ([100] * (B - 1)),
            page_size=page_size,
            block_tables=block_tables,
        )
    )

    if attn_bias not in get_supported_attn_bias_types(op):
        pytest.skip("Not supported bias")

    attn_chunk, lse_chunk = fmha.memory_efficient_attention_partial(
        q,
        k,
        v,
        attn_bias,
        op=op,
        output_dtype=output_dtype,
    )
    if bmghk:
        assert attn_chunk.shape == (1, B_T, G, N_H_L, D_H)
        assert lse_chunk.shape == (
            1,
            G,
            N_H_L,
            B_T,
        ), f"{lse_chunk.shape=}, {(1, G, N_H_L, B_T)=}"
    else:
        assert attn_chunk.shape == (1, B_T, N_H_L, D_H)
        assert lse_chunk.shape == (
            1,
            N_H_L,
            B_T,
        ), f"{lse_chunk.shape=}, {(1, N_H_L, B_T)=}"


@disable_on_rocm
@sm80_or_better_only
@pytest.mark.parametrize(
    "dtype,op",
    [
        (torch.bfloat16, fmha.triton_splitk.FwOp_S1),
        (torch.bfloat16, fmha.triton_splitk.FwOp_S32),
        # Cutlass's LSE is not consistent
        # (torch.float32, fmha.cutlass.FwOp),
        (torch.bfloat16, fmha.flash.FwOp),
    ],
    ids=lambda o: f"{o.NAME}" if hasattr(o, "NAME") else str(o),
)
@pytest.mark.parametrize("num_queries", [1, 2])
@pytest.mark.parametrize("bmghk", [True, False], ids=lambda x: "bmghk" if x else "")
@pytest.mark.parametrize(
    "stack_inputs", (False, True), ids=lambda x: "stack_inputs" if x else ""
)
def test_merge_attentions_decoding(
    dtype: torch.dtype,
    op: Type[AttentionFwOpBase],
    num_queries: int,
    bmghk: bool,
    stack_inputs: bool,
):
    """
    Compute decoding attention on chunks of K/V and merge them together.
    Compare with computing attention on the whole K/V.
    """
    MAX_T = 8192
    B = 128
    N_H_L = 8
    D_H = 128
    G = 2 if bmghk else 1
    torch.manual_seed(1)
    output_dtype = torch.float32 if op.SUPPORTS_OUTPUT_DTYPE else None

    num_chunks = 10

    chunk_starts = sorted(
        torch.randint(low=1, high=MAX_T // 2, size=(num_chunks,)).tolist()
    )
    chunk_starts[0] = 0
    chunk_starts.append(MAX_T)

    # We construct sequences so that even the last chunk has a non-empty part of every sequence
    # as long as the number of queries.
    # Otherwise the corresponding LSE will be -inf and that'll propagate to the whole sum.
    # It is possible to teach the kernel to ignore infinite LSEs, but in practical use cases
    # of merging attention, e.g. a batch of sequences with a common prefix, this condition should be satisfied.
    k_lens = torch.randint(
        low=chunk_starts[-2] + num_queries, high=MAX_T, size=(B,)
    ).tolist()
    q_lens = [num_queries] * B
    B_T = num_queries * B

    q = torch.randn((1, B_T, G, N_H_L, D_H), dtype=dtype, device="cuda")
    k = torch.randn((B, MAX_T, G, 1, D_H), dtype=dtype, device="cuda")
    v = torch.randn_like(k)
    if not bmghk:
        q = q[:, :, 0]

    # Compute per-chunk attention
    chunks_output: List[Tuple[torch.Tensor, torch.Tensor]] = []
    for i in range(num_chunks):
        chunk_start, chunk_end = chunk_starts[i], chunk_starts[i + 1]
        k_chunk = k[:, chunk_start:chunk_end, ...]
        v_chunk = v[:, chunk_start:chunk_end, ...]
        axk = k_chunk.reshape(-1, G, 1, D_H).expand(1, -1, G, N_H_L, D_H)
        axv = v_chunk.reshape(-1, G, 1, D_H).expand(1, -1, G, N_H_L, D_H)
        if not bmghk:
            axk = axk[:, :, 0]
            axv = axv[:, :, 0]

        bias_type = fmha.attn_bias.BlockDiagonalPaddedKeysMask
        if i + 1 == num_chunks:
            bias_type = fmha.attn_bias.BlockDiagonalCausalWithOffsetPaddedKeysMask
        attn_bias = bias_type.from_seqlens(
            q_seqlen=q_lens,
            kv_padding=chunk_end - chunk_start,
            kv_seqlen=[max(min(x, chunk_end) - chunk_start, 0) for x in k_lens],
        )

        attn_chunk, lse_chunk = fmha.memory_efficient_attention_partial(
            q,
            axk,
            axv,
            attn_bias,
            op=op,
            output_dtype=output_dtype,
        )
        if bmghk:
            assert attn_chunk.shape == (1, B_T, G, N_H_L, D_H)
            assert lse_chunk.shape == (1, G, N_H_L, B_T)
        else:
            assert attn_chunk.shape == (1, B_T, N_H_L, D_H)
            assert lse_chunk.shape == (1, N_H_L, B_T)
        chunks_output.append((attn_chunk, lse_chunk))

    # Merge attention from all chunks
    attn_split = [attn_chunk for attn_chunk, _ in chunks_output]
    lse_split = [lse_chunk for _, lse_chunk in chunks_output]
    if stack_inputs:
        attn_out, lse_out = fmha.merge_attentions(
            torch.stack(attn_split), torch.stack(lse_split), output_dtype=dtype
        )
    else:
        attn_out, lse_out = fmha.merge_attentions(
            attn_split, lse_split, output_dtype=dtype
        )
    assert lse_out is not None

    # Compute attention on the full K/V
    attn_bias = fmha.attn_bias.BlockDiagonalCausalWithOffsetPaddedKeysMask.from_seqlens(
        q_seqlen=q_lens,
        kv_padding=MAX_T,
        kv_seqlen=k_lens,
    )
    axk = k.view(1, -1, G, 1, D_H).expand(1, -1, G, N_H_L, D_H)
    axv = v.view(1, -1, G, 1, D_H).expand(1, -1, G, N_H_L, D_H)
    if not bmghk:
        axk = axk[:, :, 0]
        axv = axv[:, :, 0]
    attn_full, lse_full = fmha.memory_efficient_attention_partial(
        q,
        axk,
        axv,
        attn_bias,
        op=op,
        output_dtype=output_dtype,
    )

    assert_allclose(
        lse_out.to(lse_full.dtype), lse_full, rtol=1e-3, atol=1e-3, msg="lse"
    )
    assert_allclose(
        attn_out.to(attn_full.dtype), attn_full, rtol=1e-3, atol=1e-3, msg="out"
    )

    attn_full2 = fmha.memory_efficient_attention_forward(
        q,
        axk,
        axv,
        attn_bias,
        op=op,
        output_dtype=output_dtype,
    )
    assert_allclose(attn_full2, attn_full, rtol=1e-3, atol=1e-3, msg="out2")


@disable_on_rocm
@sm80_or_better_only
@pytest.mark.parametrize(
    "dtype,op",
    [
        (torch.bfloat16, fmha.triton_splitk.FwOp_S1),
        (torch.bfloat16, fmha.triton_splitk.FwOp_S32),
    ],
    ids=lambda o: f"{o.NAME}" if hasattr(o, "NAME") else str(o),
)
@pytest.mark.parametrize("gqa", [False, True], ids=lambda x: "gqa" if x else "")
def test_merge_attentions_sharedinput(
    dtype: torch.dtype,
    op: Type[AttentionFwOpBase],
    gqa: bool,
):
    """
    Compute decoding attention on chunks of K/V and merge them together.
    Compare with computing attention on the whole K/V.
    """
    MAX_T = 8192
    N_H_L = 16
    D_H = 128
    G = 2
    torch.manual_seed(1)
    output_dtype = torch.float32 if op.SUPPORTS_OUTPUT_DTYPE else None

    shared_length = 20
    full_lengths = [30, 35, 40]

    attn_bias = fmha.attn_bias.BlockDiagonalCausalWithOffsetPaddedKeysMask.from_seqlens(
        q_seqlen=[1, 1, 1],
        kv_padding=MAX_T,
        kv_seqlen=full_lengths,
    )
    attn_bias1 = fmha.attn_bias.BlockDiagonalPaddedKeysMask.from_seqlens(
        q_seqlen=[2, 1],
        kv_padding=MAX_T,
        kv_seqlen=[shared_length, 0],
    )
    attn_bias2 = fmha.attn_bias.BlockDiagonalGappyKeysMask.from_seqlens(
        q_seqlen=[1, 1, 1],
        kv_seqstarts=[shared_length, MAX_T + shared_length, 2 * MAX_T, 3 * MAX_T],
        kv_seqlen=[
            full_lengths[0] - shared_length,
            full_lengths[1] - shared_length,
            full_lengths[2],
        ],
    )

    q = torch.randn((1, 3, G, N_H_L, D_H), dtype=dtype, device="cuda")
    k = torch.randn((3, MAX_T, G, 1 if gqa else N_H_L, D_H), dtype=dtype, device="cuda")
    v = torch.randn_like(k)
    k[1, :shared_length] = k[0, :shared_length]
    v[1, :shared_length] = v[0, :shared_length]
    k = k.flatten(end_dim=1)[None]
    v = v.flatten(end_dim=1)[None]
    k = k.expand((1, 3 * MAX_T, G, N_H_L, D_H))
    v = v.expand((1, 3 * MAX_T, G, N_H_L, D_H))

    attn_chunk1, lse_chunk1 = fmha.memory_efficient_attention_partial(
        q,
        k,
        v,
        attn_bias1,
        op=op,
        output_dtype=output_dtype,
    )
    assert attn_chunk1.shape == (1, 3, G, N_H_L, D_H)
    assert lse_chunk1.shape == (1, G, N_H_L, 3)
    if gqa:
        attn_chunk1a, lse_chunk1a = fmha.memory_efficient_attention_partial(
            q,
            k.contiguous(),
            v,
            attn_bias1,
            op=op,
            output_dtype=output_dtype,
        )
        assert attn_chunk1a.shape == (1, 3, G, N_H_L, D_H)
        assert lse_chunk1a.shape == (1, G, N_H_L, 3)
        assert_allclose(
            attn_chunk1a.nan_to_num(0, 0, 0), attn_chunk1.nan_to_num(0, 0, 0)
        )
        assert_allclose(lse_chunk1a.nan_to_num(0, 0, 0), lse_chunk1.nan_to_num(0, 0, 0))

    attn_chunk2, lse_chunk2 = fmha.memory_efficient_attention_partial(
        q,
        k,
        v,
        attn_bias2,
        op=op,
        output_dtype=output_dtype,
    )
    assert attn_chunk2.shape == (1, 3, G, N_H_L, D_H)
    assert lse_chunk2.shape == (1, G, N_H_L, 3)
    # Merge attention from all chunks

    attn_out, lse_out = fmha.merge_attentions(
        [attn_chunk1, attn_chunk2],
        [lse_chunk1, lse_chunk2],
        output_dtype=dtype,  # type: ignore
    )
    assert lse_out is not None

    # Compute attention on the full K/V
    attn_full, lse_full = fmha.memory_efficient_attention_partial(
        q,
        k,
        v,
        attn_bias,
        op=op,
        output_dtype=output_dtype,
    )
    assert_allclose(
        attn_out.to(attn_full.dtype), attn_full, rtol=1e-2, atol=2e-3, msg="out"
    )
    assert_allclose(
        lse_out.to(lse_full.dtype), lse_full, rtol=1e-3, atol=1e-3, msg="lse"
    )


@sm80_or_better_only
@pytest.mark.parametrize("bmghk", (False, True))
def test_merge_attentions_against_ref(bmghk: bool):
    split_k = 16
    B = 12
    M = 137
    G = 2 if bmghk else 1
    N_H_L = 8
    D_H = 128
    dtype = torch.float32

    attn_split = torch.randn([split_k, B, M, G, N_H_L, D_H], dtype=dtype, device="cuda")
    lse_split = torch.randn([split_k, B, G, N_H_L, M], dtype=dtype, device="cuda")

    if not bmghk:
        attn_split = attn_split[:, :, :, 0]
        lse_split = lse_split[:, :, 0]

    attn_out_ref, lse_out_ref = _merge_attentions_ref(attn_split, lse_split)
    attn_out, lse_out = fmha.merge_attentions(attn_split, lse_split)

    torch.testing.assert_close(lse_out, lse_out_ref, rtol=1e-4, atol=1e-4)
    torch.testing.assert_close(attn_out, attn_out_ref, rtol=1e-4, atol=1e-4)


def _merge_attentions_ref(attn_split, lse_split):
    """
    attn_split: [split_k, B, M, (G,) H, Kq]
    lse_split: [split_k, B, (G,) H, M]
    """
    is_bmghk = len(attn_split.shape) == 6
    if not is_bmghk:
        attn_split = attn_split.unsqueeze(3)
        lse_split = lse_split.unsqueeze(2)

    lse_split = lse_split[..., None].moveaxis(4, 2)  # [split_k, B, M, G, H, 1]

    lse_max, _ = torch.max(lse_split, dim=0)  # [B, M, G, H, 1]
    sumexp_normalized = torch.exp(lse_split - lse_max)  # [split_k, B, M, G, H, 1]
    denominator = sumexp_normalized.sum(dim=0)  # [B, M, G, H, 1]
    numerator = (sumexp_normalized * attn_split).sum(dim=0)  # [B, M, G, H, K]

    attn_out = numerator / denominator  # [B, M_ceil, G, H, Kq]
    lse_out = lse_max + torch.log(denominator)
    lse_out = lse_out.squeeze(4).permute(0, 2, 3, 1)  # [B, G, H, M]

    if not is_bmghk:
        attn_out = attn_out.squeeze(2)
        lse_out = lse_out.squeeze(1)

    return attn_out, lse_out


@sm80_or_better_only
def test_merge_attention_with_compile() -> None:
    op = fmha.flash3.FwOp
    if not op.is_available():
        pytest.skip("Op is not available")
    dtype = torch.bfloat16
    B, M, H, K = 1, 256, 2, 128
    q, k, v = [
        (3 * torch.rand(B, M, H, K, dtype=dtype, device="cuda")) for _ in range(3)
    ]

    out1, lse1 = fmha.memory_efficient_attention_partial(q, k, v, op=op)

    def run_code() -> torch.Tensor:
        out, _ = fmha.merge_attentions([out1], [lse1], write_lse=True)
        return out

    out_ref = run_code()
    out_c = torch.compile(run_code, fullgraph=True)()

    assert torch.allclose(out_ref, out_c, atol=1e-2, rtol=1e-2)

    q.requires_grad_(True)
    out1, lse1 = fmha.memory_efficient_attention_partial(q, k, v, op=op)
    loss = fmha.merge_attentions([out1], [lse1])[0].sum()
    with pytest.raises(
        NotImplementedError,
        match="Backward pass is not implemented for merge_attentions",
    ):
        loss.backward()


@sm80_or_better_only
@disable_on_rocm
def test_merge_training():
    torch.manual_seed(1)
    B, M, H, K = 1, 50, 1, 128
    dtype = torch.bfloat16
    op = (fmha.flash3.FwOp, fmha.flash3.BwOp)
    q = 3 * torch.rand((B, M, H, K), device="cuda", dtype=dtype)
    k = 3 * torch.rand((B, M, H, K), device="cuda", dtype=dtype)
    v = 3 * torch.rand((B, M, H, K), device="cuda", dtype=dtype)
    grad_out = 3 * torch.rand((B, M, H, K), device="cuda", dtype=dtype)

    total_attention, (q_grad, k_grad, v_grad) = torch.autograd.functional.vjp(
        fmha.memory_efficient_attention, (q, k, v), grad_out
    )

    def total_attn_via_Partial(q_, k_, v_):
        return merge_attentions_autograd(
            memory_efficient_attention_partial_autograd(q_, k_, v_)
        )

    attn, (q_grad_, k_grad_, v_grad_) = torch.autograd.functional.vjp(
        total_attn_via_Partial, (q, k, v), grad_out
    )

    assert_allclose(attn, total_attention, rtol=1e-1, atol=1e-3, msg="out")
    assert_allclose(k_grad_, k_grad, rtol=1e-2, atol=1e-3, msg="dk_")
    assert_allclose(q_grad_, q_grad, rtol=1e-2, atol=1e-3, msg="dq_")
    assert_allclose(v_grad_, v_grad, rtol=1e-2, atol=1e-3, msg="dv_")

    def attn_via_Partial(q_, k_, v_):
        split = M // 2
        k1 = k_[:, :split]
        k2 = k_[:, split:]
        v1 = v_[:, :split]
        v2 = v_[:, split:]

        partial1 = memory_efficient_attention_partial_autograd(q_, k1, v1, op=op)
        partial2 = memory_efficient_attention_partial_autograd(q_, k2, v2, op=op)
        return merge_attentions_autograd(partial1, partial2)

    merged, (q_grad_, k_grad_, v_grad_) = torch.autograd.functional.vjp(
        attn_via_Partial, (q, k, v), grad_out
    )
    assert_allclose(
        merged.to(total_attention.dtype),
        total_attention,
        rtol=1e-1,
        atol=1e-3,
        msg="out",
    )
    assert_allclose(k_grad_, k_grad, rtol=0.1, atol=0.9, msg="dk")
    assert_allclose(q_grad_, q_grad, rtol=0.1, atol=0.9, msg="dq")
    assert_allclose(v_grad_, v_grad, rtol=0.1, atol=0.9, msg="dv")


def _pad_seqdim(partial: Partial, left: int, right: int) -> Partial:
    padding = (0, 0) * (3 if partial.is_bmghk() else 2) + (left, right)
    return partial.apply(lambda x: torch.nn.functional.pad(x, padding))


def _slice(partial: Partial, a: int, b: int) -> Partial:
    return partial.apply(lambda x: x[:, a:b])


@sm80_or_better_only
@disable_on_rocm
def test_merge_training_compile():
    torch.manual_seed(1)
    B, M, H, K = 1, 50, 1, 128
    dtype = torch.bfloat16
    q = 3 * torch.rand((B, M, H, K), device="cuda", dtype=dtype)
    k = 3 * torch.rand((B, M, H, K), device="cuda", dtype=dtype)
    v = 3 * torch.rand((B, M, H, K), device="cuda", dtype=dtype)

    def f(
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
    ) -> torch.Tensor:
        k1 = k[:, : M // 2]
        k2 = k[:, M // 2 :]
        v1 = v[:, : M // 2]
        v2 = v[:, M // 2 :]
        partial1 = memory_efficient_attention_partial_autograd(q, k1, v1)
        partial2 = memory_efficient_attention_partial_autograd(q, k2, v2)
        partial2 = _pad_seqdim(_pad_seqdim(partial2, 2, 3), -2, -3)
        partial2 = _slice(_pad_seqdim(partial2, 2, 3), 2, -3)
        merged = merge_attentions_autograd(partial1, partial2)
        return merged.sum()

    out, grads = torch.autograd.functional.vjp(f, (q, k, v))
    outc, gradsc = torch.autograd.functional.vjp(
        torch.compile(fullgraph=True)(f), (q, k, v)
    )
    assert_allclose(out, outc, rtol=1e-1, atol=1e-3)
    for i, (grad, gradc) in enumerate(zip(grads, gradsc)):
        assert_allclose(grad, gradc, rtol=1e-1, atol=1e-3, msg=f"grad{i}")

    def g(q: torch.Tensor, k1: torch.Tensor, v1: torch.Tensor) -> torch.Tensor:
        partial1 = memory_efficient_attention_partial_autograd(q, k1, v1)
        return merge_attentions_autograd(partial1).sum()

    out, grads = torch.autograd.functional.vjp(g, (q, k, v))
    outc, gradsc = torch.autograd.functional.vjp(
        torch.compile(fullgraph=True)(g), (q, k, v)
    )
    assert_allclose(out, outc, rtol=1e-1, atol=1e-3)
    for i, (grad, gradc) in enumerate(zip(grads, gradsc)):
        assert_allclose(grad, gradc, rtol=1e-1, atol=1e-3, msg=f"grad{i}")


@sm80_or_better_only
def test_merge_training_zilch():
    with pytest.raises(ValueError, match="No partials to merge"):
        merge_attentions_autograd()


@sm80_or_better_only
def test_merge_training_undilate():
    torch.manual_seed(1)

    def undilate(factor: int) -> Callable[[torch.Tensor], torch.Tensor]:
        """
        For a given factor, operate on BMHK attention output as follows:
        If sequence length is M and there are H * factor heads,
        redistribute so there are H heads and each
        original sequence position is factor times inflated.
        """

        def inner(x: torch.Tensor) -> torch.Tensor:
            M = x.shape[1]
            H = x.shape[2] // factor
            return x.flatten(1, 2).unflatten(1, (M * factor, H))

        return inner

    B, M, F, H, K = 1, 2, 3, 5, 128
    dtype = torch.bfloat16
    q = 3 * torch.rand((B, M, F * H, K), device="cuda", dtype=dtype)
    k = 3 * torch.rand((B, M, F * H, K), device="cuda", dtype=dtype)
    v = 3 * torch.rand((B, M, F * H, K), device="cuda", dtype=dtype)

    bias = fmha.BlockDiagonalMask.from_seqlens([1] * M)
    partial = memory_efficient_attention_partial_autograd(q, k, v, bias)

    q = q.reshape(B, M * F, H, K)
    k = k.reshape(B, M * F, H, K)
    v = v.reshape(B, M * F, H, K)

    bias = fmha.BlockDiagonalMask.from_seqlens([1] * (M * F))
    expected = memory_efficient_attention_partial_autograd(q, k, v, bias)
    undilated = partial.apply(undilate(F))

    assert_allclose(undilated._attn, expected._attn)
    assert_allclose(undilated._lse, expected._lse)
