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.attention.MHADecodePagedWithKVCacheFwdOp ¶

Paged Multi-Head Attention (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].

A causal call aligns the queries to the end of each request's cache: query i sees the keys up to position i + real_seqlen_kv - seqlen_q. A query that sees no key outputs zeros.

init ¶

__init__(
    page_size,
    is_causal=False,
    *,
    target=None
)

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

Parameters:

page_size (int) –

Manifest params.page_size, int.
is_causal (bool, default: False ) –

Manifest params.is_causal, bool, default False.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

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, same dtype as q.
v (Tensor) –

Input tensor, same dtype 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.attention.GQABwdOp ¶

Grouped-Query Attention (GQA) backward. Layout: BSHD.

init ¶

__init__(
    is_causal=True,
    *,
    target=None
)

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

Parameters:

is_causal (bool, default: True ) –

Manifest params.is_causal, bool, default True.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

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, same dtype as q.
v (Tensor) –

Input tensor, same dtype as q.
o (Tensor) –

Input tensor, same dtype as q.
do (Tensor) –

Input tensor, same dtype 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.attention.GQADenseFwdOp ¶

Grouped-Query Attention (GQA) over dense \(Q\)/\(K\)/\(V\) tensors.

By default the op computes causal attention,

\[ O = \operatorname{softmax}\!\left( \tfrac{1}{\sqrt D}\,QK^{\mathsf T} + M^{\mathrm{causal}}\right) V , \]

where \(M^{\mathrm{causal}}\) is \(0\) on visible entries and \(-\infty\) elsewhere. Accumulation is FP32 and the output takes the input dtype. FP8 input is the one exception: it has no FP8 output, so out_dtype names a 16-bit one.

Every dimension is inferred from the forward tensors and none is fixed at construction, so one instance serves any shape. forward documents the shapes; these are the names it and the equations below use:

Dimension	Meaning
\(B\)	Batch size
\(S_q\), \(S_{kv}\)	Query and KV sequence length
\(H\), \(H_{kv}\)	Query and KV head count
\(D\)	Head dimension

Dense means every batch entry shares one query length and one KV length, and its keys and values live in one contiguous tensor: neither ragged (varlen) batching nor a paged KV cache.

Query heads are partitioned into \(H_{kv}\) groups of \(g = H / H_{kv}\), and the heads of one group share a KV head: writing \(h\) for a query head and \(r\) for a KV head, head \(h\) attends to \(r(h) = \lfloor h / g \rfloor\). \(H\) must be divisible by \(H_{kv}\).

Each capability below is opt-in, enabled through a constructor parameter, the shape or dtype of the tensors passed to forward, or both:

Capability	Enabled by
Attention without the causal mask	`is_causal=False`
Sliding-window visibility	`window_size_left`, `window_size_right`
A custom score scale	`sm_scale`
A logit softcap	`softcap`
Rectangular attention, \(S_q \ne S_{kv}\)	Nothing to set; read from the tensor shapes
RoPE fused into \(Q\) and \(K\)	`pos_encoding_mode="rope"`, plus `rope_cos` and `rope_sin` on the call
FP8 \(Q\)/\(K\)/\(V\), dequantized per KV head	`float8_e4m3fn` inputs, `q_scale`/`k_scale`/`v_scale` on the call, and `out_dtype` naming the 16-bit output

A call proceeds in the stages below, one per row of the table above. Using none of the optional capabilities leaves stages 2 and 3 doing nothing. Below, \(i\) indexes a query row and \(j\) a KV row; every quantity belongs to one batch entry, which the subscripts leave out.

1. Query positions. Query row \(i\) is bottom-right aligned with the KV sequence, so its position in the KV coordinate system is

\[ p_i = i + S_{kv} - S_q . \]

Causal masking, windowing, and the query-side rotation all use \(p_i\), not \(i\). Causal and fused-RoPE calls therefore require \(S_q \le S_{kv}\).

2. FP8 dequantization. Each per-KV-head scale multiplies its tensor; for 16-bit inputs all three are implicitly one:

\[ \begin{aligned} \hat Q_{i,h} &= Q_{i,h} \cdot \mathrm{qscale}_{r(h)}, \\ \hat K_{j,r} &= K_{j,r} \cdot \mathrm{kscale}_{r}, \\ V'_{j,r} &= V_{j,r} \cdot \mathrm{vscale}_{r}. \end{aligned} \]

3. Fused rotation. Let \(R_x^{(d_r,\,\ell)}\) rotate the first rotary_dim dimensions at sequence position \(x\) using rope_layout \(\ell\); without pos_encoding_mode="rope" it is the identity:

\[ Q'_{i,h} = R_{p_i}^{(d_r,\,\ell)}(\hat Q_{i,h}), \qquad K'_{j,r} = R_j^{(d_r,\,\ell)}(\hat K_{j,r}). \]

4. Scores. With \(\alpha\) = sm_scale, default \(1/\sqrt D\):

\[ Z_{h,i,j} = \alpha\, \langle Q'_{i,h}, K'_{j,r(h)} \rangle . \]

5. Visibility. The causal flag and the window together decide which keys query row \(i\) attends to; every other key contributes nothing to its output. With \(w_L\) = window_size_left and \(w_R\) = window_size_right, each -1 when unlimited, the visible set \(\mathcal V_i\) holds key \(j\) exactly when

\[ (\neg\mathrm{causal} \;\lor\; j \le p_i) \;\land\; (w_L=-1 \;\lor\; j \ge p_i-w_L) \;\land\; (w_R=-1 \;\lor\; j \le p_i+w_R). \]

6. Logits. With \(c\) = softcap, capping and masking give

\[ L_{h,i,j} = \begin{cases} c\tanh(Z_{h,i,j}/c), & j\in\mathcal V_i \text{ and } c>0, \\ Z_{h,i,j}, & j\in\mathcal V_i \text{ and } c=0, \\ -\infty, & j\notin\mathcal V_i. \end{cases} \]

7. Output. Normalizing over the KV axis and reducing the values gives

\[ \begin{aligned} P_{h,i,j} &= \operatorname{softmax}_{j}(L_{h,i,j}), \\ O_{i,h} &= \sum_j P_{h,i,j}\,V'_{j,r(h)}. \end{aligned} \]

init ¶

__init__(
    is_causal=True,
    window_size_left=-1,
    window_size_right=-1,
    sm_scale=None,
    softcap=None,
    pos_encoding_mode="none",
    rotary_dim=None,
    rope_layout="neox",
    out_dtype=None,
    *,
    target=None
)

Configure the op. Tensor shapes and input dtype come from each call.

Parameters:

is_causal (bool, default: True ) –

Apply the causal mask, bottom-right aligned.
window_size_left (int, default: -1 ) –

Keys admitted left of the query row's KV position \(p_i\); -1 means unlimited.
window_size_right (int, default: -1 ) –

Keys admitted right of \(p_i\); -1 means unlimited.
sm_scale (Optional[float], default: None ) –

Score scale \(\alpha\). None resolves to \(1 / \sqrt{D}\) using the current call's head dimension.
softcap (Optional[float], default: None ) –

Positive cap \(c\) replacing each raw score \(z\) with \(c \tanh(z / c)\). None or 0 disables capping.
pos_encoding_mode (str, default: 'none' ) –

"none", or "rope" to fuse the rotary embedding into attention.
rotary_dim (Optional[int], default: None ) –

Rotated width of each head; even, at most \(D\), default the full head dimension. Valid only with pos_encoding_mode="rope".
rope_layout (str, default: 'neox' ) –

"neox" (rotate split halves) or "interleaved" (rotate adjacent pairs).
out_dtype (Optional[dtype], default: None ) –

Manifest params.out_dtype. An FP8 call cannot return FP8 and the two 16-bit types are equally valid, so float16 or bfloat16 must be named here; a 16-bit call has nothing to choose and accepts only None or its own input dtype.
target (Target, default: None ) –

Backend target to serve this op, or None to decide from the input device.

Raises:

ValueError –

softcap is negative.

forward ¶

forward(
    q,
    k,
    v,
    q_scale=None,
    k_scale=None,
    v_scale=None,
    rope_cos=None,
    rope_sin=None,
)

Run dense GQA attention over one set of \(Q\)/\(K\)/\(V\) tensors.

\(Q\), \(K\), and \(V\) are laid out row-major in BSHD axis order, so the head dimension is the contiguous one. Every input must be on the same device as q; a non-contiguous input is copied to a contiguous one before the kernel runs.

Parameters:

q (Tensor) –

Queries, \([B \times S_q \times H \times D]\); float16, bfloat16, or float8_e4m3fn.
k (Tensor) –

Keys, \([B \times S_{kv} \times H_{kv} \times D]\), same dtype as q.
v (Tensor) –

Values, \([B \times S_{kv} \times H_{kv} \times D]\), same dtype as q.
q_scale (Optional[Tensor], default: None ) –

FP8 dequantization scales for q, one per batch and KV head, \([B \times H_{kv}]\), float32. The three scales are required together for FP8 input and invalid otherwise.
k_scale (Optional[Tensor], default: None ) –

Scales for k, one per batch and KV head, \([B \times H_{kv}]\), float32.
v_scale (Optional[Tensor], default: None ) –

Scales for v, one per batch and KV head, \([B \times H_{kv}]\), float32.
rope_cos (Optional[Tensor], default: None ) –

RoPE cosine table indexed by KV position and rotated pair, \([P \times d_r / 2]\) with \(P \ge S_{kv}\) and \(d_r\) = rotary_dim, in the output dtype. The two tables are required together with pos_encoding_mode="rope" and invalid otherwise.
rope_sin (Optional[Tensor], default: None ) –

RoPE sine table, same layout, shape, and dtype as rope_cos.

Returns:

Tensor –

Attention output, \([B \times S_q \times H \times D]\), laid out
Tensor –

like q and contiguous. Its dtype is out_dtype for FP8 input,
Tensor –

and the input dtype otherwise.

Raises:

ValueError –

Shapes, dtypes, devices, or optional-input combinations violate the contract above, or no in-tree kernel serves the call; the message names the limit each kernel refused.

tileops.attention.GQAPagedFwdOp ¶

Grouped-Query Attention (GQA) over a caller-owned paged KV cache.

Packed Q and its cumulative sequence lengths cover both prefill and decode. page_table maps logical pages to physical entries in k_pages and v_pages. This Op reads the cache only: allocation, append, and mutation remain runtime responsibilities. The in-tree kernels serve a call in which Q and the cache share a float16 or bfloat16 dtype, over any positive page size and any mix of per-request query lengths, a request with no query token included. Optional float32 sinks contribute a logit per query head to the softmax denominator only: they are not scaled or softcapped; -inf disables a sink and +inf makes that head's output zero. A query row with no visible key returns zero, including when its sink is disabled. RoPE rotates Q and cached K in NeoX or interleaved layout, including a partial rotary width. The table must cover every request's cache length. FP8 calls still require an external target implementation.

The Hopper RoPE path may allocate up to 512 MiB of temporary rotated Q/K/V storage per call. Under CUDA Graph capture, the graph pool retains that storage. Its allocation is bounded by the page-table and rotary-table extents; preparation writes only live cache pages, including the last page's padding. Calls exceeding this budget use fused rotation without expanded KV storage. Device lengths stay live replay inputs; choosing and preparing this path requires no host readback.

init ¶

__init__(
    is_causal=True,
    window_size_left=-1,
    window_size_right=-1,
    sm_scale=None,
    softcap=None,
    out_dtype=None,
    pos_encoding_mode="none",
    rotary_dim=None,
    rope_layout="neox",
    *,
    target=None
)

Configure paged GQA semantics without owning or mutating the cache.

Parameters:

is_causal (bool, default: True ) –

Apply a bottom-right-aligned causal mask per request.
window_size_left (int, default: -1 ) –

Visible keys to the left; -1 is unlimited.
window_size_right (int, default: -1 ) –

Visible keys to the right; -1 is unlimited.
sm_scale (Optional[float], default: None ) –

Score scale, or None for 1 / sqrt(head_dim).
softcap (Optional[float], default: None ) –

Positive score cap; None or zero disables it.
out_dtype (Optional[dtype], default: None ) –

Output dtype, inferred from the input when omitted.
pos_encoding_mode (str, default: 'none' ) –

"none" or "rope".
rotary_dim (Optional[int], default: None ) –

Even rotated width; None uses the full head dimension.
rope_layout (str, default: 'neox' ) –

"neox" or "interleaved".
target (Target, default: None ) –

Backend target, or None to resolve from the input device.

forward ¶

forward(
    q,
    k_pages,
    v_pages,
    page_table,
    cache_seqlens,
    cu_seqlens_q,
    q_scale=None,
    k_scale=None,
    v_scale=None,
    rope_cos=None,
    rope_sin=None,
    sinks=None,
)

Run read-only paged GQA over packed Q and rank-4 KV pages.

Parameters:

q (Tensor) –

Packed queries [total_q, heads, dim].
k_pages (Tensor) –

Key pages [num_pages, page_size, heads_kv, dim].
v_pages (Tensor) –

Value pages [num_pages, page_size, heads_kv, dim].
page_table (Tensor) –

Physical page of each request's logical page [batch, pages].
cache_seqlens (Tensor) –

Each request's cached length, its query tokens included [batch].
cu_seqlens_q (Tensor) –

Request boundaries in the packed queries [batch + 1].
q_scale (Optional[Tensor], default: None ) –

Dequantization scale of an FP8 query [batch, heads_kv].
k_scale (Optional[Tensor], default: None ) –

Dequantization scale of an FP8 key cache.
v_scale (Optional[Tensor], default: None ) –

Dequantization scale of an FP8 value cache.
rope_cos (Optional[Tensor], default: None ) –

Rotary cosine table when pos_encoding_mode='rope'.
rope_sin (Optional[Tensor], default: None ) –

Rotary sine table when pos_encoding_mode='rope'.
sinks (Optional[Tensor], default: None ) –

Optional float32 logits [heads], added to the softmax denominator only.

Returns:

Tensor –

The attention output [total_q, heads, dim].

tileops.attention.GQAPrefillPagedWithKVCacheFwdOp ¶

Packed Grouped-Query Attention (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.

The in-tree kernels refuse a page_size that is not a power of two and fused RoPE over an FP8 cache.

By default the op does not check tensor contents. An FP8 cache scale that is not finite and positive can corrupt the output, and with fused RoPE a request whose cached plus new tokens exceed max_position reads the RoPE table out of bounds. The caller guarantees both, or passes validate_inputs=True, which checks them on every call at the cost of device synchronizations and cannot run inside CUDA Graph capture.

init ¶

__init__(
    page_size,
    max_seqlen_q,
    is_causal=True,
    cache_dtype=None,
    sm_scale=None,
    softcap=None,
    fuse_rope=False,
    rope_base=10000.0,
    max_position=None,
    rotary_dim=None,
    *,
    validate_inputs=False,
    target=None
)

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

Parameters:

page_size (int) –

Manifest params.page_size, int.
max_seqlen_q (int) –

Manifest params.max_seqlen_q, the launch bound the kernel is built for; callers must keep every request within it.
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, which resolves to 1 / sqrt(D) from each call's head dimension.
softcap (Optional[float], default: None ) –

Manifest params.softcap, float | None, default None.
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, which rotates each call's full head dimension.
validate_inputs (bool, default: False ) –

Check scale values and RoPE positions on the CPU. Synchronizes the device; enable only outside CUDA Graph capture.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

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

Attend the packed chunk to each request's cache, appending its keys and values.

Parameters:

q (Tensor) –

Queries of the new tokens, packed by request [total_q, heads, dim].
k_new (Tensor) –

Keys of the new tokens [total_q, heads_kv, dim].
v_new (Tensor) –

Values of the new tokens [total_q, heads_kv, dim].
k_pages (Tensor) –

Key page pool [physical_tokens, heads_kv, dim], written in place.
v_pages (Tensor) –

Value page pool [physical_tokens, heads_kv, dim], written in place.
k_scale (Tensor) –

Key dequantization scale of an FP8 pool [1].
v_scale (Tensor) –

Value dequantization scale of an FP8 pool [1].
cu_seqlens_q (Tensor) –

Request boundaries in the packed chunk [batch + 1].
cache_seqlens (Tensor) –

Each request's cache length before the append [batch].
block_table (Tensor) –

Physical page of each request's logical page [batch, pages].

Returns:

Tensor –

The attention output [total_q, heads, dim].

Raises:

ValueError –

An FP8 pool's scales are not finite and positive, a fused-RoPE call reaches past max_position, or no in-tree kernel serves the call.

tileops.attention.GQAVarlenFwdOp ¶

Grouped-Query Attention (GQA) over packed THD tensors.

cu_seqlens_q and cu_seqlens_kv delimit each request. The interface covers both prefill and decode; tensor geometry and sequence metadata come from each call, while mask, score, out_dtype, and RoPE semantics are fixed at construction. The BUILTIN path implements 16-bit regular and sliding-window attention, with or without fused RoPE. Optional float32 sinks contribute a logit per query head to the softmax denominator only: they are not scaled or softcapped; -inf disables a sink and +inf makes that head's output zero. A query row with no visible key returns zero, including when its sink is disabled.

float8_e4m3fn Q/K/V are dequantized by one q_scale, k_scale and v_scale per request and KV head, and name a 16-bit out_dtype; the FP8 result carries the error of a single e4m3 rounding of the softmax weights, so it agrees with the 16-bit path to about 2% relative.

By default the op does not check the contents of cu_seqlens_q and cu_seqlens_kv. The kernels read and write the packed tensors at the positions the offsets name. Malformed offsets can cause out-of-bounds accesses or incorrect results. The caller guarantees well-formed offsets, or passes validate_inputs=True, which checks them on every call at the cost of a device synchronization and cannot run inside CUDA Graph capture.

init ¶

__init__(
    is_causal=True,
    window_size_left=-1,
    window_size_right=-1,
    sm_scale=None,
    softcap=None,
    out_dtype=None,
    pos_encoding_mode="none",
    rotary_dim=None,
    rope_layout="neox",
    validate_inputs=False,
    *,
    target=None
)

Configure packed variable-length GQA semantics.

Parameters:

is_causal (bool, default: True ) –

Apply a bottom-right-aligned causal mask per request.
window_size_left (int, default: -1 ) –

Visible keys to the left; -1 is unlimited.
window_size_right (int, default: -1 ) –

Visible keys to the right; -1 is unlimited.
sm_scale (Optional[float], default: None ) –

Score scale, or None for 1 / sqrt(head_dim).
softcap (Optional[float], default: None ) –

Positive score cap; None or zero disables it.
out_dtype (Optional[dtype], default: None ) –

Output dtype, inferred from the input when omitted.
pos_encoding_mode (str, default: 'none' ) –

"none" or "rope".
rotary_dim (Optional[int], default: None ) –

Even rotated width; None uses the full head dimension.
rope_layout (str, default: 'neox' ) –

"neox" or "interleaved".
validate_inputs (bool, default: False ) –

Check cumulative offsets against packed tensors on the CPU.
target (Target, default: None ) –

Backend target, or None to resolve from the input device.

forward ¶

forward(
    q,
    k,
    v,
    cu_seqlens_q,
    cu_seqlens_kv,
    q_scale=None,
    k_scale=None,
    v_scale=None,
    rope_cos=None,
    rope_sin=None,
    sinks=None,
)

Run packed Varlen GQA; Q/K/V use [total_tokens, heads, dim].

Multi-Head Latent Attention¶

tileops.attention.MLADecodeWithKVCacheFwdOp ¶

Multi-Head Latent Attention (MLA) decode against a per-request KV cache. Layout: BSHD.

init ¶

__init__(
    *,
    target=None
)

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

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

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, same dtype as q.
k (Tensor) –

Input tensor, same dtype as q.
k_pe (Tensor) –

Input tensor, same dtype as q.

Returns:

Tensor –

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

Native Sparse Attention¶

tileops.attention.NSACompressedVarlenFwdOp ¶

Native Sparse Attention (NSA) compression forward over a ragged batch.

Attends each token to the compressed chunk summaries of its own request and returns both the output and the log-sum-exp NSATopKVarlenFwdOp scores against.

Sequence layout is packed: offsets marks the request boundaries, so the batch size and the chunk count come from the call rather than from construction.

init ¶

__init__(
    scale,
    bs,
    *,
    target=None
)

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

Parameters:

scale (float) –

Softmax scale applied to the QK product.
bs (int) –

Compression block size.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

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

Attend each token to its request's compressed chunks.

Parameters:

q (Tensor) –

Queries, packed over the batch [c_seq_len, heads, dim_k].
k_cmp (Tensor) –

Compressed keys [chunk_num, head_kv, dim_k].
v_cmp (Tensor) –

Compressed values [chunk_num, head_kv, dim_v].
offsets (Tensor) –

Request boundaries into the packed sequence [seq_num + 1].
chunk_offsets (Tensor) –

Per-request chunk boundaries [seq_num + 1].
token_indices (Tensor) –

Request id and in-request position per token [c_seq_len, 2].

Returns:

Tuple[Tensor, Tensor] –

Tuple of (o, lse).

tileops.attention.NSATopKVarlenFwdOp ¶

Native Sparse Attention (NSA) block selection over a ragged batch.

Scores each compressed chunk against the query and returns, per token and per KV head, the selected_block_num block ids the sparse forward will attend to. Follows FLA: the first, previous and current blocks score one per query head; other blocks use normalized attention summed across the GQA group. Rank raw scores descending, with unspecified tie selection/order and trailing -1 padding.

Sequence layout is packed: q holds every request's tokens back to back and offsets marks the boundaries, so the batch size and the chunk count come from the call rather than from construction.

init ¶

__init__(
    scale,
    selected_block_num,
    bs,
    *,
    target=None
)

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

Parameters:

scale (float) –

Multiplies Q in its input dtype before the QK product, as in FLA.
selected_block_num (int) –

Blocks to keep per token and KV head.
bs (int) –

Compression block size.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

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

Select the blocks each token attends to.

Parameters:

q (Tensor) –

Queries, packed over the batch [c_seq_len, heads, dim].
k_cmp (Tensor) –

Compressed keys [chunk_num, head_kv, dim].
offsets (Tensor) –

Request boundaries into the packed sequence [seq_num + 1].
chunk_offsets (Tensor) –

Per-request chunk boundaries [seq_num + 1].
token_indices (Tensor) –

Request id and in-request position per token [c_seq_len, 2].

Returns:

Tensor –

Selected block ids [c_seq_len, head_kv, selected_block_num].

tileops.attention.NSAVarlenFwdOp ¶

Native Sparse Attention (NSA) sparse forward over a ragged batch.

Attends each token to the blocks NSATopKVarlenFwdOp selected for it. Sequence layout is packed: offsets marks the request boundaries, so the batch size and the block count come from the call rather than from construction.

A causal token sees the keys of its selected blocks up to its own position, a non-causal one up to the end of its sequence; a block starting past that bound contributes nothing, and a token no block gives a key outputs zeros.

init ¶

__init__(
    is_causal,
    scale,
    block_size,
    *,
    target=None
)

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

Parameters:

is_causal (bool) –

Whether a token sees keys only up to its own position.
scale (float) –

Softmax scale applied to the QK product.
block_size (int) –

Tokens per selected block.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

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

Attend each token to its selected blocks.

Parameters:

q (Tensor) –

Queries, packed over the batch [c_seq_len, heads, dim].
k (Tensor) –

Keys [c_seq_len, head_kv, dim].
v (Tensor) –

Values [c_seq_len, head_kv, dim].
block_indices (Tensor) –

Selected block ids [c_seq_len, head_kv, selected_blocks].
block_counts (Tensor) –

Valid block count per token and KV head [c_seq_len, head_kv].
offsets (Tensor) –

Request boundaries into the packed sequence [batch + 1].
token_indices (Tensor) –

Request id and in-request position per token [c_seq_len, 2].

Returns:

Tensor –

Attention output [c_seq_len, heads, dim].

DeepSeek Sparse Attention¶

tileops.attention.DSADecodeWithKVCacheFwdOp ¶

DeepSeek Sparse Attention (DSA) decode.

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.

The in-tree kernels serve causal calls with a power-of-two value dimension and a zero or power-of-two tail dimension, subject to shared-memory limits.

init ¶

__init__(
    dim_tail,
    stride_kv,
    q_start_index_s,
    sm_scale=None,
    is_causal=True,
    *,
    target=None
)

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

Parameters:

dim_tail (int) –

The dimension of the tail portion of the attention vectors.
stride_kv (int) –

The stride for the key-value sequence.
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).
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

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.

Attention indexing¶

tileops.attention.FP8LightningIndexerFwdOp ¶

Lightning indexer logits over FP8 index keys.

For query s and key t of group g, the logit sums weights[s, h] times relu(q[s, h] . k[t, g]) over the heads of group g, for keys inside the query's window [cu_seqlen_ks[s], cu_seqlen_ke[s]). A bf16 call is quantized to FP8 by the op; an FP8 call passes the per-key scales in index_k_scale.

init ¶

__init__(
    clean_logits=True,
    *,
    target=None
)

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

Parameters:

clean_logits (bool, default: True ) –

Manifest params.clean_logits, bool, default True.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    index_q,
    index_k,
    weights,
    cu_seqlen_ks,
    cu_seqlen_ke,
    index_k_scale=None,
)

Run the op on the inputs the manifest declares.

Parameters:

index_q (Tensor) –

Input tensor, dtype bfloat16 | float8_e4m3fn.
index_k (Tensor) –

Input tensor, dtype bfloat16 | float8_e4m3fn.
weights (Tensor) –

Input tensor, dtype float32.
cu_seqlen_ks (Tensor) –

Input tensor, dtype int32.
cu_seqlen_ke (Tensor) –

Input tensor, dtype int32.
index_k_scale (Optional[Tensor], default: None ) –

Input tensor, dtype float32. Optional.

Returns:

Tensor –

logits, as the manifest declares.

Paged KV Cache¶

tileops.attention.PagedKVCacheWriteFwdOp ¶

Scatter new K/V rows into a page pool, the write half of a paged attention cache.

The slots come from the caller as one flattened row index per token, -1 skipping the token, so the pages a call writes need not be consecutive, need not be new, and need not belong to one request: a prefix-caching server reuses a page it already holds and a chunked prefill writes the pages it was just given. Distinct non-negative slots are the caller's obligation, and a call writes nothing outside the slots it names. With an FP8 cache the rows are divided by their scale and cast to E4M3FN. NaN is preserved; finite overflow and infinities saturate to the largest finite magnitude with the same sign.

init ¶

__init__(
    *,
    target=None
)

Build the op; the cache geometry and dtypes arrive with each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    k,
    v,
    k_pages,
    v_pages,
    slot_mapping,
    k_scale=None,
    v_scale=None,
)

