import pytest
import torch
from torch.nn.attention.flex_attention import flex_attention
from flash_attn.cute.interface import (
    _flash_attn_fwd,
    _flash_attn_bwd,
    flash_attn_varlen_func,
)
from test_score_mod import _generate_block_kvcache
from score_mod_definitions import (
    # TensorSSA-based score mods
    score_mod_alibi,
    score_mod_batch_bias,
    score_mod_block_diagonal,
    score_mod_causal,
    score_mod_causal_v2,
    score_mod_debug_global_idx,
    score_mod_dual_buffer,
    score_mod_global_kv_bias,
    score_mod_global_logical_rel_plus_kv_bias,
    score_mod_global_q_and_kv_bias,
    score_mod_global_q_bias,
    score_mod_global_rel_plus_kv_bias,
    score_mod_identity,
    score_mod_rel_bias,
    score_mod_rel_bias_x2,
    score_mod_sliding_window,
    score_mod_stress_complex_arithmetic,
    score_mod_stress_conditional_mask,
    score_mod_stress_global_offset,
    score_mod_stress_multi_buffer,
    score_mod_stress_xor_pattern,
    score_mod_times_two,
)  # isort: split
from score_mod_definitions import (
    score_mod_identity_vectorized,
    score_mod_causal_vectorized,
    score_mod_rel_bias as score_mod_rel_bias_vectorized,
    score_mod_rel_bias_x2_vectorized,
    score_mod_times_two_vectorized,
    score_mod_alibi_vectorized,
    score_mod_batch_bias_vectorized,
    score_mod_dual_buffer_vectorized,
)  # isort: split
from score_mod_definitions import (
    # Eager (torch) reference score mods
    identity_eager,
    causal_eager,
    rel_bias_eager,
    rel_bias_x2_eager,
    times_two_eager,
    alibi_eager,
    sliding_window_eager,
    block_diagonal_eager,
    causal_v2_eager,
    batch_bias_factory,
    dual_buffer_factory,
    packed_kv_bias_factory,
    packed_q_bias_factory,
    packed_rel_plus_kv_bias_factory,
    packed_q_and_kv_bias_factory,
    packed_logical_rel_plus_kv_bias_factory,
    stress_complex_arithmetic_factory,
    stress_conditional_mask_factory,
    stress_multi_buffer_factory,
    stress_global_offset_factory,
    stress_xor_pattern_factory,
    debug_global_idx_factory,
)  # isort: split
from score_mod_definitions import (
    # Forward + backward score mods used in bwd tests
    score_mod_squared,
    score_mod_scaled_squared,
    score_mod_rel_bias_clamped,
    score_mod_bwd_identity,
    score_mod_bwd_times_two,
    score_mod_bwd_rel_bias_clamped,
    score_mod_bwd_causal,
    score_mod_bwd_squared,
    score_mod_bwd_dual_buffer,
    score_mod_bwd_scaled_squared,
    squared_eager,
    rel_bias_clamped_eager,
    scaled_squared_factory,
)

IS_SM90 = torch.cuda.get_device_capability()[0] == 9
IS_SM100 = torch.cuda.get_device_capability()[0] == 10

# =============================================================================
# Test pairs
# =============================================================================

# (cute_score_mod, eager_factory_or_fn, aux_type)
# aux_type: None, "batch", "dual_buffer"
# All score_mods use 7-arg signature: (tSrS_ssa, b_idx, h_idx, q_idx, kv_idx, seqlen_info, aux_tensors)
TEST_PAIRS_NO_GLOBAL = [
    (score_mod_identity, identity_eager, None),
    (score_mod_causal, causal_eager, None),
    (score_mod_rel_bias, rel_bias_eager, None),
    (score_mod_rel_bias_x2, rel_bias_x2_eager, None),
    (score_mod_times_two, times_two_eager, None),
    (score_mod_alibi, alibi_eager, None),
    (score_mod_sliding_window, sliding_window_eager, None),
    (score_mod_block_diagonal, block_diagonal_eager, None),
    (score_mod_causal_v2, causal_v2_eager, None),
    (score_mod_batch_bias, batch_bias_factory, "batch"),
    (score_mod_dual_buffer, dual_buffer_factory, "dual_buffer"),
]

# Test pairs to compare vectorized score_mods: (cute_jit_function, cute_jit_function_vectorized)
TEST_PAIRS_VECTORIZED_NO_GLOBAL = [
    (score_mod_identity, score_mod_identity_vectorized, None),
    (score_mod_causal, score_mod_causal_vectorized, None),
    (score_mod_rel_bias, score_mod_rel_bias_vectorized, None),
    (score_mod_rel_bias_x2, score_mod_rel_bias_x2_vectorized, None),
    (score_mod_times_two, score_mod_times_two_vectorized, None),
    (score_mod_alibi, score_mod_alibi_vectorized, None),
    (score_mod_batch_bias, score_mod_batch_bias_vectorized, "batch"),
    (score_mod_dual_buffer, score_mod_dual_buffer_vectorized, "dual_buffer"),
]
# (cute_score_mod, eager_factory, aux_type, requires_global)
# aux_type: "kv", "q", "q_and_kv", "q_concat", "kv_with_cu", "multi_buffer"
# requires_global: "q" (needs varlen_q), "kv" (needs varlen_k), "both" (needs both)
# All score_mods use 7-arg signature and compute global indices from seqlen_info
TEST_PAIRS_WITH_GLOBAL = [
    (score_mod_global_kv_bias, packed_kv_bias_factory, "kv", "kv"),
    (score_mod_global_q_bias, packed_q_bias_factory, "q", "q"),
    (score_mod_global_rel_plus_kv_bias, packed_rel_plus_kv_bias_factory, "kv", "kv"),
    (score_mod_global_q_and_kv_bias, packed_q_and_kv_bias_factory, "q_and_kv", "both"),
    (
        score_mod_global_logical_rel_plus_kv_bias,
        packed_logical_rel_plus_kv_bias_factory,
        "kv",
        "kv",
    ),
    (
        score_mod_stress_complex_arithmetic,
        stress_complex_arithmetic_factory,
        "q_concat",
        "q",
    ),
    (
        score_mod_stress_conditional_mask,
        stress_conditional_mask_factory,
        "kv_with_cu",
        "both",
    ),
    (
        score_mod_stress_multi_buffer,
        stress_multi_buffer_factory,
        "multi_buffer",
        "both",
    ),
    (score_mod_stress_global_offset, stress_global_offset_factory, "kv", "kv"),
    (score_mod_stress_xor_pattern, stress_xor_pattern_factory, "kv_with_cu", "kv"),
    (score_mod_debug_global_idx, debug_global_idx_factory, "kv", "kv"),
]

SEQLEN_CONFIGS = [
    ([1], [1]),
    ([1, 1], [1, 1]),
    ([2, 3], [2, 3]),
    ([8, 16], [8, 16]),
    ([32, 32], [32, 32]),
    ([64, 128], [64, 128]),
    ([64, 56, 128], [64, 56, 128]),
    ([256, 512], [256, 512]),
    ([113, 203], [113, 203]),
    ([239, 1], [239, 1]),
    ([64], [64]),
    ([128, 128], [128, 128]),
    ([32, 32, 32, 32], [32, 32, 32, 32]),
    ([16, 32, 64, 128, 256], [16, 32, 64, 128, 256]),
    ([1, 1024], [1, 1024]),
    ([1024, 1], [1024, 1]),
    ([1, 256, 1], [1, 256, 1]),
    ([256, 1, 256], [256, 1, 256]),
    ([17, 33, 65], [17, 33, 65]),
    ([64, 128], [32, 64]),
    ([100, 100], [50, 50]),
    ([256, 512, 256], [128, 256, 128]),
    ([2, 1], [16384, 32 * 1024]),
    ([1, 1], [128 * 1024] * 2),
    ([2, 1], [8192, 8192]),
    ([1, 3], [8192, 8192]),
    ([3, 3], [8192, 8192]),
    ([128, 128], [8192, 8192]),
    ([2, 2, 2], [8 * 1024] * 3),
    ([2, 1], [1024 * 32, 16384]),
    ([1, 2], [1024 * 32, 16384]),
    ([1, 1, 1], [128 * 1024] * 3),
    ([1, 1, 1], [256 * 1024] * 3),
]

