### Install FlashMLA

Source: https://github.com/deepseek-ai/flashmla/blob/main/README.md

Follow these steps to clone the FlashMLA repository, initialize submodules, and install the library. Ensure you have the necessary dependencies like PyTorch installed.

```bash
git clone https://github.com/deepseek-ai/FlashMLA.git flash-mla
cd flash-mla
git submodule update --init --recursive
pip install -v .
```

--------------------------------

### Seesaw Scheduling: Step 1 - Compute p0

Source: https://github.com/deepseek-ai/flashmla/blob/main/docs/20250422-new-kernel-deep-dive.md

Calculates the initial product vector p0 using query (q) and key (K0) matrices, scaled by qk_scale. This is the first step in the seesaw scheduling algorithm.

```math
`\vec p_0 = \vec q K_0^\intercal / qk\_scale`
```

--------------------------------

### Benchmark MLA Prefill (Sparse)

Source: https://github.com/deepseek-ai/flashmla/blob/main/README.md

Run this command to benchmark the sparse MLA prefill kernel. This test is crucial for evaluating performance on different GPU architectures and CUDA versions.

```bash
python tests/test_flash_mla_sparse_prefill.py
```

--------------------------------

### Benchmark MHA Prefill (Dense)

Source: https://github.com/deepseek-ai/flashmla/blob/main/README.md

Execute this command to benchmark the dense Multi-Head Attention (MHA) prefill kernel. Performance metrics are reported by NVIDIA for specific GPU architectures.

```bash
python tests/test_fmha_sm100.py
```

--------------------------------

### Seesaw Scheduling: Step 2 - Compute p1

Source: https://github.com/deepseek-ai/flashmla/blob/main/docs/20250422-new-kernel-deep-dive.md

Calculates the second initial product vector p1 using query (q) and key (K1) matrices, scaled by qk_scale. This step is performed concurrently with step 1 by a different warpgroup.

```math
`\vec p_1 = \vec q K_1^\intercal / qk\_scale`
```

--------------------------------

### Seesaw Scheduling: Step 8 - Update oR

Source: https://github.com/deepseek-ai/flashmla/blob/main/docs/20250422-new-kernel-deep-dive.md

Updates the right output matrix (oR) by scaling the previous oR with the combined scaling factors (scale0 * scale1) and adding the product of the softmax-normalized p1 and the V1R matrix. This accumulates results into the right output.

```math
`\vec o_R \gets \vec o_R \cdot (scale_0 \cdot scale_1) + \vec p_1 V_{1R}`
```

--------------------------------

### Benchmark MLA Decoding (Sparse & Dense)

Source: https://github.com/deepseek-ai/flashmla/blob/main/README.md

Run these commands to benchmark the performance of both sparse and dense MLA decoding kernels. Ensure your environment meets the specified requirements for accurate results.

```bash
python tests/test_flash_mla_dense_decoding.py
```

```bash
python tests/test_flash_mla_sparse_decoding.py
```

--------------------------------

### Seesaw Scheduling: Step 5 - Update oL

Source: https://github.com/deepseek-ai/flashmla/blob/main/docs/20250422-new-kernel-deep-dive.md

Updates the left output matrix (oL) by scaling the previous oL with scale0 and adding the product of the softmax-normalized p0 and the V0L matrix. This accumulates results into the left output.

```math
`\vec o_L \gets \vec o_L \cdot scale_0 + \vec p_0 V_{0L}`
```

--------------------------------

### Seesaw Scheduling: Step 4 - Softmax on p0

Source: https://github.com/deepseek-ai/flashmla/blob/main/docs/20250422-new-kernel-deep-dive.md

Applies the softmax function to the p0 vector using the calculated maximum (m_new0) to normalize the values. This prepares p0 for accumulation into the output matrix.

```math
`\vec p_0 \gets \exp(\vec p_0 - m\_new_0)`
```

--------------------------------

### Seesaw Scheduling: Step 10 - Update oR with scaled p0

Source: https://github.com/deepseek-ai/flashmla/blob/main/docs/20250422-new-kernel-deep-dive.md

Further updates the right output matrix (oR) by adding the product of the scaled p0 vector and the V0R matrix. This completes the accumulation for the right output.

```math
`\vec o_R \gets \vec o_R + \vec p_0 V_{0R}`
```

--------------------------------

### FP8 KV Cache Quantization and Dequantization

Source: https://context7.com/deepseek-ai/flashmla/llms.txt

Utilities for quantizing and dequantizing KV cache to FP8 format, essential for sparse attention kernels. Supports different layout configurations like V32_FP8Sparse.

```python
import torch
from tests.quant import (
    FP8KVCacheLayout,
    quantize_k_cache,
    dequantize_k_cache,
    abs_indices2indices_in_kvcache
)

# KV cache layout configurations
# V32_FP8Sparse: d=576, d_nope=512, d_rope=64, tile_size=128, num_tiles=4
# MODEL1_FP8Sparse: d=512, d_nope=448, d_rope=64, tile_size=64, num_tiles=7
layout = FP8KVCacheLayout.V32_FP8Sparse

# Original BF16 KV cache
num_blocks = 1024
block_size = 64
head_dim = 576
k_cache_bf16 = torch.randn(num_blocks, block_size, 1, head_dim,
                           dtype=torch.bfloat16, device='cuda')

# Quantize to FP8 format
# Output format for V32_FP8Sparse (656 bytes per token):
# - Bytes 0-511: quantized NoPE (512 x float8_e4m3)
# - Bytes 512-527: scale factors (4 x float32)
# - Bytes 528-655: RoPE part (64 x bfloat16, unquantized)
k_cache_fp8 = quantize_k_cache(k_cache_bf16, layout)

# Dequantize back to BF16 for verification
k_cache_recovered = dequantize_k_cache(k_cache_fp8, layout)
```

