> ## Documentation Index
> Fetch the complete documentation index at: https://bench.flashinfer.ai/docs/llms.txt
> Use this file to discover all available pages before exploring further.

# MLA Paged

Multi-head Latent Attention (MLA) with paged memory layout. MLA is an advanced attention mechanism that decomposes the key-value representation into separate compressed key-value (CKV) and key positional encoding (KPE) components to reduce memory usage while maintaining model performance. The paged layout enables efficient memory management for variable-length sequences.

Variants:

* prefill
* decode

## prefill

Axes (8 dimensions):

* `total_q`, `num_pages`, `len_indptr`, `num_kv_indices`: variable
* `num_qo_heads`, `head_dim_ckv`, `head_dim_kpe`, `page_size`: constant

Inputs (7 tensors + 1 scalar):

* `q_nope`: query tensor without positional encoding \[total\_q, num\_qo\_heads, head\_dim\_ckv]
* `q_pe`: query positional encoding component \[total\_q, num\_qo\_heads, head\_dim\_kpe]
* `ckv_cache`: compressed key-value cache \[num\_pages, page\_size, head\_dim\_ckv]
* `kpe_cache`: key positional encoding cache \[num\_pages, page\_size, head\_dim\_kpe]
* `qo_indptr`, `kv_indptr`, `kv_indices`: paging indices
* `sm_scale`: softmax scale (scalar)

Outputs (2 tensors):

* `output`: attention output \[total\_q, num\_qo\_heads, head\_dim\_ckv]
* `lse`: log-sum-exp values \[total\_q, num\_qo\_heads]

Constraints:

* `total_q == qo_indptr[-1]`
* `num_kv_indices = kv_indptr[-1]`

## decode

Axes (8 dimensions):

* `batch_size`, `num_pages`, `len_indptr`, `num_kv_indices`: variable
* `num_qo_heads`, `head_dim_ckv`, `head_dim_kpe`, `page_size`: constant

Inputs (6 tensors + 1 scalar):

* `q_nope`: query tensor without positional encoding \[batch\_size, num\_qo\_heads, head\_dim\_ckv]
* `q_pe`: query positional encoding \[batch\_size, num\_qo\_heads, head\_dim\_kpe]
* `ckv_cache`: compressed key-value cache \[num\_pages, page\_size, head\_dim\_ckv]
* `kpe_cache`: key positional encoding cache \[num\_pages, page\_size, head\_dim\_kpe]
* `kv_indptr`, `kv_indices`: paging indices
* `sm_scale`: softmax scale (scalar)

Outputs (2 tensors):

* `output`: attention output \[batch\_size, num\_qo\_heads, head\_dim\_ckv]
* `lse`: log-sum-exp values \[batch\_size, num\_qo\_heads]

Constraints:

* `len_indptr = num_pages + 1`
* `num_kv_indices = kv_indptr[-1]`