VEC_SIZES_TO_CHECK_EQUALITY = [1, 2, 4] if IS_SM100 else [1, 2]

# =============================================================================
# Helper functions
# =============================================================================


def run_cute_flash(
    q,
    k,
    v,
    score_mod,
    aux_tensors=None,
    pack_gqa=False,
    cu_seqlens_q=None,
    cu_seqlens_k=None,
    page_table=None,
    seqused_k=None,
):
    """Run CuTE flash attention."""
    if cu_seqlens_q is not None or cu_seqlens_k is not None:
        out = torch.empty_like(q)
        _flash_attn_fwd(
            q,
            k,
            v,
            cu_seqlens_q=cu_seqlens_q,
            cu_seqlens_k=cu_seqlens_k,
            seqused_k=seqused_k,
            page_table=page_table,
            return_lse=True,
            score_mod=score_mod,
            out=out,
            lse=None,
            aux_tensors=aux_tensors,
            pack_gqa=pack_gqa,
        )
        return out

    out = torch.empty_like(q)
    _flash_attn_fwd(
        q,
        k,
        v,
        seqused_k=seqused_k,
        page_table=page_table,
        return_lse=True,
        score_mod=score_mod,
        out=out,
        lse=None,
        aux_tensors=aux_tensors,
        pack_gqa=pack_gqa,
    )
    return out


def run_flex_varlen_ref(q, k, v, cu_seqlens_q, cu_seqlens_k, score_mod, dtype=None):
    """Run flex_attention per-sequence for varlen reference."""
    if cu_seqlens_q is not None:
        num_batches = len(cu_seqlens_q) - 1
    else:
        num_batches = len(cu_seqlens_k) - 1

    results = []
    for i in range(num_batches):
        # Get Q slice
        if cu_seqlens_q is not None:
            q_slice = (
                q[cu_seqlens_q[i] : cu_seqlens_q[i + 1]].unsqueeze(0).transpose(1, 2)
            )
        else:
            q_slice = q[i : i + 1].transpose(1, 2)

        # Get K/V slices
        if cu_seqlens_k is not None:
            k_slice = (
                k[cu_seqlens_k[i] : cu_seqlens_k[i + 1]].unsqueeze(0).transpose(1, 2)
            )
            v_slice = (
                v[cu_seqlens_k[i] : cu_seqlens_k[i + 1]].unsqueeze(0).transpose(1, 2)
            )
        else:
            k_slice = k[i : i + 1].transpose(1, 2)
            v_slice = v[i : i + 1].transpose(1, 2)

        if dtype is not None:
            q_slice, k_slice, v_slice = (
                q_slice.to(dtype),
                k_slice.to(dtype),
                v_slice.to(dtype),
            )

        def wrapped_mod(score, b, h, q_idx, kv_idx):
            return score_mod(score, i, h, q_idx, kv_idx)

        out = flex_attention(
            q_slice,
            k_slice,
            v_slice,
            score_mod=wrapped_mod,
            enable_gqa=q_slice.shape[1] != k_slice.shape[1],
        )
        results.append(out.transpose(1, 2).squeeze(0))

    return torch.cat(results, dim=0)


def setup_tensors(seqlens_q, seqlens_k, varlen_q, varlen_k, num_heads, head_dim, dtype):
    """Create Q, K, V tensors and cu_seqlens based on varlen flags."""
    batch_size = len(seqlens_q)

    if varlen_q:
        total_q = sum(seqlens_q)
        q = torch.randn(total_q, num_heads, head_dim, device="cuda", dtype=dtype)
        cu_seqlens_q = torch.tensor(
            [0] + list(torch.tensor(seqlens_q).cumsum(0).tolist()),
            device="cuda",
            dtype=torch.int32,
        )
    else:
        seqlen_q = seqlens_q[0]  # All sequences have the same length for non-varlen
        q = torch.randn(
            batch_size, seqlen_q, num_heads, head_dim, device="cuda", dtype=dtype
        )
        cu_seqlens_q = None

    if varlen_k:
        total_k = sum(seqlens_k)
        k = torch.randn(total_k, num_heads, head_dim, device="cuda", dtype=dtype)
        v = torch.randn(total_k, num_heads, head_dim, device="cuda", dtype=dtype)
        cu_seqlens_k = torch.tensor(
            [0] + list(torch.tensor(seqlens_k).cumsum(0).tolist()),
            device="cuda",
            dtype=torch.int32,
        )
    else:
        seqlen_k = seqlens_k[0]  # All sequences have the same length for non-varlen
        k = torch.randn(
            batch_size, seqlen_k, num_heads, head_dim, device="cuda", dtype=dtype
        )
        v = torch.randn(
            batch_size, seqlen_k, num_heads, head_dim, device="cuda", dtype=dtype
        )
        cu_seqlens_k = None

    return q, k, v, cu_seqlens_q, cu_seqlens_k


def prepare_ref_tensors(
    q, k, v, cu_seqlens_q, cu_seqlens_k, varlen_q, varlen_k, batch_size, seqlens_q
):
    """Prepare tensors for flex_attention reference (handle mixed varlen formats)."""
    num_heads = q.shape[1] if varlen_q else q.shape[2]

    if not varlen_q and varlen_k:
        seqlen_q = q.shape[1]
        q_packed = q.reshape(-1, num_heads, q.shape[-1])
        ref_cu_seqlens_q = torch.tensor(
            [seqlen_q * i for i in range(batch_size + 1)],
            device="cuda",
            dtype=torch.int32,
        )
        return q_packed, k, v, ref_cu_seqlens_q, cu_seqlens_k

    if varlen_q and not varlen_k:
        return q, k, v, cu_seqlens_q, None

    return q, k, v, cu_seqlens_q, cu_seqlens_k


def check_results(
    out_cute,
    out_ref_fp32,
    out_pt,
    test_name,
    rtol=2,
    extra_atol=1e-4,
    seqlens_q=None,
    cu_seqlens_q=None,
):
    """Compare CuTE output against references."""
    assert not torch.isnan(out_cute).any(), f"{test_name}: NaN in output"
    assert torch.isfinite(out_cute).all(), f"{test_name}: Inf in output"

    varlen_q = cu_seqlens_q is not None

    if varlen_q:
        # Unpack and compare per-sequence
        assert seqlens_q is not None, "varlen_q requires use of seqlens_q"
        num_seqs = len(seqlens_q)
        max_cute_error = 0.0
        max_pt_error = 0.0

        for i in range(num_seqs):
            # Extract sequences using cu_seqlens (all outputs are in packed format)
            start_q = cu_seqlens_q[i]
            end_q = cu_seqlens_q[i + 1]
            cute_seq = out_cute[start_q:end_q]
            ref_seq = out_ref_fp32[start_q:end_q]
            pt_seq = out_pt[start_q:end_q]

            max_cute_error = max(
                max_cute_error, (cute_seq - ref_seq).abs().max().item()
            )
            max_pt_error = max(max_pt_error, (pt_seq - ref_seq).abs().max().item())

        cute_error = max_cute_error
        pt_error = max_pt_error
    else:
        # Direct comparison
        pt_error = (out_pt - out_ref_fp32).abs().max().item()
        cute_error = (out_cute - out_ref_fp32).abs().max().item()

    fwd_atol = 2 * (out_ref_fp32 + 0.3 - 0.3 - out_ref_fp32).abs().max().item()

    print(f"\n{test_name}:")
    print(f"  PyTorch vs FP32 ref: {pt_error:.2e}")
    print(f"  CuTE vs FP32 ref: {cute_error:.2e}")

    tol = rtol * pt_error + fwd_atol + extra_atol
    assert cute_error <= tol, (
        f"{test_name}: CuTE error {cute_error:.2e} exceeds tolerance {tol:.2e}"
    )