--------------------------------

### flash_attn_varlen_qkvpacked_func

Source: https://context7.com/deepseek-ai/flashmla/llms.txt

Optimized attention with QKV packed into a single tensor. Useful when Q, K, V have the same sequence length and are stored contiguously.

```APIDOC
## flash_attn_varlen_qkvpacked_func

### Description
Optimized attention with QKV packed into a single tensor. Useful when Q, K, V have the same sequence length and are stored contiguously.

### Method
N/A (Python function)

### Endpoint
N/A (Python function)

### Parameters
#### Path Parameters
N/A

#### Query Parameters
N/A

#### Request Body
- **qkv** (torch.Tensor) - Packed QKV tensor: (total_len, num_heads, head_dim_qk * 2 + head_dim_v)
- **cu_seqlens** (torch.Tensor) - Cumulative sequence lengths.
- **max_seqlen** (int) - Maximum sequence length.
- **head_dim_qk** (int) - Dimension of the query and key heads.
- **dropout_p** (float) - Dropout probability.
- **softmax_scale** (float) - Scaling factor for softmax.
- **causal** (bool) - Whether to apply causal attention.
- **is_varlen** (bool) - Whether the input is variable length.

### Request Example
```python
import torch
from flash_mla import flash_attn_varlen_qkvpacked_func

# Configuration
batch_size = 2
seqlens = torch.tensor([4096, 8192], dtype=torch.int32)
total_len = seqlens.sum().item()
cu_seqlens = torch.cumsum(
    torch.nn.functional.pad(seqlens, (1, 0)), 0, dtype=torch.int32
).cuda()

num_heads = 128
head_dim_qk = 128
head_dim_v = 128
max_seqlen = seqlens.max().item()

# Packed QKV tensor: (total_len, num_heads, head_dim_qk + head_dim_qk + head_dim_v)
qkv = torch.randn(total_len, num_heads, head_dim_qk * 2 + head_dim_v,
                  dtype=torch.bfloat16, device='cuda')

# Run attention with packed input
out, lse = flash_attn_varlen_qkvpacked_func(
    qkv=qkv,
    cu_seqlens=cu_seqlens,
    max_seqlen=max_seqlen,
    head_dim_qk=head_dim_qk,
    dropout_p=0.0,
    softmax_scale=head_dim_qk ** (-0.5),
    causal=True,
    is_varlen=True
)
```

### Response
#### Success Response (200)
- **out** (torch.Tensor) - Output tensor.
- **lse** (torch.Tensor) - Log-sum-exp values.

#### Response Example
N/A (Python function return values)
```

--------------------------------

### Seesaw Scheduling: Step 3 - Compute max and scale for p0

Source: https://github.com/deepseek-ai/flashmla/blob/main/docs/20250422-new-kernel-deep-dive.md

Determines the maximum value in p0 (mp0), updates the running maximum (m_new0), and calculates the scaling factor (scale0) based on the previous maximum (m). This is crucial for the softmax operation.

```math
mp0 = max(\vec p_0), `m_new_0 = max(m, mp_0)`, and $`scale_0 = \exp(m\_new_0 - m)`$. Update $`m \gets m\_new_0`$.
```

--------------------------------

### Initialize FlashMLA Scheduler Metadata

Source: https://context7.com/deepseek-ai/flashmla/llms.txt

Call get_mla_metadata once before the decoding loop to initialize scheduler metadata. This metadata can be reused across layers and decoding steps if tensor shapes and cache sequence lengths remain consistent.

```python
import torch
from flash_mla import get_mla_metadata, flash_mla_with_kvcache

# Initialize scheduler metadata (call once before decoding loop)
tile_scheduler_metadata, num_splits = get_mla_metadata()

# The metadata can be reused across multiple layers and decoding steps
# as long as tensor shapes and cache_seqlens remain consistent
for layer_idx in range(num_layers):
    out, lse = flash_mla_with_kvcache(
        q=query_tensor,
        k_cache=kv_cache,
        block_table=block_table,
        cache_seqlens=cache_seqlens,
        head_dim_v=512,
        tile_scheduler_metadata=tile_scheduler_metadata,
        num_splits=num_splits,
        causal=False
    )
```

--------------------------------

### Seesaw Scheduling: Step 6 - Compute max and scale for p1

Source: https://github.com/deepseek-ai/flashmla/blob/main/docs/20250422-new-kernel-deep-dive.md

Determines the maximum value in p1 (mp1), updates the running maximum (m_new1), and calculates the scaling factor (scale1) based on the current maximum (m). This is similar to step 3 but for the p1 vector.

```math
mp1 = max(\vec p_1), $`m\_new_1 = max(m, mp_1)`$, and $`scale_1 = \exp(m\_new_1 - m)`$. Update $`m \gets m\_new_1`$.
```

--------------------------------

### FP8 KV Cache Quantization Utilities

Source: https://context7.com/deepseek-ai/flashmla/llms.txt

Helper functions for quantizing and dequantizing KV cache to FP8 format with per-tile scale factors. Essential for using sparse attention kernels which require FP8 KV cache format.

