test_extend_attention.py

# SPDX-License-Identifier: MIT

import torch
import torch.nn.functional as F
import pytest

from aiter.ops.triton.extend_attention import extend_attention_fwd


def extend_attention_fwd_torch_swa(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    o: torch.Tensor,
    k_cache: torch.Tensor,
    v_cache: torch.Tensor,
    qo_indptr: torch.Tensor,
    kv_indptr: torch.Tensor,
    kv_indices: torch.Tensor,
    sliding_window_size: int,
    *,
    k_scale: float = 1.0,
    v_scale: float = 1.0,
    sm_scale: float | None = None,
):
    """Reference for causal + sliding-window extend attention (sglang test style).

    Runs the heavy matmul/softmax on CPU float32 for numerical stability and to avoid
    ROCm aborts on large bf16 einsum after GPU kernels.

    v2 与 Triton 一致：``k_scale`` / ``v_scale`` **只作用在 prefix（cache）键位**；
    extend 段 logits 与 V 不额外乘这两个标量。
    """
    B = qo_indptr.size(0) - 1
    _, H_Q, D = q.shape
    _, H_KV, _ = k.shape

    group_size = H_Q // H_KV
    scale = float(sm_scale) if sm_scale is not None else 1.0 / D**0.5
    out_dev = o.device
    out_dtype = o.dtype

    for i in range(B):
        q_start = int(qo_indptr[i].item())
        q_end = int(qo_indptr[i + 1].item())
        kv_start = int(kv_indptr[i].item())
        kv_end = int(kv_indptr[i + 1].item())

        prefix_indices = kv_indices[kv_start:kv_end]
        k_prefix = k_cache[prefix_indices]
        v_prefix = v_cache[prefix_indices]

        k_extend = k[q_start:q_end]
        v_extend = v[q_start:q_end]
        q_extend = q[q_start:q_end]

        k_full = torch.cat([k_prefix, k_extend], dim=0)
        v_full = torch.cat([v_prefix, v_extend], dim=0)

        if group_size != 1:
            k_full_hq = k_full.repeat_interleave(group_size, dim=1)
            v_full_hq = v_full.repeat_interleave(group_size, dim=1)
        else:
            k_full_hq = k_full
            v_full_hq = v_full

        prefix_len = k_prefix.size(0)
        extend_len = k_extend.size(0)
        total_len = prefix_len + extend_len

        q_e = q_extend.detach().float().cpu()
        k_h = k_full_hq.detach().float().cpu()
        v_h = v_full_hq.detach().float().cpu()

        pos_keys = torch.arange(total_len)
        t = prefix_len + torch.arange(extend_len)
        causal_mask = pos_keys.unsqueeze(0) <= t.unsqueeze(1)

        if sliding_window_size is not None and sliding_window_size > 0:
            start = (t - sliding_window_size).clamp_min(0)
        else:
            start = torch.zeros_like(t)
        window_mask = pos_keys.unsqueeze(0) >= start.unsqueeze(1)

        final_mask = causal_mask & window_mask

        attn_scores = torch.einsum("qhd,khd->qhk", q_e, k_h) * scale
        if k_scale != 1.0:
            attn_scores[:, :, :prefix_len] = attn_scores[:, :, :prefix_len] * k_scale
        attn_scores = attn_scores.masked_fill(~final_mask.unsqueeze(1), float("-inf"))

        attn_weights = F.softmax(attn_scores, dim=-1)
        if v_scale != 1.0:
            v_prefix = v_h[:prefix_len] * v_scale
            v_h_scaled = torch.cat([v_prefix, v_h[prefix_len:]], dim=0)
        else:
            v_h_scaled = v_h
        out_cpu = torch.einsum("qhk,khd->qhd", attn_weights, v_h_scaled)
        o[q_start:q_end] = out_cpu.to(device=out_dev, dtype=out_dtype)