# =============================================================================
# Tests
# =============================================================================


@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("varlen_q", [True, False])
@pytest.mark.parametrize("varlen_k", [True, False])
@pytest.mark.parametrize("qhead_per_kvhead,num_kv_heads", [(4, 2)])
@pytest.mark.parametrize("seqlens_q,seqlens_k", SEQLEN_CONFIGS)
@pytest.mark.parametrize("score_mod_tuple", TEST_PAIRS_NO_GLOBAL)
def test_varlen_with_score_mod(
    seqlens_q,
    seqlens_k,
    varlen_q,
    varlen_k,
    qhead_per_kvhead,
    num_kv_heads,
    dtype,
    score_mod_tuple,
):
    """Test varlen attention with score_mod functions that don't use global indices.

    Covers: both varlen, varlen Q only, varlen K only.
    Skips: neither varlen
    """
    if not varlen_q and not varlen_k:
        pytest.skip(
            "At least one of varlen_q or varlen_k must be True for varlen tests"
        )

    # For non-varlen dimension, all sequences must have same length
    if not varlen_q:
        seqlens_q = [seqlens_q[0]] * len(seqlens_q)
    if not varlen_k:
        seqlens_k = [seqlens_k[0]] * len(seqlens_k)

    torch.random.manual_seed(42)
    cute_score_mod, eager_factory, aux_type = score_mod_tuple

    num_heads = num_kv_heads * qhead_per_kvhead
    pack_gqa = qhead_per_kvhead > 1
    head_dim = 128
    batch_size = len(seqlens_q)

    q, k, v, cu_seqlens_q, cu_seqlens_k = setup_tensors(
        seqlens_q, seqlens_k, varlen_q, varlen_k, num_heads, head_dim, dtype
    )

    if pack_gqa:
        if varlen_k:
            k = k[:, :num_kv_heads, :].clone()
            v = v[:, :num_kv_heads, :].clone()
        else:
            k = k[:, :, :num_kv_heads, :].clone()
            v = v[:, :, :num_kv_heads, :].clone()

    aux_tensors = None
    if aux_type == "batch":
        bias = torch.zeros(batch_size, device="cuda", dtype=dtype) * 0.1
        aux_tensors = [bias]
        eager_score_mod = eager_factory(bias)
    elif aux_type == "dual_buffer":
        seqlen_q = seqlens_q[0] if not varlen_q else max(seqlens_q)
        head_bias = torch.randn(num_heads, device="cuda", dtype=dtype) * 0.2
        pos_bias = torch.arange(seqlen_q, device="cuda", dtype=dtype) * 0.01
        aux_tensors = [head_bias, pos_bias]
        eager_score_mod = eager_factory(head_bias, pos_bias)
    else:
        eager_score_mod = eager_factory

    # Prepare reference tensors
    q_ref, k_ref, v_ref, ref_cu_q, ref_cu_k = prepare_ref_tensors(
        q, k, v, cu_seqlens_q, cu_seqlens_k, varlen_q, varlen_k, batch_size, seqlens_q
    )

    out_ref_fp32 = run_flex_varlen_ref(
        q_ref, k_ref, v_ref, ref_cu_q, ref_cu_k, eager_score_mod, dtype=torch.float32
    )
    out_pt = run_flex_varlen_ref(
        q_ref, k_ref, v_ref, ref_cu_q, ref_cu_k, eager_score_mod, dtype=dtype
    )
    out_cute = run_cute_flash(
        q,
        k,
        v,
        cute_score_mod,
        aux_tensors=aux_tensors,
        pack_gqa=pack_gqa,
        cu_seqlens_q=cu_seqlens_q,
        cu_seqlens_k=cu_seqlens_k,
    )

    if not varlen_q and varlen_k:
        seqlen_q = q.shape[1]
        out_ref_fp32 = out_ref_fp32.reshape(batch_size, seqlen_q, num_heads, head_dim)
        out_pt = out_pt.reshape(batch_size, seqlen_q, num_heads, head_dim)

    assert out_cute.shape == out_ref_fp32.shape, (
        f"Shape mismatch: {out_cute.shape} vs {out_ref_fp32.shape}"
    )

    test_name = f"{cute_score_mod.__name__} (varlen_q={varlen_q}, varlen_k={varlen_k})"
    extra_atol = 2e-3
    check_results(
        out_cute,
        out_ref_fp32,
        out_pt,
        test_name,
        extra_atol=extra_atol,
        seqlens_q=seqlens_q if varlen_q else None,
        cu_seqlens_q=cu_seqlens_q if varlen_q else None,
    )


@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("varlen_q", [True, False])
@pytest.mark.parametrize("varlen_k", [True, False])
@pytest.mark.parametrize("qhead_per_kvhead,num_kv_heads", [(4, 2)])
@pytest.mark.parametrize("seqlens_q,seqlens_k", SEQLEN_CONFIGS)
@pytest.mark.parametrize("score_mod_vec_tuple", TEST_PAIRS_VECTORIZED_NO_GLOBAL)
def test_varlen_with_score_mod_vectorized(
    seqlens_q,
    seqlens_k,
    varlen_q,
    varlen_k,
    qhead_per_kvhead,
    num_kv_heads,
    dtype,
    score_mod_vec_tuple,
):
    """Tests equality between original and vectorized versions of score mods"""
    if not varlen_q and not varlen_k:
        pytest.skip(
            "At least one of varlen_q or varlen_k must be True for varlen tests"
        )

    # For non-varlen dimension, all sequences must have same length
    if not varlen_q:
        seqlens_q = [seqlens_q[0]] * len(seqlens_q)
    if not varlen_k:
        seqlens_k = [seqlens_k[0]] * len(seqlens_k)
    torch.random.manual_seed(42)
    cute_score_mod, cute_vectorized_score_mod, aux_type = score_mod_vec_tuple

    num_heads = num_kv_heads * qhead_per_kvhead
    pack_gqa = qhead_per_kvhead > 1
    head_dim = 128
    batch_size = len(seqlens_q)

    q, k, v, cu_seqlens_q, cu_seqlens_k = setup_tensors(
        seqlens_q, seqlens_k, varlen_q, varlen_k, num_heads, head_dim, dtype
    )
    aux_tensors = None
    if aux_type == "batch":
        bias = torch.zeros(batch_size, device="cuda", dtype=dtype) * 0.1
        aux_tensors = [bias]
    elif aux_type == "dual_buffer":
        seqlen_q = seqlens_q[0] if not varlen_q else max(seqlens_q)
        head_bias = torch.randn(num_heads, device="cuda", dtype=dtype) * 0.2
        pos_bias = torch.arange(seqlen_q, device="cuda", dtype=dtype) * 0.01
        aux_tensors = [head_bias, pos_bias]

    if pack_gqa:
        if varlen_k:
            k = k[:, :num_kv_heads, :].clone()
            v = v[:, :num_kv_heads, :].clone()
        else:
            k = k[:, :, :num_kv_heads, :].clone()
            v = v[:, :, :num_kv_heads, :].clone()

    out_ref = run_cute_flash(
        q,
        k,
        v,
        cute_score_mod,
        aux_tensors=aux_tensors,
        pack_gqa=pack_gqa,
        cu_seqlens_q=cu_seqlens_q,
        cu_seqlens_k=cu_seqlens_k,
    )

    for vec_size in VEC_SIZES_TO_CHECK_EQUALITY:
        cute_vectorized_score_mod.__vec_size__ = vec_size
        out = run_cute_flash(
            q,
            k,
            v,
            cute_vectorized_score_mod,
            aux_tensors=aux_tensors,
            pack_gqa=pack_gqa,
            cu_seqlens_q=cu_seqlens_q,
            cu_seqlens_k=cu_seqlens_k,
        )
        # Vectorized codegen reorders float ops vs the scalar path; softmax amplifies
        # the resulting 1-ulp score differences (measured max 4e-4 fp16 / 2.9e-3 bf16 on sm103).
        rtol, atol = (2e-3, 1e-3) if dtype == torch.float16 else (1.6e-2, 5e-3)
        assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)