```APIDOC
## FP8 KV Cache Quantization Utilities

### Description
Helper functions for quantizing and dequantizing KV cache to FP8 format with per-tile scale factors. Essential for using sparse attention kernels which require FP8 KV cache format.

### Method
N/A (Python functions)

### Endpoint
N/A (Python functions)

### Parameters
#### Path Parameters
N/A

#### Query Parameters
N/A

#### Request Body
- **quantize_k_cache**: 
  - **k_cache_bf16** (torch.Tensor) - Original BF16 KV cache.
  - **layout** (FP8KVCacheLayout) - KV cache layout configuration.
- **dequantize_k_cache**: 
  - **k_cache_fp8** (torch.Tensor) - FP8 KV cache.
  - **layout** (FP8KVCacheLayout) - KV cache layout configuration.
- **abs_indices2indices_in_kvcache**: 
  - **abs_indices** (torch.Tensor) - Logical token positions.
  - **block_table** (torch.Tensor) - Block table mapping.
  - **block_size** (int) - Size of each block.

### Request Example
```python
import torch
from tests.quant import (
    FP8KVCacheLayout,
    quantize_k_cache,
    dequantize_k_cache,
    abs_indices2indices_in_kvcache
)

# KV cache layout configurations
# V32_FP8Sparse: d=576, d_nope=512, d_rope=64, tile_size=128, num_tiles=4
# MODEL1_FP8Sparse: d=512, d_nope=448, d_rope=64, tile_size=64, num_tiles=7
layout = FP8KVCacheLayout.V32_FP8Sparse

# Original BF16 KV cache
num_blocks = 1024
block_size = 64
head_dim = 576
k_cache_bf16 = torch.randn(num_blocks, block_size, 1, head_dim,
                           dtype=torch.bfloat16, device='cuda')

# Quantize to FP8 format
# Output format for V32_FP8Sparse (656 bytes per token):
# - Bytes 0-511: quantized NoPE (512 x float8_e4m3)
# - Bytes 512-527: scale factors (4 x float32)
# - Bytes 528-655: RoPE part (64 x bfloat16, unquantized)
k_cache_fp8 = quantize_k_cache(k_cache_bf16, layout)

# Dequantize back to BF16 for verification
k_cache_recovered = dequantize_k_cache(k_cache_fp8, layout)

# Convert logical indices to physical KV cache indices
# abs_indices: logical token positions (0 to seq_len-1)
# Returns: physical indices = (block_index * block_size) + offset
batch_size = 32
seq_len_q = 2
topk = 2048
abs_indices = torch.randint(0, num_blocks * block_size,
                            (batch_size, seq_len_q, topk),
                            dtype=torch.int32, device='cuda')

block_table = torch.arange(num_blocks, dtype=torch.int32, device='cuda').view(
    batch_size, -1)

indices_in_kvcache = abs_indices2indices_in_kvcache(
    abs_indices=abs_indices,
    block_table=block_table,
    block_size=block_size
)
# Use indices_in_kvcache with flash_mla_with_kvcache indices parameter
```

### Response
#### Success Response (200)
- **quantize_k_cache**: Returns FP8 quantized KV cache.
- **dequantize_k_cache**: Returns BF16 dequantized KV cache.
- **abs_indices2indices_in_kvcache**: Returns physical indices in KV cache.

#### Response Example
N/A (Python function return values)
```

--------------------------------

### Seesaw Scheduling: Step 7 - Softmax on p1

Source: https://github.com/deepseek-ai/flashmla/blob/main/docs/20250422-new-kernel-deep-dive.md

Applies the softmax function to the p1 vector using the calculated maximum (m_new1). This normalizes p1 for accumulation into the right output matrix.

```math
`\vec p_1 \gets \exp(\vec p_1 - m\_new_1)`
```

--------------------------------

### Seesaw Scheduling: Step 11 - Update oL with scaled p1

Source: https://github.com/deepseek-ai/flashmla/blob/main/docs/20250422-new-kernel-deep-dive.md

Updates the left output matrix (oL) by scaling the previous oL with scale1 and adding the product of the softmax-normalized p1 and the V1L matrix. This completes the accumulation for the left output.

```math
`\vec o_L \gets \vec o_L \cdot scale_1 + \vec p_1 V_{1L}`
```

--------------------------------

### Sparse Attention Prefill with Flash-MLA

Source: https://context7.com/deepseek-ai/flashmla/llms.txt

This kernel performs sparse attention prefill, allowing queries to attend to specific KV tokens defined by indices. It's designed for long-context scenarios and does not support batch dimensions directly.

```python
import torch
from flash_mla import flash_mla_sparse_fwd

# Sparse prefill parameters
seq_len_q = 4096
seq_len_kv = 32768
topk = 2048
num_heads_q = 128
num_heads_kv = 1
head_dim = 576  # d_qk
head_dim_v = 512

# Input tensors
q = torch.randn(seq_len_q, num_heads_q, head_dim,
                dtype=torch.bfloat16, device='cuda')
k = torch.randn(seq_len_kv, num_heads_kv, head_dim,
                 dtype=torch.bfloat16, device='cuda')

# Indices specifying which KV tokens each query attends to
# Shape: (seq_len_q, num_heads_kv, topk)
# Use -1 or values >= seq_len_kv for invalid indices
indices = torch.randint(0, seq_len_kv, (seq_len_q, num_heads_kv, topk),
                        dtype=torch.int32, device='cuda')

# Softmax scale factor
sm_scale = head_dim ** (-0.5)

# Optional: attention sink for output scaling
attn_sink = torch.zeros(num_heads_q, dtype=torch.float32, device='cuda')

# Optional: variable topk length per query
topk_length = torch.full((seq_len_q,), topk, dtype=torch.int32, device='cuda')

# Run sparse prefill
out, max_logits, lse = flash_mla_sparse_fwd(
    q=q,
    kv=kv,
    indices=indices,
    sm_scale=sm_scale,
    d_v=head_dim_v,
    attn_sink=attn_sink,
    topk_length=topk_length
)
# out: (seq_len_q, num_heads_q, head_dim_v) - attention output
# max_logits: (seq_len_q, num_heads_q) - maximum logit per head
# lse: (seq_len_q, num_heads_q) - log-sum-exp of attention scores

# Note: This kernel does not support batch dimension
# For multi-batch, reshape inputs and adjust indices accordingly
```