def input_helper(
    B,
    H,
    prefix_length,
    extend_length,
    kv_lora_rank,
    qk_rope_head_dim,
    v_head_dim,
    dtype,
    device,
    attn_impl="normal",
    equal_seqlens=False,
    requires_grad=False,
):
    torch.manual_seed(0)

    if not equal_seqlens:
        max_extend_length = extend_length
        max_prefix_length = prefix_length

        seqlens_extend = torch.randint(
            1, max_extend_length + 1, (B,), dtype=torch.int32
        )
        if prefix_length == 0:
            seqlens_prefix = torch.full((B,), prefix_length, dtype=torch.int32)
        else:
            seqlens_prefix = torch.randint(
                1, max_prefix_length + 1, (B,), dtype=torch.int32
            )

    else:
        seqlens_extend = torch.full((B,), extend_length, dtype=torch.int32)
        seqlens_prefix = torch.full((B,), prefix_length, dtype=torch.int32)

    cu_seqlens_extend = torch.cat(
        [
            torch.tensor([0], dtype=torch.int32),
            seqlens_extend.cumsum(dim=0, dtype=torch.int32),
        ]
    )
    cu_seqlens_prefix = torch.cat(
        [
            torch.tensor([0], dtype=torch.int32),
            seqlens_prefix.cumsum(dim=0, dtype=torch.int32),
        ]
    )

    cu_seqlens_extend = cu_seqlens_extend.to(device="cuda")
    cu_seqlens_prefix = cu_seqlens_prefix.to(device="cuda")

    total_extend = cu_seqlens_extend[-1].item()
    total_prefix = cu_seqlens_prefix[-1].item()

    if attn_impl == "absorb":
        Lq = kv_lora_rank + qk_rope_head_dim
        Lk = kv_lora_rank + qk_rope_head_dim
        Lv = kv_lora_rank
    else:
        Lq = v_head_dim + qk_rope_head_dim
        Lk = v_head_dim + qk_rope_head_dim
        Lv = v_head_dim

    q_extend = torch.randn(
        total_extend, H, Lq, dtype=dtype, device=device
    ).requires_grad_(requires_grad)

    # extend parts
    k_extend = torch.randn(
        total_extend, 1, Lk, dtype=dtype, device=device
    ).requires_grad_(requires_grad)
    v_extend = k_extend[..., :Lv]

    # extend indexing
    qo_indptr = cu_seqlens_extend

    # prefix parts
    k_buffer = torch.randn(
        total_prefix, 1, Lk, dtype=dtype, device=device
    ).requires_grad_(requires_grad)
    v_buffer = k_buffer[..., :Lv]

    if attn_impl != "absorb":
        # simulate v = kv_latent * w_vc which changes the values compared to k
        v_extend = torch.randn_like(v_extend, dtype=v_extend.dtype)
        v_buffer = torch.randn_like(v_buffer, dtype=v_buffer.dtype)

    # prefix indexing
    kv_indptr = cu_seqlens_prefix
    kv_indices = torch.arange(total_prefix, device=device, dtype=torch.int32)

    custom_mask = None
    mask_indptr = None
    max_len_extend = extend_length

    return (
        q_extend,
        k_extend,
        v_extend,
        k_buffer,
        v_buffer,
        kv_indptr,
        kv_indices,
        qo_indptr,
        custom_mask,
        mask_indptr,
        max_len_extend,
    )