@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("varlen_q", [True, False])
@pytest.mark.parametrize("varlen_k", [True, False])
@pytest.mark.parametrize("qhead_per_kvhead,num_kv_heads", [(1, 1), (4, 2)])
@pytest.mark.parametrize("seqlens_q,seqlens_k", SEQLEN_CONFIGS)
@pytest.mark.parametrize("score_mod_tuple", TEST_PAIRS_WITH_GLOBAL)
def test_varlen_with_global_idx_score_mod(
    seqlens_q,
    seqlens_k,
    varlen_q,
    varlen_k,
    qhead_per_kvhead,
    num_kv_heads,
    dtype,
    score_mod_tuple,
):
    """Test varlen attention with score_mod functions that use global indices.

    These score_mods compute q_idx_global and/or kv_idx_global from seqlen_info for packed tensor indexing.
    Skips tests where required global indices aren't available.
    """
    if not varlen_q and not varlen_k:
        pytest.skip(
            "At least one of varlen_q or varlen_k must be True for varlen tests"
        )

    cute_score_mod, eager_factory, aux_type, requires_global = score_mod_tuple

    # Skip if score_mod requires global indices we can't provide
    if requires_global == "q" and not varlen_q:
        pytest.skip(f"{cute_score_mod.__name__} requires varlen_q for q_idx_global")
    if requires_global == "kv" and not varlen_k:
        pytest.skip(f"{cute_score_mod.__name__} requires varlen_k for kv_idx_global")
    if requires_global == "both" and (not varlen_q or not varlen_k):
        pytest.skip(f"{cute_score_mod.__name__} requires both varlen_q and varlen_k")

    # For non-varlen dimension, all sequences must have same length
    if not varlen_q:
        seqlens_q = [seqlens_q[0]] * len(seqlens_q)
    if not varlen_k:
        seqlens_k = [seqlens_k[0]] * len(seqlens_k)

    torch.random.manual_seed(42)

    num_heads = num_kv_heads * qhead_per_kvhead
    pack_gqa = qhead_per_kvhead > 1
    head_dim = 128
    batch_size = len(seqlens_q)
    max_rel_pos = 512

    total_q = sum(seqlens_q)
    total_k = sum(seqlens_k)

    cu_seqlens_q = torch.tensor(
        [0] + list(torch.tensor(seqlens_q).cumsum(0).tolist()),
        device="cuda",
        dtype=torch.int32,
    )
    cu_seqlens_k = torch.tensor(
        [0] + list(torch.tensor(seqlens_k).cumsum(0).tolist()),
        device="cuda",
        dtype=torch.int32,
    )

    if varlen_q:
        q = torch.randn(total_q, num_heads, head_dim, device="cuda", dtype=dtype)
    else:
        seqlen_q = seqlens_q[0]
        q = torch.randn(
            batch_size, seqlen_q, num_heads, head_dim, device="cuda", dtype=dtype
        )

    if varlen_k:
        k = torch.randn(total_k, num_heads, head_dim, device="cuda", dtype=dtype)
        v = torch.randn(total_k, num_heads, head_dim, device="cuda", dtype=dtype)
    else:
        seqlen_k = seqlens_k[0]
        k = torch.randn(
            batch_size, seqlen_k, num_heads, head_dim, device="cuda", dtype=dtype
        )
        v = torch.randn(
            batch_size, seqlen_k, num_heads, head_dim, device="cuda", dtype=dtype
        )

    if pack_gqa:
        if varlen_k:
            k = k[:, :num_kv_heads, :].clone()
            v = v[:, :num_kv_heads, :].clone()
        else:
            k = k[:, :, :num_kv_heads, :].clone()
            v = v[:, :, :num_kv_heads, :].clone()

    # Setup aux tensors based on indexing type
    if aux_type == "kv":
        bias = torch.randn(total_k, device="cuda", dtype=dtype) * 0.1
        aux_tensors = [bias]
        eager_score_mod = eager_factory(bias, cu_seqlens_k)
    elif aux_type == "q":
        bias = torch.randn(total_q, device="cuda", dtype=dtype) * 0.1
        aux_tensors = [bias]
        eager_score_mod = eager_factory(bias, cu_seqlens_q)
    elif aux_type == "q_and_kv":
        q_bias = torch.randn(total_q, device="cuda", dtype=dtype) * 0.1
        kv_bias = torch.randn(total_k, device="cuda", dtype=dtype) * 0.1
        aux_tensors = [q_bias, kv_bias]
        eager_score_mod = eager_factory(q_bias, kv_bias, cu_seqlens_q, cu_seqlens_k)
    elif aux_type == "q_concat":
        bias = torch.randn(total_q, device="cuda", dtype=dtype) * 0.1
        aux_tensors = [bias]
        eager_score_mod = eager_factory(bias, cu_seqlens_q)
    elif aux_type == "kv_with_cu":
        kv_bias = torch.randn(total_k, device="cuda", dtype=dtype) * 0.1
        aux_tensors = [kv_bias]
        eager_score_mod = eager_factory(kv_bias, cu_seqlens_q, cu_seqlens_k)
    elif aux_type == "multi_buffer":
        batch_bias = torch.randn(batch_size, device="cuda", dtype=dtype) * 0.1
        head_scale = torch.randn(num_heads, device="cuda", dtype=dtype) * 0.1 + 1.0
        q_pos_bias = torch.randn(total_q, device="cuda", dtype=dtype) * 0.1
        kv_pos_bias = torch.randn(total_k, device="cuda", dtype=dtype) * 0.1
        rel_pos_scale = (
            torch.randn(max_rel_pos * 2 + 1, device="cuda", dtype=dtype) * 0.1
        )
        aux_tensors = [batch_bias, head_scale, q_pos_bias, kv_pos_bias, rel_pos_scale]
        eager_score_mod = eager_factory(
            batch_bias,
            head_scale,
            q_pos_bias,
            kv_pos_bias,
            rel_pos_scale,
            cu_seqlens_q,
            cu_seqlens_k,
            max_rel_pos,
        )
    else:
        raise ValueError(f"Unknown aux_type: {aux_type}")

    # Prepare reference tensors for flex_attention
    q_ref, k_ref, v_ref, ref_cu_q, ref_cu_k = prepare_ref_tensors(
        q, k, v, cu_seqlens_q, cu_seqlens_k, varlen_q, varlen_k, batch_size, seqlens_q
    )

    out_ref_fp32 = run_flex_varlen_ref(
        q_ref, k_ref, v_ref, ref_cu_q, ref_cu_k, eager_score_mod, dtype=torch.float32
    )
    out_pt = run_flex_varlen_ref(
        q_ref, k_ref, v_ref, ref_cu_q, ref_cu_k, eager_score_mod, dtype=dtype
    )

    kernel_cu_seqlens_q = cu_seqlens_q if varlen_q else None
    kernel_cu_seqlens_k = cu_seqlens_k if varlen_k else None
    out_cute = run_cute_flash(
        q,
        k,
        v,
        cute_score_mod,
        aux_tensors=aux_tensors,
        pack_gqa=pack_gqa,
        cu_seqlens_q=kernel_cu_seqlens_q,
        cu_seqlens_k=kernel_cu_seqlens_k,
    )

    if varlen_q:
        out_ref_final = out_ref_fp32
        out_pt_final = out_pt
        out_cute_final = out_cute
    else:
        seqlen_q = seqlens_q[0]
        out_ref_final = out_ref_fp32.reshape(batch_size, seqlen_q, num_heads, head_dim)
        out_pt_final = out_pt.reshape(batch_size, seqlen_q, num_heads, head_dim)
        out_cute_final = out_cute

    assert out_cute_final.shape == out_ref_final.shape, (
        f"Shape mismatch: {out_cute_final.shape} vs {out_ref_final.shape}"
    )

    test_name = f"{cute_score_mod.__name__} (varlen_q={varlen_q}, varlen_k={varlen_k}, {aux_type})"

    check_results(
        out_cute_final,
        out_ref_final,
        out_pt_final,
        test_name,
        extra_atol=1e-3,
        seqlens_q=seqlens_q if varlen_q else None,
        cu_seqlens_q=cu_seqlens_q if varlen_q else None,
    )