--------------------------------

### Seesaw Scheduling: Step 9 - Scale p0

Source: https://github.com/deepseek-ai/flashmla/blob/main/docs/20250422-new-kernel-deep-dive.md

Scales the softmax-normalized p0 vector by scale1. This adjustment is necessary before its contribution to the right output matrix is calculated.

```math
`\vec p_0 \gets \vec p_0 \cdot scale_1`
```

--------------------------------

### Sparse MLA Prefill Kernel

Source: https://github.com/deepseek-ai/flashmla/blob/main/README.md

Use flash_mla_sparse_fwd for sparse MLA prefill. This kernel does not support a batch dimension; simulate batch processing by reshaping inputs and adjusting indices. Invalid indices should be set to -1 or any number >= s_kv.

```python
Q: [s_q, h_q, d_qk], bfloat16
kv: [s_kv, h_kv, d_qk], bfloat16
indices: [s_q, h_kv, topk], int32

kv = kv.squeeze(1)  # [s_kv, d_qk], h_kv must be 1
indices = indices.squeeze(1)    # [s_q, topk]
focused_kv = kv[indices]    # For the i-th sequence (s_q), the corresponding KV tokens are selected from the KV cache based on indices[i, :]. This operation results in a tensor of shape [s_q, topk, d_qk].

P = (Q @ focused_kv.transpose(-1, -2)) * sm_scale * math.log2(math.e)    # [s_q, h_q, topk]
max_logits = P.max(dim=-1) # [s_q, h_q]
lse = log2sumexp2(P, dim=-1, base=2)   # [s_q, h_q]，"log2sumexp2" means that the exponentiation and logarithm are base-2
S = exp2(P - lse)      # [s_q, h_q, topk]
out = S @ focused_kv  # [s_q, h_q, d_qk]

return (out, max_logits, lse)
```

--------------------------------

### flash_attn_varlen_kvpacked_func

Source: https://context7.com/deepseek-ai/flashmla/llms.txt

Attention with separate Q and packed KV tensors. Enables efficient memory layout when K and V share the same sequence but Q has different length.

```APIDOC
## flash_attn_varlen_kvpacked_func

### Description
Attention with separate Q and packed KV tensors. Enables efficient memory layout when K and V share the same sequence but Q has different length.

### Method
N/A (Python function)

### Endpoint
N/A (Python function)

### Parameters
#### Path Parameters
N/A

#### Query Parameters
N/A

#### Request Body
- **q** (torch.Tensor) - Separate Q tensor.
- **kv** (torch.Tensor) - Packed KV tensor: (total_kv, num_heads, head_dim_qk + head_dim_v).
- **cu_seqlens_qo** (torch.Tensor) - Cumulative sequence lengths for query.
- **cu_seqlens_kv** (torch.Tensor) - Cumulative sequence lengths for key/value.
- **max_seqlen_qo** (int) - Maximum sequence length for query.
- **max_seqlen_kv** (int) - Maximum sequence length for key/value.
- **head_dim_qk** (int) - Dimension of the query and key heads.
- **dropout_p** (float) - Dropout probability.
- **softmax_scale** (float) - Scaling factor for softmax.
- **causal** (bool) - Whether to apply causal attention.
- **is_varlen** (bool) - Whether the input is variable length.

### Request Example
```python
import torch
from flash_mla import flash_attn_varlen_kvpacked_func

# Configuration with different Q and KV lengths
batch_size = 2
seqlens_q = torch.tensor([1024, 2048], dtype=torch.int32)
seqlens_kv = torch.tensor([4096, 8192], dtype=torch.int32)
total_q = seqlens_q.sum().item()
total_kv = seqlens_kv.sum().item()

cu_seqlens_q = torch.cumsum(
    torch.nn.functional.pad(seqlens_q, (1, 0)), 0, dtype=torch.int32
).cuda()
cu_seqlens_kv = torch.cumsum(
    torch.nn.functional.pad(seqlens_kv, (1, 0)), 0, dtype=torch.int32
).cuda()

num_heads = 128
head_dim_qk = 128
head_dim_v = 128

# Separate Q tensor
q = torch.randn(total_q, num_heads, head_dim_qk,
                dtype=torch.bfloat16, device='cuda')

# Packed KV tensor: (total_kv, num_heads, head_dim_qk + head_dim_v)
kv = torch.randn(total_kv, num_heads, head_dim_qk + head_dim_v,
                 dtype=torch.bfloat16, device='cuda')

# Run attention with KV-packed input
out, lse = flash_attn_varlen_kvpacked_func(
    q=q,
    kv=kv,
    cu_seqlens_qo=cu_seqlens_q,
    cu_seqlens_kv=cu_seqlens_kv,
    max_seqlen_qo=seqlens_q.max().item(),
    max_seqlen_kv=seqlens_kv.max().item(),
    head_dim_qk=head_dim_qk,
    dropout_p=0.0,
    softmax_scale=head_dim_qk ** (-0.5),
    causal=True,
    is_varlen=True
)
```

### Response
#### Success Response (200)
- **out** (torch.Tensor) - Output tensor.
- **lse** (torch.Tensor) - Log-sum-exp values.

#### Response Example
N/A (Python function return values)
```

