vllm.models.inkling.amd.ops.rel_attention_decode ¶

def decode_split_count(max_kv_len: int, window_left: int) -> int:
    """Return the number of parallel KV partitions for decode."""
    if window_left >= 0:
        effective_len = max(1, min(max_kv_len, window_left + 1))
        return min(4, max(1, triton.cdiv(effective_len, 128)))
    return min(32, max(1, triton.cdiv(max(1, max_kv_len), 2048)))

@torch.no_grad()
def inkling_rel_attention_split_kv_decode(
    q: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    *,
    block_table: torch.Tensor,
    cache_seqlens: torch.Tensor,
    softmax_scale: float,
    window_left: int,
    rel_extent: int,
    rel_logits: torch.Tensor,
    max_kv_len: int,
    out: torch.Tensor,
) -> torch.Tensor:
    """Run split-KV relative attention for single-token decode."""
    num_kv_heads = key_cache.shape[2]
    gqa_group_size = q.shape[1] // num_kv_heads
    max_kv_splits = decode_split_count(max_kv_len, window_left)
    block_d = triton.next_power_of_2(q.shape[2])
    block_h = min(8, gqa_group_size)

    mid_out = torch.empty(
        q.shape[0],
        q.shape[1],
        max_kv_splits,
        q.shape[2],
        dtype=torch.float32,
        device=q.device,
    )
    mid_lse = torch.empty(
        q.shape[0],
        q.shape[1],
        max_kv_splits,
        dtype=torch.float32,
        device=q.device,
    )
    stage1_grid = (
        q.shape[0],
        triton.cdiv(q.shape[1], block_h),
        max_kv_splits,
    )
    _split_kv_stage1[stage1_grid](
        q,
        rel_logits,
        key_cache,
        value_cache,
        block_table,
        cache_seqlens,
        mid_out,
        mid_lse,
        softmax_scale,
        q.stride(0),
        q.stride(1),
        key_cache.stride(0),
        key_cache.stride(1),
        key_cache.stride(2),
        key_cache.stride(3),
        value_cache.stride(0),
        value_cache.stride(1),
        value_cache.stride(2),
        value_cache.stride(3),
        rel_logits.stride(0),
        rel_logits.stride(1),
        rel_logits.stride(2),
        mid_out.stride(0),
        mid_out.stride(1),
        mid_out.stride(2),
        mid_lse.stride(0),
        mid_lse.stride(1),
        mid_lse.stride(2),
        block_table.stride(0),
        page_size=key_cache.shape[1],
        window_left=window_left,
        gqa_group_size=gqa_group_size,
        num_q_heads=q.shape[1],
        head_dim=q.shape[2],
        rel_extent=rel_extent,
        max_kv_splits=max_kv_splits,
        BLOCK_D=block_d,
        BLOCK_H=block_h,
        BLOCK_N=key_cache.shape[1],
        num_warps=4,
        num_stages=1,
    )
    stage2_grid = (q.shape[0], q.shape[1])
    _split_kv_stage2[stage2_grid](
        mid_out,
        mid_lse,
        out,
        cache_seqlens,
        mid_out.stride(0),
        mid_out.stride(1),
        mid_out.stride(2),
        mid_lse.stride(0),
        mid_lse.stride(1),
        mid_lse.stride(2),
        out.stride(0),
        out.stride(1),
        window_left=window_left,
        head_dim=q.shape[2],
        max_kv_splits=max_kv_splits,
        BLOCK_D=block_d,
        num_warps=4,
        num_stages=2,
    )
    return out

def use_split_kv_decode(
    *,
    max_query_len: int,
    max_kv_len: int,
    page_size: int,
    window_left: int,
) -> bool:
    """Select split-KV only where it outperforms the single-pass kernel."""
    if os.getenv("INKLING_SPLIT_KV", "1") != "1":
        return False
    if max_query_len != 1:
        return False
    if page_size >= 64:
        return True
    if window_left >= 0:
        return page_size >= 32
    return max_kv_len >= 8192