@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("page_size", [None, 128])
@pytest.mark.parametrize("varlen_q", [True, False])
@pytest.mark.parametrize("varlen_k", [True, False])
@pytest.mark.parametrize("qhead_per_kvhead,num_kv_heads", [(4, 2)])
@pytest.mark.parametrize("seqlens_q,seqlens_k", SEQLEN_CONFIGS)
@pytest.mark.parametrize("score_mod_tuple", TEST_PAIRS_NO_GLOBAL)
def test_varlen_score_mod_kvcache(
    seqlens_q,
    seqlens_k,
    varlen_q,
    varlen_k,
    qhead_per_kvhead,
    num_kv_heads,
    page_size,
    dtype,
    score_mod_tuple,
):
    """Test varlen attention with score_mod and paged KV cache."""
    if IS_SM90 and page_size is not None:
        pytest.xfail("paged KV not supported on SM90")

    if not varlen_q and not varlen_k:
        pytest.skip(
            "At least one of varlen_q or varlen_k must be True for varlen tests"
        )

    if page_size is not None and varlen_k:
        pytest.skip("Paged KV requires batched (non-varlen) K")

    if not varlen_q:
        seqlens_q = [seqlens_q[0]] * len(seqlens_q)
    if not varlen_k:
        seqlens_k = [seqlens_k[0]] * len(seqlens_k)

    # Skip if page_size doesn't divide seqlens evenly (for simplicity)
    if page_size is not None and not varlen_k:
        if seqlens_k[0] % page_size != 0:
            pytest.skip("page_size must divide seqlen_k")

    torch.random.manual_seed(42)
    cute_score_mod, eager_factory, aux_type = score_mod_tuple

    num_heads = num_kv_heads * qhead_per_kvhead
    pack_gqa = qhead_per_kvhead > 1
    head_dim = 128
    batch_size = len(seqlens_q)
    device = "cuda"

    # Setup tensors
    q, k, v, cu_seqlens_q, cu_seqlens_k = setup_tensors(
        seqlens_q, seqlens_k, varlen_q, varlen_k, num_heads, head_dim, dtype
    )

    if pack_gqa:
        if varlen_k:
            k = k[:, :num_kv_heads, :].clone()
            v = v[:, :num_kv_heads, :].clone()
        else:
            k = k[:, :, :num_kv_heads, :].clone()
            v = v[:, :, :num_kv_heads, :].clone()

    page_table = None
    k_cache_paged = None
    v_cache_paged = None
    k_cache = k
    v_cache = v

    if page_size is not None:
        seqlen_k = seqlens_k[0]
        (
            k_cache_bhsd,
            v_cache_bhsd,
            page_table,
            k_cache_paged,
            v_cache_paged,
            num_blocks,
        ) = _generate_block_kvcache(
            seqlen_k, page_size, batch_size, num_kv_heads, head_dim, device, dtype
        )
        k_cache = k_cache_bhsd.transpose(1, 2)  # BHSD -> BSHD
        v_cache = v_cache_bhsd.transpose(1, 2)
        seqused_k = torch.tensor(seqlens_k, dtype=torch.int32, device=device)
    else:
        seqused_k = None

    # Setup aux tensors and eager score_mod
    aux_tensors = None
    if aux_type == "batch":
        bias = torch.zeros(batch_size, device=device, dtype=dtype) * 0.1
        aux_tensors = [bias]
        eager_score_mod = eager_factory(bias)
    elif aux_type == "dual_buffer":
        seqlen_q = seqlens_q[0] if not varlen_q else max(seqlens_q)
        head_bias = torch.randn(num_heads, device=device, dtype=dtype) * 0.2
        pos_bias = torch.arange(seqlen_q, device=device, dtype=dtype) * 0.01
        aux_tensors = [head_bias, pos_bias]
        eager_score_mod = eager_factory(head_bias, pos_bias)
    else:
        eager_score_mod = eager_factory

    # Prepare reference tensors
    q_ref, k_ref, v_ref, ref_cu_q, ref_cu_k = prepare_ref_tensors(
        q,
        k_cache,
        v_cache,
        cu_seqlens_q,
        cu_seqlens_k,
        varlen_q,
        varlen_k,
        batch_size,
        seqlens_q,
    )

    out_ref_fp32 = run_flex_varlen_ref(
        q_ref, k_ref, v_ref, ref_cu_q, ref_cu_k, eager_score_mod, dtype=torch.float32
    )
    out_pt = run_flex_varlen_ref(
        q_ref, k_ref, v_ref, ref_cu_q, ref_cu_k, eager_score_mod, dtype=dtype
    )

    k_input = k_cache_paged if page_size is not None else k_cache
    v_input = v_cache_paged if page_size is not None else v_cache

    out_cute = run_cute_flash(
        q,
        k_input,
        v_input,
        cute_score_mod,
        aux_tensors=aux_tensors,
        pack_gqa=pack_gqa,
        cu_seqlens_q=cu_seqlens_q if varlen_q else None,
        cu_seqlens_k=cu_seqlens_k if (varlen_k and page_size is None) else None,
        page_table=page_table if page_size is not None else None,
        seqused_k=seqused_k if page_size is not None else None,
    )

    if not varlen_q and varlen_k:
        seqlen_q = q.shape[1]
        out_ref_fp32 = out_ref_fp32.reshape(batch_size, seqlen_q, num_heads, head_dim)
        out_pt = out_pt.reshape(batch_size, seqlen_q, num_heads, head_dim)

    assert out_cute.shape == out_ref_fp32.shape, (
        f"Shape mismatch: {out_cute.shape} vs {out_ref_fp32.shape}"
    )

    test_name = f"{cute_score_mod.__name__} (varlen_q={varlen_q}, varlen_k={varlen_k}, paged={page_size is not None})"
    extra_atol = 2e-3
    check_results(
        out_cute,
        out_ref_fp32,
        out_pt,
        test_name,
        extra_atol=extra_atol,
        seqlens_q=seqlens_q if varlen_q else None,
        cu_seqlens_q=cu_seqlens_q if varlen_q else None,
    )