--------------------------------

### Flash Attention with Packed QKV Input

Source: https://context7.com/deepseek-ai/flashmla/llms.txt

Use this function when Q, K, and V have the same sequence length and are stored contiguously in a single tensor. Ensure cu_seqlens and max_seqlen are correctly configured.

```python
import torch
from flash_mla import flash_attn_varlen_qkvpacked_func

# Configuration
batch_size = 2
seqlens = torch.tensor([4096, 8192], dtype=torch.int32)
total_len = seqlens.sum().item()
cu_seqlens = torch.cumsum(
    torch.nn.functional.pad(seqlens, (1, 0)), 0, dtype=torch.int32
).cuda()

num_heads = 128
head_dim_qk = 128
head_dim_v = 128
max_seqlen = seqlens.max().item()

# Packed QKV tensor: (total_len, num_heads, head_dim_qk + head_dim_qk + head_dim_v)
qkv = torch.randn(total_len, num_heads, head_dim_qk * 2 + head_dim_v,
                  dtype=torch.bfloat16, device='cuda')

# Run attention with packed input
out, lse = flash_attn_varlen_qkvpacked_func(
    qkv=qkv,
    cu_seqlens=cu_seqlens,
    max_seqlen=max_seqlen,
    head_dim_qk=head_dim_qk,
    dropout_p=0.0,
    softmax_scale=head_dim_qk ** (-0.5),
    causal=True,
    is_varlen=True
)
```

--------------------------------

### get_mla_metadata

Source: https://context7.com/deepseek-ai/flashmla/llms.txt

Initializes and returns scheduler metadata for tile scheduler management. This should be called once before the decoding loop.

```APIDOC
## GET /get_mla_metadata

### Description
Returns an empty FlashMLASchedMeta instance for tile scheduler metadata management. The actual scheduling metadata is generated during the first invocation of flash_mla_with_kvcache, enabling automatic optimization based on input tensor shapes and sequence lengths.

### Method
GET

### Endpoint
/get_mla_metadata

### Parameters
None

### Request Example
```python
import torch
from flash_mla import get_mla_metadata, flash_mla_with_kvcache

# Initialize scheduler metadata (call once before decoding loop)
tile_scheduler_metadata, num_splits = get_mla_metadata()

# The metadata can be reused across multiple layers and decoding steps
# as long as tensor shapes and cache_seqlens remain consistent
for layer_idx in range(num_layers):
    out, lse = flash_mla_with_kvcache(
        q=query_tensor,
        k_cache=kv_cache,
        block_table=block_table,
        cache_seqlens=cache_seqlens,
        head_dim_v=512,
        tile_scheduler_metadata=tile_scheduler_metadata,
        num_splits=num_splits,
        causal=False
    )
```

### Response
#### Success Response (200)
- **tile_scheduler_metadata** (FlashMLASchedMeta) - Metadata for tile scheduling.
- **num_splits** (int) - Number of splits for scheduling.

#### Response Example
```json
{
  "tile_scheduler_metadata": "<metadata_object>",
  "num_splits": 16
}
```
```

--------------------------------

### flash_mla_with_kvcache (Dense Attention)

Source: https://context7.com/deepseek-ai/flashmla/llms.txt

Performs MLA decoding attention with paged KV cache support for dense attention patterns.

```APIDOC
## POST /flash_mla_with_kvcache (Dense Attention)

### Description
Performs MLA decoding attention with paged KV cache support. Supports dense attention (standard MLA). Returns attention output and log-sum-exp values for each query head.

### Method
POST

### Endpoint
/flash_mla_with_kvcache

### Parameters
#### Query Parameters
- **head_dim_v** (int) - Required - The dimension of the value head.
- **softmax_scale** (float) - Optional - Scaling factor for softmax.
- **causal** (bool) - Optional - Whether to apply causal mask. Defaults to False.
- **is_fp8_kvcache** (bool) - Optional - Whether the KV cache is in FP8 format. Defaults to False.

#### Request Body
- **q** (torch.Tensor) - Required - Query tensor. Shape: (batch_size, seq_len_q, num_heads_q, head_dim).
- **k_cache** (torch.Tensor) - Required - Key cache tensor. Shape: (num_blocks, block_size, num_heads_kv, head_dim).
- **block_table** (torch.Tensor) - Required - Block table mapping batch sequences to cache blocks. Shape: (batch_size, num_blocks_per_sequence).
- **cache_seqlens** (torch.Tensor) - Required - Current sequence lengths in the KV cache. Shape: (batch_size,).
- **tile_scheduler_metadata** (FlashMLASchedMeta) - Required - Metadata for tile scheduling.
- **num_splits** (int) - Required - Number of splits for scheduling.

### Request Example
```python
import torch
from flash_mla import get_mla_metadata, flash_mla_with_kvcache

# Setup for dense MLA decoding
batch_size = 128
seq_len_q = 1  # Single token decoding (or >1 for MTP speculative decoding)
num_heads_q = 128
num_heads_kv = 1
head_dim = 576  # d_qk for DeepSeek V3
head_dim_v = 512
block_size = 64
max_seq_len = 4096

# Initialize tensors
q = torch.randn(batch_size, seq_len_q, num_heads_q, head_dim,
                dtype=torch.bfloat16, device='cuda')
cache_seqlens = torch.full((batch_size,), max_seq_len, dtype=torch.int32, device='cuda')