Write each token's key and value row into the page-pool row its slot names.

Parameters:

k (Tensor) –

New keys, (tokens, heads_kv, dim) in float16 or bfloat16.
v (Tensor) –

New values, the same shape and dtype.
k_pages (Tensor) –

The key page pool, (num_pages, page_size, heads_kv, dim), written in place. Its dtype is the cache dtype: float16, bfloat16, or float8_e4m3fn.
v_pages (Tensor) –

The value page pool, the same shape and dtype, written in place.
slot_mapping (Tensor) –

int64 (tokens,); entry n is page * page_size + offset, the flattened row token n writes, or -1 to skip it.
k_scale (Optional[Tensor], default: None ) –

float32 (1,) divisor of the stored keys, required exactly when the cache is float8_e4m3fn.
v_scale (Optional[Tensor], default: None ) –

The same divisor for the stored values.

Returns:

None –

None; both page pools are mutated and no other tensor is touched.

Attention Operators¶

Multi-Head Attention¶

tileops.attention.MHADecodePagedWithKVCacheFwdOp ¶

__init__ ¶

forward ¶

Grouped-Query Attention¶

tileops.attention.GQABwdOp ¶

__init__ ¶

forward ¶

tileops.attention.GQADenseFwdOp ¶

__init__ ¶

forward ¶

tileops.attention.GQAPagedFwdOp ¶

__init__ ¶

forward ¶

tileops.attention.GQAPrefillPagedWithKVCacheFwdOp ¶

__init__ ¶

forward ¶

tileops.attention.GQAVarlenFwdOp ¶

__init__ ¶

forward ¶

Multi-Head Latent Attention¶

tileops.attention.MLADecodeWithKVCacheFwdOp ¶

__init__ ¶

forward ¶

Native Sparse Attention¶

tileops.attention.NSACompressedVarlenFwdOp ¶

__init__ ¶

forward ¶

tileops.attention.NSATopKVarlenFwdOp ¶

__init__ ¶

forward ¶

tileops.attention.NSAVarlenFwdOp ¶

__init__ ¶

forward ¶

DeepSeek Sparse Attention¶

tileops.attention.DSADecodeWithKVCacheFwdOp ¶

__init__ ¶

forward ¶

Attention indexing¶

tileops.attention.FP8LightningIndexerFwdOp ¶

__init__ ¶

forward ¶

Paged KV Cache¶

tileops.attention.PagedKVCacheWriteFwdOp ¶

__init__ ¶

forward ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