@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("page_size", [None, 128])
@pytest.mark.parametrize("varlen_q", [True, False])
@pytest.mark.parametrize("varlen_k", [True, False])
@pytest.mark.parametrize("qhead_per_kvhead,num_kv_heads", [(1, 1), (4, 2)])
@pytest.mark.parametrize("seqlens_q,seqlens_k", SEQLEN_CONFIGS)
@pytest.mark.parametrize("score_mod_tuple", TEST_PAIRS_WITH_GLOBAL)
def test_varlen_score_mod_with_paged_kvcache_global(
    seqlens_q,
    seqlens_k,
    varlen_q,
    varlen_k,
    qhead_per_kvhead,
    num_kv_heads,
    page_size,
    dtype,
    score_mod_tuple,
):
    """Test varlen attention with global idx score_mod and paged KV cache."""
    if IS_SM90 and page_size is not None:
        pytest.xfail("paged KV not supported on SM90")

    if page_size is not None and varlen_k:
        pytest.skip("Paged KV cache requires batched (non-varlen) K")

    if not varlen_q and not varlen_k:
        pytest.skip(
            "At least one of varlen_q or varlen_k must be True for varlen tests"
        )

    if not varlen_q:
        seqlens_q = [seqlens_q[0]] * len(seqlens_q)
    if not varlen_k:
        seqlens_k = [seqlens_k[0]] * len(seqlens_k)

    if page_size is not None and not varlen_k:
        if seqlens_k[0] % page_size != 0:
            pytest.skip("page_size must divide seqlen_k")

    cute_score_mod, eager_factory, aux_type, requires_global = score_mod_tuple

    if requires_global == "q" and not varlen_q:
        pytest.skip(f"{cute_score_mod.__name__} requires varlen_q for q_idx_global")
    if requires_global == "kv" and not varlen_k:
        pytest.skip(f"{cute_score_mod.__name__} requires varlen_k for kv_idx_global")
    if requires_global == "both" and (not varlen_q or not varlen_k):
        pytest.skip(f"{cute_score_mod.__name__} requires both varlen_q and varlen_k")

    torch.random.manual_seed(42)

    num_heads = num_kv_heads * qhead_per_kvhead
    pack_gqa = qhead_per_kvhead > 1
    head_dim = 128
    batch_size = len(seqlens_q)
    max_rel_pos = 512
    device = "cuda"

    total_q = sum(seqlens_q)
    total_k = sum(seqlens_k)

    cu_seqlens_q = torch.tensor(
        [0] + list(torch.tensor(seqlens_q).cumsum(0).tolist()),
        device=device,
        dtype=torch.int32,
    )
    cu_seqlens_k = torch.tensor(
        [0] + list(torch.tensor(seqlens_k).cumsum(0).tolist()),
        device=device,
        dtype=torch.int32,
    )
    cu_seqlens_k_for_kernel = cu_seqlens_k if varlen_k else None

    q = torch.randn(total_q, num_heads, head_dim, device=device, dtype=dtype)
    if varlen_k:
        k = torch.randn(total_k, num_heads, head_dim, device=device, dtype=dtype)
        v = torch.randn(total_k, num_heads, head_dim, device=device, dtype=dtype)
    else:
        seqlen_k = seqlens_k[0]
        k = torch.randn(
            batch_size, seqlen_k, num_heads, head_dim, device=device, dtype=dtype
        )
        v = torch.randn(
            batch_size, seqlen_k, num_heads, head_dim, device=device, dtype=dtype
        )

    if pack_gqa:
        if varlen_k:
            k = k[:, :num_kv_heads, :].clone()
            v = v[:, :num_kv_heads, :].clone()
        else:
            k = k[:, :, :num_kv_heads, :].clone()
            v = v[:, :, :num_kv_heads, :].clone()

    page_table = None
    k_cache_paged = None
    v_cache_paged = None
    k_cache = k
    v_cache = v

    if page_size is not None:
        seqlen_k = seqlens_k[0]
        (
            k_cache_bhsd,
            v_cache_bhsd,
            page_table,
            k_cache_paged,
            v_cache_paged,
            num_blocks,
        ) = _generate_block_kvcache(
            seqlen_k, page_size, batch_size, num_kv_heads, head_dim, device, dtype
        )
        k_cache = k_cache_bhsd.transpose(1, 2)  # BHSD -> BSHD
        v_cache = v_cache_bhsd.transpose(1, 2)
        seqused_k = torch.tensor(seqlens_k, dtype=torch.int32, device=device)
    else:
        seqused_k = None

    if aux_type == "kv":
        bias = torch.randn(total_k, device=device, dtype=dtype) * 0.1
        aux_tensors = [bias]
        eager_score_mod = eager_factory(bias, cu_seqlens_k)
    elif aux_type == "q":
        bias = torch.randn(total_q, device=device, dtype=dtype) * 0.1
        aux_tensors = [bias]
        eager_score_mod = eager_factory(bias, cu_seqlens_q)
    elif aux_type == "q_and_kv":
        q_bias = torch.randn(total_q, device=device, dtype=dtype) * 0.1
        kv_bias = torch.randn(total_k, device=device, dtype=dtype) * 0.1
        aux_tensors = [q_bias, kv_bias]
        eager_score_mod = eager_factory(q_bias, kv_bias, cu_seqlens_q, cu_seqlens_k)
    elif aux_type == "q_concat":
        bias = torch.randn(total_q, device=device, dtype=dtype) * 0.1
        aux_tensors = [bias]
        eager_score_mod = eager_factory(bias, cu_seqlens_q)
    elif aux_type == "kv_with_cu":
        kv_bias = torch.randn(total_k, device=device, dtype=dtype) * 0.1
        aux_tensors = [kv_bias]
        eager_score_mod = eager_factory(kv_bias, cu_seqlens_q, cu_seqlens_k)
    elif aux_type == "multi_buffer":
        batch_bias = torch.randn(batch_size, device=device, dtype=dtype) * 0.1
        head_scale = torch.randn(num_heads, device=device, dtype=dtype) * 0.1 + 1.0
        q_pos_bias = torch.randn(total_q, device=device, dtype=dtype) * 0.1
        kv_pos_bias = torch.randn(total_k, device=device, dtype=dtype) * 0.1
        rel_pos_scale = (
            torch.randn(max_rel_pos * 2 + 1, device=device, dtype=dtype) * 0.1
        )
        aux_tensors = [batch_bias, head_scale, q_pos_bias, kv_pos_bias, rel_pos_scale]
        eager_score_mod = eager_factory(
            batch_bias,
            head_scale,
            q_pos_bias,
            kv_pos_bias,
            rel_pos_scale,
            cu_seqlens_q,
            cu_seqlens_k,
            max_rel_pos,
        )
    else:
        raise ValueError(f"Unknown aux_type: {aux_type}")

    q_ref, k_ref, v_ref, ref_cu_q, ref_cu_k = prepare_ref_tensors(
        q,
        k_cache,
        v_cache,
        cu_seqlens_q,
        cu_seqlens_k,
        True,
        varlen_k,
        batch_size,
        seqlens_q,
    )

    out_ref_fp32 = run_flex_varlen_ref(
        q_ref, k_ref, v_ref, ref_cu_q, ref_cu_k, eager_score_mod, dtype=torch.float32
    )
    out_pt = run_flex_varlen_ref(
        q_ref, k_ref, v_ref, ref_cu_q, ref_cu_k, eager_score_mod, dtype=dtype
    )

    # Run CuTE
    k_input = k_cache_paged if page_size is not None else k_cache
    v_input = v_cache_paged if page_size is not None else v_cache

    out_cute = torch.empty_like(q)
    _flash_attn_fwd(
        q,
        k_input,
        v_input,
        cu_seqlens_q=cu_seqlens_q,
        cu_seqlens_k=cu_seqlens_k_for_kernel if page_size is None else None,
        seqused_k=seqused_k if page_size is not None else None,
        page_table=page_table,
        return_lse=True,
        score_mod=cute_score_mod,
        out=out_cute,
        lse=None,
        aux_tensors=aux_tensors,
        pack_gqa=pack_gqa,
    )

    assert out_cute.shape == out_ref_fp32.shape, (
        f"Shape mismatch: {out_cute.shape} vs {out_ref_fp32.shape}"
    )

    test_name = f"{cute_score_mod.__name__} (paged={page_size is not None}, {aux_type})"
    check_results(
        out_cute,
        out_ref_fp32,
        out_pt,
        test_name,
        extra_atol=1e-3,
        seqlens_q=seqlens_q,
        cu_seqlens_q=cu_seqlens_q,
    )