@pytest.mark.parametrize(
    "B, H, prefix, extend, kv_lora_rank, qk_rope_head_dim, v_head_dim",
    [
        (2, 4, 0, 512, 32, 16, 32),
        (3, 5, 0, 333, 18, 13, 17),
        (3, 5, 512, 333, 18, 0, 17),
        (3, 5, 110, 333, 18, 0, 19),
        # (8, 16, 0, 1024, 128, 0, 128), # this one passes
        # (8, 16, 0, 16324, 128, 0, 128), # this one fails, numeric precision is likely the issue
        (2, 1, 64, 32, 128, 64, 128),
        (2, 1, 64, 32, 128, 64, 128),
        (4, 16, 64, 96, 128, 64, 128),
        (1, 16, 0, 7, 512, 64, 512),
        (1, 16, 7, 4, 512, 64, 512),
        (1, 16, 32, 4, 512, 64, 512),
        (1, 16, 64, 3, 512, 64, 512),
        (1, 16, 127, 4, 512, 64, 512),
        (1, 16, 255, 15, 512, 64, 512),
        (3, 16, 452, 16, 512, 64, 512),
        (4, 16, 512, 14, 512, 64, 512),
        (4, 16, 1024, 16, 512, 64, 512),
        (4, 16, 2048, 13, 512, 64, 512),
    ],
)
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
@pytest.mark.parametrize("causal", [False, True])
@pytest.mark.parametrize("ref_attn_impl", ["normal", "absorb"])
def test_op_fwd(
    B,
    H,
    prefix,
    extend,
    kv_lora_rank,
    qk_rope_head_dim,
    v_head_dim,
    dtype,
    ref_attn_impl,
    causal,
    sm_scale=1.0,
    logit_cap=0.0,
    device="cuda",
):
    torch.manual_seed(0)
    torch.set_default_device(device)
    torch.set_default_dtype(dtype)

    (
        q_extend,
        k_extend,
        v_extend,
        k_buffer,
        v_buffer,
        kv_indptr,
        kv_indices,
        qo_indptr,
        custom_mask,
        mask_indptr,
        max_len_extend,
    ) = input_helper(
        B,
        H,
        prefix,
        extend,
        kv_lora_rank,
        qk_rope_head_dim,
        v_head_dim,
        dtype,
        device,
        ref_attn_impl,
    )
    tri_out = torch.empty(
        (*q_extend.shape[:-1], v_extend.shape[-1]),
        dtype=q_extend.dtype,
        device=q_extend.device,
    )

    # Reference
    extend_attention_fwd(
        q_extend,
        k_extend,
        v_extend,
        tri_out,
        k_buffer,
        v_buffer,
        qo_indptr,
        kv_indptr,
        kv_indices,
        custom_mask,
        causal,
        mask_indptr,
        max_len_extend,
        sm_scale=sm_scale,
        logit_cap=logit_cap,
    )

    ref_out = torch.empty_like(tri_out, dtype=q_extend.dtype, device=q_extend.device)
    # ref implementation
    for i in range(0, B):
        start_q, start_k = qo_indptr[i], kv_indptr[i]
        end_q, end_k = qo_indptr[i + 1], kv_indptr[i + 1]

        # Get query, prefix key/values, and extend key/values
        q = q_extend[start_q:end_q]  # [seq_len, H, C]
        k_prefix = k_buffer[start_k:end_k]  # [prefix_len, 1, C]
        v_prefix = v_buffer[start_k:end_k]  # [prefix_len, 1, C]
        k_ext = k_extend[start_q:end_q]  # [seq_len, 1, C]
        v_ext = v_extend[start_q:end_q]  # [seq_len, 1, C]

        prefix_len = end_k - start_k
        seq_len = end_q - start_q

        # Calculate attention scores for prefix tokens
        scores_prefix = torch.einsum(
            "qhc,khc->hqk", q.float(), k_prefix.float()
        )  # .float()

        # Calculate attention scores for extend tokens
        scores_extend = torch.einsum(
            "qhc,khc->hqk", q.float(), k_ext.float()
        )  # .float()

        # Apply causal mask only to the extend part if needed
        if causal:
            causal_mask = torch.triu(
                torch.ones(
                    (seq_len, seq_len), dtype=torch.bool, device=scores_extend.device
                ),
                diagonal=1,
            )
            causal_mask = causal_mask.unsqueeze(0).expand(
                scores_extend.shape[0], -1, -1
            )
            scores_extend = scores_extend.masked_fill(causal_mask, float("-inf"))

        # Combine scores and apply softmax
        scores_combined = torch.cat([scores_prefix, scores_extend], dim=-1) * sm_scale
        p_combined = torch.softmax(scores_combined, dim=-1).to(dtype)

        # Split the attention weights back
        p_prefix = p_combined[:, :, :prefix_len]
        p_extend = p_combined[:, :, prefix_len:]

        # Calculate output separately and combine
        out_prefix = torch.einsum(
            "hqk,khd->qhd", p_prefix.to(dtype).float(), v_prefix.float()
        )
        out_extend = torch.einsum(
            "hqk,khd->qhd", p_extend.to(dtype).float(), v_ext.float()
        )

        ref_out[start_q:end_q] = out_prefix.to(dtype) + out_extend.to(dtype)

    torch.testing.assert_close(ref_out, tri_out, rtol=2e-2, atol=2e-2)


def test_extend_attention_v2_identity_scales_match_v1():
    """v2 with fp32 1.0 scales should match v1 (k_scale/v_scale None)."""
    device = "cuda"
    dtype = torch.float16
    torch.manual_seed(0)
    (
        q_extend,
        k_extend,
        v_extend,
        k_buffer,
        v_buffer,
        kv_indptr,
        kv_indices,
        qo_indptr,
        custom_mask,
        mask_indptr,
        max_len_extend,
    ) = input_helper(
        2, 4, 64, 32, 128, 64, 128, dtype, device, "normal"
    )
    out_v1 = torch.empty(
        (*q_extend.shape[:-1], v_extend.shape[-1]),
        dtype=q_extend.dtype,
        device=device,
    )
    out_v2 = torch.empty_like(out_v1)

    extend_attention_fwd(
        q_extend,
        k_extend,
        v_extend,
        out_v1,
        k_buffer,
        v_buffer,
        qo_indptr,
        kv_indptr,
        kv_indices,
        custom_mask,
        True,
        mask_indptr,
        max_len_extend,
        sm_scale=None,
        logit_cap=0.0,
        skip_prefix_custom_mask=True,
        config=None,
    )

    extend_attention_fwd(
        q_extend,
        k_extend,
        v_extend,
        out_v2,
        k_buffer,
        v_buffer,
        qo_indptr,
        kv_indptr,
        kv_indices,
        custom_mask,
        True,
        mask_indptr,
        max_len_extend,
        sm_scale=None,
        logit_cap=0.0,
        skip_prefix_custom_mask=True,
        config=None,
        k_scale=1.0,
        v_scale=1.0,
        sliding_window_size=-1,
        sinks=None,
        window_kv_offsets=None,
        xai_temperature_len=-1,
    )

    torch.testing.assert_close(out_v1, out_v2, rtol=2e-2, atol=2e-2)