# Paged KV cache: (num_blocks, block_size, num_heads_kv, head_dim)
num_blocks = batch_size * (max_seq_len // block_size)
k_cache = torch.randn(num_blocks, block_size, num_heads_kv, head_dim,
                        dtype=torch.bfloat16, device='cuda')

# Block table maps batch sequences to cache blocks
block_table = torch.arange(num_blocks, dtype=torch.int32, device='cuda').view(
    batch_size, -1)

# Get scheduler metadata
tile_scheduler_metadata, num_splits = get_mla_metadata()

# Run dense MLA decoding
out, lse = flash_mla_with_kvcache(
    q=q,
    k_cache=k_cache,
    block_table=block_table,
    cache_seqlens=cache_seqlens,
    head_dim_v=head_dim_v,
    tile_scheduler_metadata=tile_scheduler_metadata,
    num_splits=num_splits,
    softmax_scale=head_dim ** (-0.5),
    causal=False,  # Set True for causal attention mask
    is_fp8_kvcache=False
)
# out: (batch_size, seq_len_q, num_heads_q, head_dim_v)
# lse: (batch_size, num_heads_q, seq_len_q) - log-sum-exp of attention scores
```

### Response
#### Success Response (200)
- **out** (torch.Tensor) - Attention output tensor.
- **lse** (torch.Tensor) - Log-sum-exp values of attention scores.

#### Response Example
```json
{
  "out": "<tensor_data>",
  "lse": "<tensor_data>"
}
```
```

--------------------------------

### Flash Attention with Packed KV Input

Source: https://context7.com/deepseek-ai/flashmla/llms.txt

Use this function when K and V share the same sequence length but Q has a different length. This enables efficient memory layout. Ensure sequence lengths and cumulative sequence lengths for both Q and KV are correctly provided.

```python
import torch
from flash_mla import flash_attn_varlen_kvpacked_func

# Configuration with different Q and KV lengths
batch_size = 2
seqlens_q = torch.tensor([1024, 2048], dtype=torch.int32)
seqlens_kv = torch.tensor([4096, 8192], dtype=torch.int32)
total_q = seqlens_q.sum().item()
total_kv = seqlens_kv.sum().item()

cu_seqlens_q = torch.cumsum(
    torch.nn.functional.pad(seqlens_q, (1, 0)), 0, dtype=torch.int32
).cuda()
cu_seqlens_kv = torch.cumsum(
    torch.nn.functional.pad(seqlens_kv, (1, 0)), 0, dtype=torch.int32
).cuda()

num_heads = 128
head_dim_qk = 128
head_dim_v = 128

# Separate Q tensor
q = torch.randn(total_q, num_heads, head_dim_qk,
                dtype=torch.bfloat16, device='cuda')

# Packed KV tensor: (total_kv, num_heads, head_dim_qk + head_dim_v)
kv = torch.randn(total_kv, num_heads, head_dim_qk + head_dim_v,
                 dtype=torch.bfloat16, device='cuda')

# Run attention with KV-packed input
out, lse = flash_attn_varlen_kvpacked_func(
    q=q,
    kv=kv,
    cu_seqlens_qo=cu_seqlens_q,
    cu_seqlens_kv=cu_seqlens_kv,
    max_seqlen_qo=seqlens_q.max().item(),
    max_seqlen_kv=seqlens_kv.max().item(),
    head_dim_qk=head_dim_qk,
    dropout_p=0.0,
    softmax_scale=head_dim_qk ** (-0.5),
    causal=True,
    is_varlen=True
)
```

--------------------------------

### flash_mla_with_kvcache (Sparse Attention with FP8 KV Cache)

Source: https://context7.com/deepseek-ai/flashmla/llms.txt

Performs MLA decoding attention with paged KV cache support, specifically for sparse attention patterns using an FP8 quantized KV cache.

```APIDOC
## POST /flash_mla_with_kvcache (Sparse Attention FP8)

### Description
Performs MLA decoding attention with paged KV cache support. Supports sparse attention (DSA) with FP8 quantized KV cache. Enables token-level sparse attention using the indices tensor for selective attention computation.

### Method
POST

### Endpoint
/flash_mla_with_kvcache

### Parameters
#### Query Parameters
- **head_dim_v** (int) - Required - The dimension of the value head.
- **softmax_scale** (float) - Optional - Scaling factor for softmax.
- **causal** (bool) - Optional - Whether to apply causal mask. Defaults to False.
- **is_fp8_kvcache** (bool) - Required - Indicates that the KV cache is in FP8 format.
- **topk** (int) - Required - Number of tokens to attend to per query for sparse attention.

#### Request Body
- **q** (torch.Tensor) - Required - Query tensor. Shape: (batch_size, seq_len_q, num_heads_q, head_dim).
- **k_cache** (torch.Tensor) - Required - FP8 quantized Key cache tensor. Specific format required.
- **block_table** (torch.Tensor) - Required - Block table mapping batch sequences to cache blocks.
- **cache_seqlens** (torch.Tensor) - Required - Current sequence lengths in the KV cache.
- **indices** (torch.Tensor) - Required - Tensor containing indices for sparse attention.
- **tile_scheduler_metadata** (FlashMLASchedMeta) - Required - Metadata for tile scheduling.
- **num_splits** (int) - Required - Number of splits for scheduling.

### Request Example
```python
import torch
from flash_mla import get_mla_metadata, flash_mla_with_kvcache

# Sparse attention with FP8 KV cache for DeepSeek V3.2
batch_size = 128
seq_len_q = 2  # MTP speculative decoding
num_heads_q = 128
topk = 2048  # Number of tokens to attend to per query

# Query tensor
q = torch.randn(batch_size, seq_len_q, num_heads_q, 576,
                dtype=torch.bfloat16, device='cuda')

# FP8 KV cache format (656 bytes per token for V3.2):
# - First 512 bytes: quantized NoPE part (512 x float8_e4m3)
# - Next 16 bytes: scale factors (4 x float32)
# - Last 128 bytes: RoPE part (64 x bfloat16, not quantized)
num_blocks = 1024
block_size = 64
bytes_per_token = 656
k_cache_fp8 = torch.empty(num_blocks, block_size, 1, bytes_per_token,
                          dtype=torch.float8_e4m3fn, device='cuda')

# Indices tensor for sparse attention
# Example: indices = torch.randint(0, block_size, (batch_size, seq_len_q, topk), device='cuda')

# Assume block_table, cache_seqlens, tile_scheduler_metadata, num_splits are initialized

# Run sparse MLA decoding with FP8 KV cache
out, lse = flash_mla_with_kvcache(
    q=q,
    k_cache=k_cache_fp8,
    block_table=block_table,
    cache_seqlens=cache_seqlens,
    head_dim_v=512,
    tile_scheduler_metadata=tile_scheduler_metadata,
    num_splits=num_splits,
    is_fp8_kvcache=True,
    topk=topk,
    causal=False
)
```

### Response
#### Success Response (200)
- **out** (torch.Tensor) - Attention output tensor.
- **lse** (torch.Tensor) - Log-sum-exp values of attention scores.

#### Response Example
```json
{
  "out": "<tensor_data>",
  "lse": "<tensor_data>"
}
```
```

--------------------------------

### MLA Decoding with KV Cache

Source: https://github.com/deepseek-ai/flashmla/blob/main/README.md

Call get_mla_metadata once before the decoding loop and flash_mla_with_kvcache in each decoding step. Ensure correct parameters for cache_seqlens, head dimensions, and FP8 usage.

```python
from flash_mla import get_mla_metadata, flash_mla_with_kvcache

tile_scheduler_metadata, num_splits = get_mla_metadata(
    cache_seqlens,
    s_q * h_q // h_kv,
    h_kv,
    h_q,
    is_fp8,
    topk,
)

for i in range(num_layers):
    ...
    o_i, lse_i = flash_mla_with_kvcache(
        q_i, kvcache_i, block_table, cache_seqlens, dv,
        tile_scheduler_metadata, num_splits,
        is_causal, is_fp8_kvcache, indices,
    )
    ...
```

--------------------------------

### Sparse Attention Decoding with KV Cache

Source: https://context7.com/deepseek-ai/flashmla/llms.txt

This snippet demonstrates decoding with a KV cache using sparse attention. It requires pre-calculated tile scheduler metadata and is not suitable for causal attention.

```python
indices = torch.randint(0, num_blocks * block_size,
                        (batch_size, seq_len_q, topk),
                        dtype=torch.int32, device='cuda')

attn_sink = torch.zeros(num_heads_q, dtype=torch.float32, device='cuda')

tile_scheduler_metadata, num_splits = get_mla_metadata()

out, lse = flash_mla_with_kvcache(
    q=q,
    k_cache=k_cache_fp8,
    block_table=None,  # Not needed for sparse attention
    cache_seqlens=None,  # Not needed for sparse attention
    head_dim_v=512,
    tile_scheduler_metadata=tile_scheduler_metadata,
    num_splits=num_splits,
    causal=False,  # Must be False for sparse attention
    is_fp8_kvcache=True,
    indices=indices,
    attn_sink=attn_sink
)
```

--------------------------------

### flash_attn_varlen_func

Source: https://context7.com/deepseek-ai/flashmla/llms.txt

Standard dense Multi-Head Attention (MHA) forward and backward operations with variable-length sequence support. Compatible with flash_attn API, optimized for SM100 (Blackwell) achieving up to 1460 TFlops in forward pass.

```APIDOC
## flash_attn_varlen_func

### Description
Standard dense Multi-Head Attention (MHA) forward and backward operations with variable-length sequence support. Compatible with flash_attn API, optimized for SM100 (Blackwell) achieving up to 1460 TFlops in forward pass.

### Method
`flash_attn_varlen_func`

### Parameters
#### Query Tensor (q)
- **q** (torch.Tensor) - Packed query tensor. Shape: (total_q, num_heads, head_dim_qk). Data type: torch.bfloat16. Device: cuda.

#### Key Tensor (k)
- **k** (torch.Tensor) - Packed key tensor. Shape: (total_k, num_heads_kv, head_dim_qk). Data type: torch.bfloat16. Device: cuda.

#### Value Tensor (v)
- **v** (torch.Tensor) - Packed value tensor. Shape: (total_k, num_heads_kv, head_dim_v). Data type: torch.bfloat16. Device: cuda.

#### Cumulative Sequence Lengths for Query (cu_seqlens_qo)
- **cu_seqlens_qo** (torch.Tensor) - Cumulative sequence lengths for query. Shape: (batch_size + 1,). Data type: torch.int32. Device: cuda.

#### Cumulative Sequence Lengths for Key-Value (cu_seqlens_kv)
- **cu_seqlens_kv** (torch.Tensor) - Cumulative sequence lengths for key-value. Shape: (batch_size + 1,). Data type: torch.int32. Device: cuda.

#### Maximum Sequence Length for Query (max_seqlen_qo)
- **max_seqlen_qo** (int) - Maximum sequence length among queries in the batch.

#### Maximum Sequence Length for Key-Value (max_seqlen_kv)
- **max_seqlen_kv** (int) - Maximum sequence length among key-values in the batch.

#### Dropout Probability
- **dropout_p** (float) - Dropout probability. Must be 0.0 for this function.

#### Softmax Scale
- **softmax_scale** (float) - Softmax scale factor. Typically `head_dim_qk ** (-0.5)`.

#### Causal Masking
- **causal** (bool) - Enable causal masking. Set to `True` for causal attention.

#### Variable Length Mode
- **is_varlen** (bool) - Enable variable-length mode. Must be `True`.

### Request Example
```python
import torch
from flash_mla import flash_attn_varlen_func

# Variable-length batch configuration
batch_size = 2
seqlens_q = torch.tensor([4096, 8192], dtype=torch.int32)
seqlens_k = torch.tensor([4096, 8192], dtype=torch.int32)
total_q = seqlens_q.sum().item()
total_k = seqlens_k.sum().item()

# Cumulative sequence lengths (starts with 0)
cu_seqlens_q = torch.cumsum(torch.nn.functional.pad(seqlens_q, (1, 0)), 0, dtype=torch.int32).cuda()
cu_seqlens_k = torch.cumsum(torch.nn.functional.pad(seqlens_k, (1, 0)), 0, dtype=torch.int32).cuda()

# MHA configuration
num_heads = 128
num_heads_kv = 128  # MHA mode (equal heads)
head_dim_qk = 128
head_dim_v = 128
max_seqlen_q = seqlens_q.max().item()
max_seqlen_k = seqlens_k.max().item()

# Input tensors (packed variable-length format)
q = torch.randn(total_q, num_heads, head_dim_qk, dtype=torch.bfloat16, device='cuda')
k = torch.randn(total_k, num_heads_kv, head_dim_qk, dtype=torch.bfloat16, device='cuda')
v = torch.randn(total_k, num_heads_kv, head_dim_v, dtype=torch.bfloat16, device='cuda')

# Run flash attention with variable-length support
out, lse = flash_attn_varlen_func(
    q=q,
    k=k,
    v=v,
    cu_seqlens_qo=cu_seqlens_q,
    cu_seqlens_kv=cu_seqlens_k,
    max_seqlen_qo=max_seqlen_q,
    max_seqlen_kv=max_seqlen_k,
    dropout_p=0.0,  # Dropout must be 0.0
    softmax_scale=head_dim_qk ** (-0.5),
    causal=True,  # Enable causal masking
    is_varlen=True  # Variable-length mode
)
```

### Response
#### Success Response
- **out** (torch.Tensor) - Attention output. Shape: (total_q, num_heads, head_dim_v).
- **lse** (torch.Tensor) - Log-sum-exp values. Shape: (total_q, num_heads).

### Notes
- Backward pass is supported for `num_heads_q == num_heads_kv`. GQA backward is not yet supported on SM100.
```

--------------------------------

### Variable-Length MHA with Flash-Attention

Source: https://context7.com/deepseek-ai/flashmla/llms.txt

This function implements standard dense Multi-Head Attention with support for variable-length sequences. It requires pre-calculated cumulative sequence lengths and is optimized for SM100 hardware. Dropout must be set to 0.0.

```python
import torch
from flash_mla import flash_attn_varlen_func

# Variable-length batch configuration
batch_size = 2
seqlens_q = torch.tensor([4096, 8192], dtype=torch.int32)
seqlens_k = torch.tensor([4096, 8192], dtype=torch.int32)
total_q = seqlens_q.sum().item()
total_k = seqlens_k.sum().item()

# Cumulative sequence lengths (starts with 0)
cu_seqlens_q = torch.cumsum(
    torch.nn.functional.pad(seqlens_q, (1, 0)), 0, dtype=torch.int32
).cuda()
cu_seqlens_k = torch.cumsum(
    torch.nn.functional.pad(seqlens_k, (1, 0)), 0, dtype=torch.int32
).cuda()

# MHA configuration
num_heads = 128
num_heads_kv = 128  # MHA mode (equal heads)
head_dim_qk = 128
head_dim_v = 128
max_seqlen_q = seqlens_q.max().item()
max_seqlen_k = seqlens_k.max().item()

# Input tensors (packed variable-length format)
q = torch.randn(total_q, num_heads, head_dim_qk,
                dtype=torch.bfloat16, device='cuda')
k = torch.randn(total_k, num_heads_kv, head_dim_qk,
                dtype=torch.bfloat16, device='cuda')
v = torch.randn(total_k, num_heads_kv, head_dim_v,
                dtype=torch.bfloat16, device='cuda')

# Run flash attention with variable-length support
out, lse = flash_attn_varlen_func(
    q=q,
    k=k,
    v=v,
    cu_seqlens_qo=cu_seqlens_q,
    cu_seqlens_kv=cu_seqlens_k,
    max_seqlen_qo=max_seqlen_q,
    max_seqlen_kv=max_seqlen_k,
    dropout_p=0.0,  # Dropout must be 0.0
    softmax_scale=head_dim_qk ** (-0.5),
    causal=True,  # Enable causal masking
    is_varlen=True  # Variable-length mode
)
# out: (total_q, num_heads, head_dim_v)
# lse: (total_q, num_heads) - log-sum-exp values

# Backward pass is supported for same num_heads_q == num_heads_kv
# GQA backward is not yet supported on SM100
```

--------------------------------

### Dense MHA Prefill Functions

Source: https://github.com/deepseek-ai/flashmla/blob/main/README.md

Implements standard dense Multi-Head Attention (MHA) forward and backward operations. Usage is similar to the flash_attn package.

```python
flash_attn_varlen_func
flash_attn_varlen_qkvpacked_func
flash_attn_varlen_kvpacked_func
```