# =============================================================================
# Backward tests
# =============================================================================

# (cute_fwd, cute_bwd, eager_ref_or_factory, aux_type)
# aux_type: None or "dual_buffer"
BWD_TEST_PAIRS = [
    (score_mod_times_two, score_mod_bwd_times_two, times_two_eager, None),
    (score_mod_rel_bias_clamped, score_mod_bwd_rel_bias_clamped, rel_bias_clamped_eager, None),
    (score_mod_squared, score_mod_bwd_squared, squared_eager, None),
    (score_mod_causal, score_mod_bwd_causal, causal_eager, None),
    (
        score_mod_dual_buffer,
        score_mod_bwd_dual_buffer,
        dual_buffer_factory,
        "dual_buffer",
    ),
]

# (cute_fwd, cute_bwd, eager_factory, aux_type, requires_global)
BWD_TEST_PAIRS_WITH_GLOBAL = [
    (
        score_mod_scaled_squared,
        score_mod_bwd_scaled_squared,
        scaled_squared_factory,
        "kv",
        "kv",
    ),
]

BWD_SEQLEN_CONFIGS = SEQLEN_CONFIGS


def run_cute_flash_bwd_varlen(
    q,
    k,
    v,
    cute_score_mod,
    cute_score_mod_bwd,
    cu_seqlens_q,
    cu_seqlens_k,
    aux_tensors=None,
    pack_gqa=False,
    use_autograd=True,
):
    """Forward + backward with score_mod for packed varlen inputs.

    Mirrors run_cute_flash_bwd in test_score_mod.py. use_autograd=True drives
    flash_attn_varlen_func + torch.autograd.grad; False calls the fwd/bwd entry
    points directly.
    """
    if use_autograd:
        q = q.detach().requires_grad_(True)
        k = k.detach().requires_grad_(True)
        v = v.detach().requires_grad_(True)
        out, lse = flash_attn_varlen_func(
            q,
            k,
            v,
            cu_seqlens_q=cu_seqlens_q,
            cu_seqlens_k=cu_seqlens_k,
            score_mod=cute_score_mod,
            score_mod_bwd=cute_score_mod_bwd,
            aux_tensors=aux_tensors,
            pack_gqa=pack_gqa,
            return_lse=True,
        )
        grad_out = torch.randn_like(out)
        dq, dk, dv = torch.autograd.grad(out, (q, k, v), grad_out)
        return out, grad_out, dq, dk, dv

    out, lse, *_ = _flash_attn_fwd(
        q,
        k,
        v,
        cu_seqlens_q=cu_seqlens_q,
        cu_seqlens_k=cu_seqlens_k,
        return_lse=True,
        score_mod=cute_score_mod,
        aux_tensors=aux_tensors,
        pack_gqa=pack_gqa,
    )
    grad_out = torch.randn_like(out)
    dq, dk, dv, *_ = _flash_attn_bwd(
        q,
        k,
        v,
        out,
        grad_out,
        lse,
        cu_seqlens_q=cu_seqlens_q,
        cu_seqlens_k=cu_seqlens_k,
        score_mod=cute_score_mod,
        score_mod_bwd=cute_score_mod_bwd,
        aux_tensors=aux_tensors,
        pack_gqa=pack_gqa,
    )
    return out, grad_out, dq, dk, dv


def run_flex_varlen_ref_bwd(
    q, k, v, cu_seqlens_q, cu_seqlens_k, eager_score_mod, grad_out, dtype=None
):
    """Per-sequence flex_attention fwd+bwd; concatenates results in packed layout.

    Uses eager flex_attention (not torch.compile). The compiled path caches
    artifacts globally and silently produces wrong fp32 gradients across calls
    with different seqlens, so the reference must stay uncompiled.
    """
    num_batches = len(cu_seqlens_q) - 1
    out_chunks, dq_chunks, dk_chunks, dv_chunks = [], [], [], []

    for i in range(num_batches):
        q_slice = q[cu_seqlens_q[i] : cu_seqlens_q[i + 1]].unsqueeze(0).transpose(1, 2)
        k_slice = k[cu_seqlens_k[i] : cu_seqlens_k[i + 1]].unsqueeze(0).transpose(1, 2)
        v_slice = v[cu_seqlens_k[i] : cu_seqlens_k[i + 1]].unsqueeze(0).transpose(1, 2)
        go_slice = (
            grad_out[cu_seqlens_q[i] : cu_seqlens_q[i + 1]].unsqueeze(0).transpose(1, 2)
        )

        if dtype is not None:
            q_slice = q_slice.to(dtype)
            k_slice = k_slice.to(dtype)
            v_slice = v_slice.to(dtype)
            go_slice = go_slice.to(dtype)

        q_slice = q_slice.detach().requires_grad_(True)
        k_slice = k_slice.detach().requires_grad_(True)
        v_slice = v_slice.detach().requires_grad_(True)

        def wrapped_mod(score, b, h, q_idx, kv_idx, _i=i):
            return eager_score_mod(score, _i, h, q_idx, kv_idx)

        out = flex_attention(
            q_slice,
            k_slice,
            v_slice,
            score_mod=wrapped_mod,
            enable_gqa=q_slice.shape[1] != k_slice.shape[1],
        )
        dq, dk, dv = torch.autograd.grad(out, (q_slice, k_slice, v_slice), go_slice)

        out_chunks.append(out.transpose(1, 2).squeeze(0))
        dq_chunks.append(dq.transpose(1, 2).squeeze(0))
        dk_chunks.append(dk.transpose(1, 2).squeeze(0))
        dv_chunks.append(dv.transpose(1, 2).squeeze(0))

    return (
        torch.cat(out_chunks, dim=0),
        torch.cat(dq_chunks, dim=0),
        torch.cat(dk_chunks, dim=0),
        torch.cat(dv_chunks, dim=0),
    )


