Attention Operators

本页暂无中文版。以下为英文原文。

Every op on this page is used the same way: construct it once, then call it. The constructor takes what the kernel is compiled with; the call takes the tensors. Both are documented under each op — __init__ and forward, where forward is what runs when you call op(...).

Multi-head attention

tileops.ops.attention.mha.MultiHeadAttentionFwdOp

Layout: BSHD.

MHA is the heads_kv == heads specialization of GQA, so route the maintained forward path through the GQA prefill dispatcher.

__init__

__init__(
    batch,
    heads,
    seq_len,
    dim,
    is_causal=True,
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• is_causal (bool, default: True ) –

  Manifest params.is_causal, bool, default True.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

forward

forward(
    q,
    k,
    v,
)

Run MHA forward.

tileops.ops.attention.mha.MultiHeadAttentionBwdOp

Layout: BSHD.

MHA backward is the heads_kv == heads specialization of GQA backward, matching the forward path's dispatch through GQA.

__init__

__init__(
    batch,
    heads,
    seq_len,
    dim,
    is_causal=True,
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• is_causal (bool, default: True ) –

  Manifest params.is_causal, bool, default True.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

forward

forward(
    q,
    k,
    v,
    o,
    do,
    lse,
)

Run the op on the inputs the manifest declares.

Parameters:

• q (Tensor) –

  Input tensor, dtype float16 | bfloat16.

• k (Tensor) –

  Input tensor, dtype same_as(q).

• v (Tensor) –

  Input tensor, dtype same_as(q).

• o (Tensor) –

  Input tensor, dtype same_as(q).

• do (Tensor) –

  Input tensor, dtype same_as(q).

• lse (Tensor) –

  Input tensor, dtype float32.

Returns:

• tuple[Tensor, Tensor, Tensor] –

  dq, dk, dv, as the manifest declares. Shape rules: dq.shape == (B, S, H, D); dk.shape == (B, S, H, D); dv.shape == (B, S, H, D).

tileops.ops.attention.mha.MultiHeadAttentionDecodeWithKVCacheFwdOp

Layout: BSHD

__init__

__init__(
    batch,
    heads,
    seqlen_q,
    seqlen_kv,
    dim,
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

forward

forward(
    q,
    k,
    v,
)

Run the op on the inputs the manifest declares.

Parameters:

• q (Tensor) –

  Input tensor, dtype float16 | bfloat16.

• k (Tensor) –

  Input tensor, dtype same_as(q).

• v (Tensor) –

  Input tensor, dtype same_as(q).

Returns:

• Tensor –

  o, as the manifest declares. Shape rules: o.shape == (B, S_q, H, D).

tileops.ops.attention.mha.MultiHeadAttentionDecodePagedWithKVCacheFwdOp

Paged MHA decode with dynamic KV cache. Layout: Q \([batch \times seqlen\_q \times heads \times dim]\) (BSHD); K, V physical cache [seqlen_kv, heads, dim]; real_seqlen_kv [batch]; block_table [batch, num_pages].

__init__

__init__(
    batch,
    heads,
    seqlen_q,
    seqlen_kv,
    dim,
    page_size,
    is_causal=False,
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• page_size (int) –

  Manifest params.page_size, int.

• is_causal (bool, default: False ) –

  Manifest params.is_causal, bool, default False.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

forward

forward(
    q,
    k,
    v,
    real_seqlen_kv,
    block_table,
)

Run the op on the inputs the manifest declares.

Parameters:

• q (Tensor) –

  Input tensor, dtype float16 | bfloat16.

• k (Tensor) –

  Input tensor, dtype same_as(q).

• v (Tensor) –

  Input tensor, dtype same_as(q).

• real_seqlen_kv (Tensor) –

  Input tensor, dtype int32.

• block_table (Tensor) –

  Input tensor, dtype int32.

Returns:

• Tensor –

  o, as the manifest declares. Shape rules: o.shape == (B, S_q, H, D).

Grouped-query attention

tileops.ops.attention.gqa.GroupedQueryAttentionFwdOp

Compatibility square GQA forward wrapper. Public layout: BSHD.

__init__

__init__(
    batch,
    heads,
    heads_kv,
    seq_len,
    dim,
    is_causal=True,
    sm_scale=None,
    softcap=None,
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• is_causal (bool, default: True ) –

  Manifest params.is_causal, bool, default True.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

forward

forward(
    q,
    k,
    v,
)

Run square GQA forward.

tileops.ops.attention.gqa.GroupedQueryAttentionBwdOp

Layout: BSHD

__init__

__init__(
    batch,
    heads,
    heads_kv,
    seq_len,
    dim,
    is_causal=True,
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• is_causal (bool, default: True ) –

  Manifest params.is_causal, bool, default True.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

forward

forward(
    q,
    k,
    v,
    o,
    do,
    lse,
)

Run the op on the inputs the manifest declares.

Parameters:

• q (Tensor) –

  Input tensor, dtype float16 | bfloat16.

• k (Tensor) –

  Input tensor, dtype same_as(q).

• v (Tensor) –

  Input tensor, dtype same_as(q).

• o (Tensor) –

  Input tensor, dtype same_as(q).

• do (Tensor) –

  Input tensor, dtype same_as(q).

• lse (Tensor) –

  Input tensor, dtype float32.

Returns:

• tuple[Tensor, Tensor, Tensor] –

  dq, dk, dv, as the manifest declares. Shape rules: dq.shape == (B, S, H, D); dk.shape == (B, S, H_kv, D); dv.shape == (B, S, H_kv, D).

tileops.ops.attention.gqa.GroupedQueryAttentionPrefillFwdOp

Canonical packed GQA prefill. Layout: THD.

Dense and square prefill are represented with uniform cu_seqlens. Ragged prefill uses the same fixed public tensor list. Scale tensors are required for manifest stability; non-FP8 kernels ignore them.

__init__

__init__(
    batch,
    heads,
    heads_kv,
    dim,
    max_seqlen_q,
    max_seqlen_kv,
    dtype=torch.float16,
    is_causal=True,
    sm_scale=None,
    softcap=None,
    window_size_left=-1,
    window_size_right=-1,
    backend="auto",
    validate_uniform_cu_seqlens=True,
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• dtype (dtype, default: float16 ) –

  Element type of o. The inputs do not determine it: identical float8_e4m3fn q/k/v admit either a float16 or a bfloat16 output, so the caller chooses here. For float16 / bfloat16 inputs it must equal their element type.

forward

forward(
    q,
    k,
    v,
    cu_seqlens_q,
    cu_seqlens_kv,
    q_scale,
    k_scale,
    v_scale,
)

Run the op on the inputs the manifest declares.

Parameters:

• q (Tensor) –

  Input tensor, dtype float16 | bfloat16 | float8_e4m3fn.

• k (Tensor) –

  Input tensor, dtype same_as(q).

• v (Tensor) –

  Input tensor, dtype same_as(q).

• cu_seqlens_q (Tensor) –

  Input tensor, dtype int32.

• cu_seqlens_kv (Tensor) –

  Input tensor, dtype int32.

• q_scale (Tensor) –

  Input tensor, dtype float32.

• k_scale (Tensor) –

  Input tensor, dtype float32.

• v_scale (Tensor) –

  Input tensor, dtype float32.

Returns:

• Tensor –

  o, as the manifest declares. Shape rules: o.shape == (total_q, H, D).

tileops.ops.attention.gqa.GroupedQueryAttentionPrefillVarlenFwdOp

Packed variable-length GQA prefill. Layout: THD.

cu_seqlens_q and cu_seqlens_kv describe packed per-request ranges. Causal prefill uses bottom-right alignment for each request independently: key position j is visible to query position i iff j <= i + (kv_len - q_len).

__init__

__init__(
    batch,
    heads,
    heads_kv,
    dim,
    max_seqlen_q,
    max_seqlen_kv,
    is_causal=True,
    sm_scale=None,
    softcap=None,
    validate_inputs=False,
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

forward

forward(
    q,
    k,
    v,
    cu_seqlens_q,
    cu_seqlens_kv,
)

Run the op on q, k, v, cu_seqlens_q and cu_seqlens_kv.

tileops.ops.attention.gqa.GroupedQueryAttentionPrefillPagedWithKVCacheFwdOp

Packed GQA prefill with paged KV cache append. Layout: THD.

The current chunk is packed by request. cache_seqlens stores each request's logical KV length before append. block_table maps logical page ids to physical pages in k_pages / v_pages.

__init__

__init__(
    batch,
    heads,
    heads_kv,
    max_pages_per_req,
    page_size,
    dim,
    is_causal=True,
    cache_dtype=None,
    sm_scale=None,
    softcap=None,
    kernel_map=None,
    tune=False,
    fuse_rope=False,
    rope_base=10000.0,
    max_position=None,
    rotary_dim=None,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• max_pages_per_req (int) –

  Manifest params.max_pages_per_req, int.

• page_size (int) –

  Manifest params.page_size, int.

• is_causal (bool, default: True ) –

  Manifest params.is_causal, bool, default True.

• cache_dtype (Optional[dtype], default: None ) –

  Manifest params.cache_dtype, dtype | None, default None.

• sm_scale (Optional[float], default: None ) –

  Manifest params.sm_scale, float | None, default None.

• softcap (Optional[float], default: None ) –

  Manifest params.softcap, float | None, default None.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

• fuse_rope (bool, default: False ) –

  Manifest params.fuse_rope, bool, default False.

• rope_base (float, default: 10000.0 ) –

  Manifest params.rope_base, float, default 10000.0.

• max_position (Optional[int], default: None ) –

  Manifest params.max_position, int | None, default None.

• rotary_dim (Optional[int], default: None ) –

  Manifest params.rotary_dim, int | None, default None.

forward

forward(
    q,
    k_new,
    v_new,
    k_pages,
    v_pages,
    k_scale,
    v_scale,
    cu_seqlens_q,
    cache_seqlens,
    block_table,
    max_seqlen_q,
)

Run the op on the inputs the manifest declares.

Parameters:

• q (Tensor) –

  Input tensor, dtype float16 | bfloat16.

• k_new (Tensor) –

  Input tensor, dtype same_as(q).

• v_new (Tensor) –

  Input tensor, dtype same_as(q).

• k_pages (Tensor) –

  Input tensor, dtype float16 | bfloat16 | float8_e4m3fn.

• v_pages (Tensor) –

  Input tensor, dtype same_as(k_pages).

• k_scale (Tensor) –

  Input tensor, dtype float32.

• v_scale (Tensor) –

  Input tensor, dtype float32.

• cu_seqlens_q (Tensor) –

  Input tensor, dtype int32.

• cache_seqlens (Tensor) –

  Input tensor, dtype int32.

• block_table (Tensor) –

  Input tensor, dtype int32.

Returns:

• Tensor –

  o, as the manifest declares. Shape rules: o.shape == (total_q, H, D).

tileops.ops.attention.gqa.GroupedQueryAttentionDecodeWithKVCacheFwdOp

Layout: BSHD

__init__

__init__(
    batch,
    heads,
    heads_kv,
    seqlen_kv,
    dim,
    sm_scale=None,
    softcap=None,
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

forward

forward(
    q,
    k,
    v,
)

Run the op on the inputs the manifest declares.

Parameters:

• q (Tensor) –

  Input tensor, dtype float16 | bfloat16.

• k (Tensor) –

  Input tensor, dtype same_as(q).

• v (Tensor) –

  Input tensor, dtype same_as(q).

Returns:

• Tensor –

  o, as the manifest declares. Shape rules: o.shape == (B, H, D).

tileops.ops.attention.gqa.GroupedQueryAttentionDecodePagedWithKVCacheFwdOp

Paged GQA decode with dynamic KV cache. Layout: Q \([batch \times heads \times dim]\) (BHD); K, V physical cache [seqlen_kv, heads_kv, dim]; real_seqlen_kv [batch]; block_table [batch, num_pages].

__init__

__init__(
    batch,
    heads,
    heads_kv,
    seqlen_kv,
    dim,
    page_size,
    sm_scale=None,
    softcap=None,
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• page_size (int) –

  Manifest params.page_size, int.

• sm_scale (Optional[float], default: None ) –

  Manifest params.sm_scale, float | None, default None.

• softcap (Optional[float], default: None ) –

  Manifest params.softcap, float | None, default None.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

forward

forward(
    q,
    k,
    v,
    real_seqlen_kv,
    block_table,
)

Run the op on the inputs the manifest declares.

Parameters:

• q (Tensor) –

  Input tensor, dtype float16 | bfloat16.

• k (Tensor) –

  Input tensor, dtype same_as(q).

• v (Tensor) –

  Input tensor, dtype same_as(q).

• real_seqlen_kv (Tensor) –

  Input tensor, dtype int32.

• block_table (Tensor) –

  Input tensor, dtype int32.

Returns:

• Tensor –

  o, as the manifest declares. Shape rules: o.shape == (B, H, D).

tileops.ops.attention.gqa.GroupedQueryAttentionSlidingWindowFwdOp

Fixed-length GQA forward with sliding window attention.

Token at q_pos attends to k_pos when ALL applicable conditions hold

• k_pos <= q_pos (is_causal=True)
• k_pos >= q_pos - window_size_left (window_size_left >= 0)
• k_pos <= q_pos + window_size_right (window_size_right >= 0)

Use window_size_left=-1 / window_size_right=-1 for no restriction.

__init__

__init__(
    batch,
    heads,
    heads_kv,
    seq_len,
    dim,
    is_causal=True,
    window_size_left=-1,
    window_size_right=-1,
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• batch (int) –

  Batch size.

• heads (int) –

  Number of query heads.

• heads_kv (int) –

  Number of KV heads (must divide heads evenly).

• seq_len (int) –

  Sequence length (same for Q, K, V).

• dim (int) –

  Head dimension.

• is_causal (bool, default: True ) –

  Whether to apply causal masking.

• window_size_left (int, default: -1 ) –

  Left window size (-1 = unlimited).

• window_size_right (int, default: -1 ) –

  Right window size (-1 = unlimited).

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for hardware-specific kernel dispatch.

• tune (bool, default: False ) –

  Whether to run autotuning on kernel instantiation.

forward

forward(
    q,
    k,
    v,
)

Run fixed-length GQA sliding window forward.

Parameters:

• q (Tensor) –

  Query tensor, shape \([batch \times seq\_len \times heads \times dim]\).

• k (Tensor) –

  Key tensor, shape \([batch \times seq\_len \times heads\_kv \times dim]\).

• v (Tensor) –

  Value tensor, shape \([batch \times seq\_len \times heads\_kv \times dim]\).

Returns:

• Tensor –

  Output tensor, shape \([batch \times seq\_len \times heads \times dim]\).

tileops.ops.attention.gqa.GroupedQueryAttentionSlidingWindowVarlenFwdOp

Variable-length GQA forward with sliding window attention.

Inputs are packed (no padding); per-sample boundaries are given via cu_seqlens arrays. seqlen_q and seqlen_k may differ per sample:

offset = seqlen_k - seqlen_q (per sample, FA3 bottom-right convention)

A token at local q_pos attends to local k_pos when ALL conditions hold

k_pos <= q_pos + offset (is_causal=True) k_pos >= q_pos + offset - window_size_left (window_size_left >= 0) k_pos <= q_pos + offset + window_size_right (window_size_right >= 0)

__init__

__init__(
    batch,
    heads,
    heads_kv,
    dim,
    is_causal=True,
    window_size_left=-1,
    window_size_right=-1,
    accum_dtype=torch.float32,
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• batch (int) –

  Number of sequences in the batch.

• heads (int) –

  Number of query heads.

• heads_kv (int) –

  Number of KV heads (must divide heads evenly).

• dim (int) –

  Head dimension.

• is_causal (bool, default: True ) –

  Whether to apply causal masking.

• window_size_left (int, default: -1 ) –

  Left window size (-1 = unlimited).

• window_size_right (int, default: -1 ) –

  Right window size (-1 = unlimited).

• accum_dtype (dtype, default: float32 ) –

  Accumulator data type for intermediate computations.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for hardware-specific kernel dispatch.

• tune (bool, default: False ) –

  Whether to run autotuning on kernel instantiation.

forward

forward(
    q,
    k,
    v,
    cu_seqlens_q,
    cu_seqlens_k,
    max_seqlen_q,
)

Run variable-length GQA sliding window forward.

Parameters:

• q (Tensor) –

  Query tensor, shape \([total\_q \times heads \times dim]\).

• k (Tensor) –

  Key tensor, shape \([total\_k \times heads\_kv \times dim]\).

• v (Tensor) –

  Value tensor, shape \([total\_k \times heads\_kv \times dim]\).

• cu_seqlens_q (Tensor) –

  Cumulative Q lengths, shape \([batch+1]\), dtype int32.

• cu_seqlens_k (Tensor) –

  Cumulative K lengths, shape \([batch+1]\), dtype int32.

• max_seqlen_q (int) –

  Maximum Q sequence length across the batch.

Returns:

• Tensor –

  Output tensor, shape \([total\_q \times heads \times dim]\).

Multi-head latent attention

tileops.ops.attention.deepseek_mla.MultiHeadLatentAttentionDecodeWithKVCacheFwdOp

Layout: BSHD

__init__

__init__(
    batch,
    heads,
    heads_kv,
    seqlen_kv,
    dim,
    pe_dim,
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• pe_dim (int) –

  Manifest params.pe_dim, int.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

forward

forward(
    q,
    q_pe,
    k,
    k_pe,
)

Run the op on the inputs the manifest declares.

Parameters:

• q (Tensor) –

  Input tensor, dtype float16 | bfloat16.

• q_pe (Tensor) –

  Input tensor, dtype same_as(q).

• k (Tensor) –

  Input tensor, dtype same_as(q).

• k_pe (Tensor) –

  Input tensor, dtype same_as(q).

Returns:

• Tensor –

  o, as the manifest declares. Shape rules: o.shape == (B, H, D).

Native sparse attention

tileops.ops.attention.deepseek_nsa.NSACmpFwdVarlenOp

__init__

__init__(
    seq_num,
    c_seq_len,
    heads,
    dim_k,
    dim_v,
    chunk_num,
    group,
    scale,
    bc,
    bs,
    accum_dtype,
    tune=False,
    kernel_map=None,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

forward

forward(
    q,
    k_cmp,
    v_cmp,
    offsets,
    chunk_offsets,
    token_indices,
)

Run the op on q, k_cmp, v_cmp, offsets, chunk_offsets and token_indices.

tileops.ops.attention.deepseek_nsa.NSATopkVarlenOp

__init__

__init__(
    seq_num,
    c_seq_len,
    heads,
    dim,
    chunk_num,
    group,
    scale,
    selected_block_num,
    bc,
    bs,
    accum_dtype,
    tune=False,
    kernel_map=None,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

forward

forward(
    q,
    k_cmp,
    lse_in,
    offsets,
    chunk_offsets,
    token_indices,
)

Run the op on q, k_cmp, lse_in, offsets, chunk_offsets and token_indices.

tileops.ops.attention.deepseek_nsa.NSAFwdVarlenOp

__init__

__init__(
    batch,
    heads,
    c_seq_len,
    dim,
    is_causal,
    scale,
    block_size,
    groups,
    selected_blocks,
    accum_dtype,
    tune=False,
    kernel_map=None,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

forward

forward(
    q,
    k,
    v,
    block_indices,
    block_counts,
    offsets,
    token_indices,
)

Run the op on q, k, v, block_indices, block_counts, offsets and token_indices.

DeepSeek sparse attention

tileops.ops.attention.deepseek_dsa.DeepSeekSparseAttentionDecodeWithKVCacheFwdOp

Sparse Attention Decode Operation with Key-Value Cache for DeepSeek.

This operation is part of a sparse attention mechanism, designed for use in decoding with key-value (KV) caching.

The layout of the operation is BSHD.

__init__

__init__(
    batch,
    heads,
    seq_len,
    seq_len_kv,
    dim,
    dim_tail,
    topk,
    stride_kv,
    heads_kv,
    q_start_index_s,
    sm_scale=None,
    is_causal=True,
    kernel_map=None,
    tune=False,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• batch (int) –

  The batch size.

• heads (int) –

  The number of attention heads.

• seq_len (int) –

  The length of the input sequence.

• seq_len_kv (int) –

  The length of the key-value sequence.

• dim (int) –

  The dimension of the attention vectors.

• dim_tail (int) –

  The dimension of the tail portion of the attention vectors.

• topk (int) –

  The number of top elements to consider in sparse attention.

• stride_kv (int) –

  The stride for the key-value sequence.

• heads_kv (int) –

  The number of key-value heads.

• q_start_index_s (int) –

  The start index for queries in the sequence.

• sm_scale (Optional[float], default=None, default: None ) –

  Scaling factor for the softmax function.

• is_causal (bool, default=True, default: True ) –

  Whether the attention is causal (True for causal, False for non-causal).

• kernel_map (Optional[Dict[str, Kernel]], default=None, default: None ) –

      Optional mapping for custom kernels.
  

• tune (bool, default=False, default: False ) –

  Whether to enable kernel tuning.

forward

forward(
    q,
    kv,
    indices,
)

Performs the forward pass of the sparse attention operation.

Parameters:

• q (Tensor) –

  The query tensor with shape (batch, seq_len, heads, dim + dim_tail).

• kv (Tensor) –

  The key-value tensor with shape (batch, seq_len_kv, heads_kv, dim + dim_tail).

• indices (Tensor) –

  Indices tensor for sparse attention.

Returns:

• Tensor –

  torch.Tensor: The result of applying the sparse attention operation on the input tensors.