def _build_extend_inputs_swa_style(B, N_CTX, H_Q, H_KV, D, device, dtype):
    """Layout aligned with sglang test_triton_attention_kernels sliding-window setup."""
    b_seq_len_prefix = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32, device=device)
    b_seq_len_extend = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32, device=device)
    b_seq_len = b_seq_len_prefix + b_seq_len_extend

    b_start_loc = torch.zeros((B,), dtype=torch.int32, device=device)
    b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0)
    b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device=device)
    b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)

    kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device)
    kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len_prefix[:B], dim=0)
    kv_indices = torch.zeros(
        (b_seq_len_prefix.sum().item(),), dtype=torch.int32, device=device
    )

    for i in range(B):
        kv_indices[kv_indptr[i] : kv_indptr[i + 1]] = torch.arange(
            b_start_loc[i], b_start_loc[i] + b_seq_len_prefix[i], device=device
        )

    total_token_num = torch.sum(b_seq_len).item()
    extend_token_num = torch.sum(b_seq_len_extend).item()
    k_buffer = torch.empty(
        (total_token_num, H_KV, D), dtype=dtype, device=device
    ).normal_(mean=0.1, std=0.2)
    v_buffer = torch.empty(
        (total_token_num, H_KV, D), dtype=dtype, device=device
    ).normal_(mean=0.1, std=0.2)

    k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device)
    v_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device)
    q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device)
    for i in range(B):
        extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i]
        extend_end_in_buffer = b_start_loc[i] + b_seq_len[i]
        extend_start = b_start_loc_extend[i]
        extend_end = b_start_loc_extend[i] + b_seq_len_extend[i]
        k_extend[extend_start:extend_end] = k_buffer[
            extend_start_in_buffer:extend_end_in_buffer
        ]
        v_extend[extend_start:extend_end] = v_buffer[
            extend_start_in_buffer:extend_end_in_buffer
        ]
        q_extend[extend_start:extend_end] = torch.empty(
            (b_seq_len_extend[i], H_Q, D), dtype=dtype, device=device
        ).normal_(mean=0.1, std=0.2)

    b_seq_len_extend = b_seq_len - b_seq_len_prefix
    max_len_extend = torch.max(b_seq_len_extend, 0)[0].item()
    qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device)
    qo_indptr[1 : B + 1] = torch.cumsum(b_seq_len_extend[:B], dim=0)

    return (
        q_extend,
        k_extend,
        v_extend,
        k_buffer,
        v_buffer,
        qo_indptr,
        kv_indptr,
        kv_indices,
        max_len_extend,
    )


@pytest.mark.parametrize("window_size", [-1, 32, 127])
def test_extend_attention_v2_sliding_window(window_size):
    """v2 + sliding_window_size vs torch reference (sglang-style construction)."""
    torch.manual_seed(42)
    device = "cuda"
    dtype = torch.bfloat16
    B, N_CTX, H_Q, H_KV, D = 4, 512, 8, 8, 128

    (
        q_extend,
        k_extend,
        v_extend,
        k_buffer,
        v_buffer,
        qo_indptr,
        kv_indptr,
        kv_indices,
        max_len_extend,
    ) = _build_extend_inputs_swa_style(B, N_CTX, H_Q, H_KV, D, device, dtype)

    extend_token_num = q_extend.shape[0]
    o_triton = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device)
    o_torch = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device)

    extend_attention_fwd(
        q_extend,
        k_extend,
        v_extend,
        o_triton,
        k_buffer,
        v_buffer,
        qo_indptr,
        kv_indptr,
        kv_indices,
        custom_mask=None,
        is_causal=True,
        mask_indptr=None,
        max_len_extend=max_len_extend,
        sm_scale=None,
        logit_cap=0.0,
        skip_prefix_custom_mask=True,
        config=None,
        k_scale=1.2,
        v_scale=1.2,
        sliding_window_size=window_size,
        sinks=None,
        window_kv_offsets=None,
        xai_temperature_len=-1,
    )

    extend_attention_fwd_torch_swa(
        q_extend,
        k_extend,
        v_extend,
        o_torch,
        k_buffer,
        v_buffer,
        qo_indptr,
        kv_indptr,
        kv_indices,
        window_size,
        k_scale=1.2,
        v_scale=1.2,
        sm_scale=None,
    )

    torch.testing.assert_close(o_triton, o_torch, rtol=2e-2, atol=2e-2)


if __name__ == "__main__":
    test_op_fwd(1, 2, 1024, 1024, 256, 0, 256, torch.bfloat16, "normal", False)
    test_op_fwd(3, 5, 110, 333, 18, 0, 17, torch.float32, "normal", True)
    test_op_fwd(4, 16, 1024, 16, 512, 64, 512, torch.float16, "normal", True)