def _check_grad(name, cute_grad, ref_fp32, pt_grad, dtype, rtol=3, extra=1e-3):
    import os
    if os.environ.get("DUMP_GRADS"):
        torch.save({"cute": cute_grad.detach().cpu(), "ref_fp32": ref_fp32.detach().cpu(),
                    "pt": pt_grad.detach().cpu()}, f"/tmp/grads_{name}.pt")
    assert not torch.isnan(cute_grad).any(), f"{name} contains NaN"
    atol = 2 * (ref_fp32 + 0.3 - 0.3 - ref_fp32).abs().max().item() + extra
    ref = ref_fp32.to(dtype)
    pt_err = (pt_grad - ref).abs().max().item()
    cute_err = (cute_grad - ref).abs().max().item()
    print(f"  {name}: PT err={pt_err:.2e}, CuTE err={cute_err:.2e}, atol={atol:.2e}")
    assert cute_err <= rtol * pt_err + atol, (
        f"{name} CuTE err {cute_err:.2e} exceeds {rtol}*PT err {pt_err:.2e} + {atol:.2e}"
    )


@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("dim", [64, 128])
@pytest.mark.parametrize("seqlens_q,seqlens_k", BWD_SEQLEN_CONFIGS)
@pytest.mark.parametrize("score_mod_tuple", BWD_TEST_PAIRS)
@pytest.mark.parametrize("use_autograd", [True, False])
def test_varlen_score_mod_backward(
    seqlens_q, seqlens_k, dim, dtype, score_mod_tuple, use_autograd
):
    """Varlen backward with score_mod (no global indices in bwd)."""
    if IS_SM90 and dim == 64:
        pytest.skip("head_dim=64 not supported on SM90 for backward")

    torch.random.manual_seed(42)
    cute_fwd, cute_bwd, eager_factory, aux_type = score_mod_tuple

    num_heads = 4
    batch_size = len(seqlens_q)
    total_q = sum(seqlens_q)
    total_k = sum(seqlens_k)

    q = torch.randn(total_q, num_heads, dim, device="cuda", dtype=dtype)
    k = torch.randn(total_k, num_heads, dim, device="cuda", dtype=dtype)
    v = torch.randn(total_k, num_heads, dim, device="cuda", dtype=dtype)
    cu_seqlens_q = torch.tensor(
        [0] + list(torch.tensor(seqlens_q).cumsum(0).tolist()),
        device="cuda",
        dtype=torch.int32,
    )
    cu_seqlens_k = torch.tensor(
        [0] + list(torch.tensor(seqlens_k).cumsum(0).tolist()),
        device="cuda",
        dtype=torch.int32,
    )

    aux_tensors = None
    if aux_type == "dual_buffer":
        max_seqlen_q = max(seqlens_q)
        head_bias = torch.randn(num_heads, device="cuda", dtype=dtype) * 0.2
        pos_bias = torch.arange(max_seqlen_q, device="cuda", dtype=dtype) * 0.01
        aux_tensors = [head_bias, pos_bias]
        eager_score_mod = eager_factory(head_bias, pos_bias)
    else:
        eager_score_mod = eager_factory

    out_cute, grad_out, dq_cute, dk_cute, dv_cute = run_cute_flash_bwd_varlen(
        q,
        k,
        v,
        cute_fwd,
        cute_bwd,
        cu_seqlens_q,
        cu_seqlens_k,
        aux_tensors=aux_tensors,
        use_autograd=use_autograd,
    )

    _, dq_ref_fp32, dk_ref_fp32, dv_ref_fp32 = run_flex_varlen_ref_bwd(
        q,
        k,
        v,
        cu_seqlens_q,
        cu_seqlens_k,
        eager_score_mod,
        grad_out,
        dtype=torch.float32,
    )
    _, dq_pt, dk_pt, dv_pt = run_flex_varlen_ref_bwd(
        q,
        k,
        v,
        cu_seqlens_q,
        cu_seqlens_k,
        eager_score_mod,
        grad_out,
        dtype=dtype,
    )

    print(f"\nVarlen backward for {cute_fwd.__name__} ({batch_size=} {dtype=} {dim=}):")
    _check_grad("dQ", dq_cute, dq_ref_fp32, dq_pt, dtype)
    _check_grad("dK", dk_cute, dk_ref_fp32, dk_pt, dtype)
    _check_grad("dV", dv_cute, dv_ref_fp32, dv_pt, dtype)


@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("dim", [64, 128])
@pytest.mark.parametrize("seqlens_q,seqlens_k", BWD_SEQLEN_CONFIGS)
@pytest.mark.parametrize("score_mod_tuple", BWD_TEST_PAIRS_WITH_GLOBAL)
@pytest.mark.parametrize("use_autograd", [True, False])
def test_varlen_score_mod_backward_with_global(
    seqlens_q, seqlens_k, dim, dtype, score_mod_tuple, use_autograd
):
    """Varlen backward with a score_mod whose bwd reads aux at global indices.

    Forward: scale[global_kv] * score**2.
    Backward: 2 * scale[global_kv] * score * grad — needs all of grad, score, and
    a global-indexed aux read.
    """
    if IS_SM90 and dim == 64:
        pytest.skip("head_dim=64 not supported on SM90 for backward")

    torch.random.manual_seed(42)
    cute_fwd, cute_bwd, eager_factory, aux_type, requires_global = score_mod_tuple

    num_heads = 4
    total_q = sum(seqlens_q)
    total_k = sum(seqlens_k)

    q = torch.randn(total_q, num_heads, dim, device="cuda", dtype=dtype)
    k = torch.randn(total_k, num_heads, dim, device="cuda", dtype=dtype)
    v = torch.randn(total_k, num_heads, dim, device="cuda", dtype=dtype)
    cu_seqlens_q = torch.tensor(
        [0] + list(torch.tensor(seqlens_q).cumsum(0).tolist()),
        device="cuda",
        dtype=torch.int32,
    )
    cu_seqlens_k = torch.tensor(
        [0] + list(torch.tensor(seqlens_k).cumsum(0).tolist()),
        device="cuda",
        dtype=torch.int32,
    )

    scale_tensor = torch.randn(total_k, device="cuda", dtype=dtype) * 0.1 + 1.0
    aux_tensors = [scale_tensor]
    eager_score_mod = eager_factory(scale_tensor, cu_seqlens_k)

    out_cute, grad_out, dq_cute, dk_cute, dv_cute = run_cute_flash_bwd_varlen(
        q,
        k,
        v,
        cute_fwd,
        cute_bwd,
        cu_seqlens_q,
        cu_seqlens_k,
        aux_tensors=aux_tensors,
        use_autograd=use_autograd,
    )

    _, dq_ref_fp32, dk_ref_fp32, dv_ref_fp32 = run_flex_varlen_ref_bwd(
        q,
        k,
        v,
        cu_seqlens_q,
        cu_seqlens_k,
        eager_score_mod,
        grad_out,
        dtype=torch.float32,
    )
    _, dq_pt, dk_pt, dv_pt = run_flex_varlen_ref_bwd(
        q,
        k,
        v,
        cu_seqlens_q,
        cu_seqlens_k,
        eager_score_mod,
        grad_out,
        dtype=dtype,
    )

    print(f"\nVarlen backward (global) for {cute_fwd.__name__} ({dtype=} {dim=}):")
    _check_grad("dQ", dq_cute, dq_ref_fp32, dq_pt, dtype)
    _check_grad("dK", dk_cute, dk_ref_fp32, dk_pt, dtype)
    _check_grad("dV", dv_cute, dv_ref_fp32, dv_pt, dtype)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
