FlashInfer

https://docs.flashinfer.ai/
Admin
FlashInfer is a library and kernel generator for Large Language Models that provides...
Tokens:209,340
Snippets:1,072
Trust Score:7.8
Update:5 months ago
Show doc for...
Tokens:
Raw
### Install FlashInfer from Source

Source: https://docs.flashinfer.ai/_sources/installation.rst.txt

Install FlashInfer after cloning the repository. Navigate to the cloned directory and run the installation command.

```bash
cd flashinfer
python -m pip install -v .
```

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

### Build and Install FlashInfer-Cubin from Source

Source: https://docs.flashinfer.ai/_sources/installation.rst.txt

Build the flashinfer-cubin package from its source directory and then install the generated wheel file.

```bash
cd flashinfer-cubin
python -m build --no-isolation --wheel
python -m pip install dist/*.whl
```

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

### Install FlashInfer Nightly Builds

Source: https://docs.flashinfer.ai/_sources/installation.rst.txt

Install the latest nightly builds for testing. This involves installing the core, cubin, and JIT cache packages from a specific nightly index URL.

```bash
# Core and cubin packages
pip install -U --pre flashinfer-python --index-url https://flashinfer.ai/whl/nightly/ --no-deps # Install the nightly package from custom index, without installing dependencies
pip install flashinfer-python  # Install flashinfer-python's dependencies from PyPI
pip install -U --pre flashinfer-cubin --index-url https://flashinfer.ai/whl/nightly/
# JIT cache package (replace cu129 with your CUDA version: cu128, cu129, or cu130)
pip install -U --pre flashinfer-jit-cache --index-url https://flashinfer.ai/whl/nightly/cu129
```

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

### Batch Decode Attention Example

Source: https://docs.flashinfer.ai/api/attention.html

A complete example demonstrating the setup and usage of `BatchDecodeWithPagedKVCacheWrapper` for batch decoding. It includes allocating workspace, initializing the wrapper, preparing KV cache data, planning the attention, and running the attention for multiple layers.

```python
import torch
import flashinfer
num_layers = 32
num_qo_heads = 64
num_kv_heads = 8
head_dim = 128
max_num_pages = 128
page_size = 16
# allocate 128MB workspace buffer
workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
decode_wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
    workspace_buffer, "NHD"
)
batch_size = 7
kv_page_indices = torch.arange(max_num_pages).int().to("cuda:0")
kv_page_indptr = torch.tensor(
    [0, 17, 29, 44, 48, 66, 100, 128], dtype=torch.int32, device="cuda:0"
)
# 1 <= kv_last_page_len <= page_size
kv_last_page_len = torch.tensor(
    [1, 7, 14, 4, 3, 1, 16], dtype=torch.int32, device="cuda:0"
)
kv_cache_at_layer = [
    torch.randn(
        max_num_pages, 2, page_size, num_kv_heads, head_dim, dtype=torch.float16, device="cuda:0"
    ) for _ in range(num_layers)
]
# create auxiliary data structures for batch decode attention
decode_wrapper.plan(
    kv_page_indptr,
    kv_page_indices,
    kv_last_page_len,
    num_qo_heads,
    num_kv_heads,
    head_dim,
    page_size,
    pos_encoding_mode="NONE",
    data_type=torch.float16
)
outputs = []
for i in range(num_layers):
    q = torch.randn(batch_size, num_qo_heads, head_dim).half().to("cuda:0")
    kv_cache = kv_cache_at_layer[i]
    # compute batch decode attention, reuse auxiliary data structures for all layers
    o = decode_wrapper.run(q, kv_cache)
    outputs.append(o)

outputs[0].shape
```

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

### Example Usage of BatchPrefillWithRaggedKVCacheWrapper

Source: https://docs.flashinfer.ai/api/attention.html

Demonstrates the complete workflow of initializing the wrapper, planning attention with causal masking, running attention for multiple layers, and verifying output shapes. Also shows an example with custom masks.

```python
import torch
import flashinfer
num_layers = 32
num_qo_heads = 64
num_kv_heads = 16
head_dim = 128
# allocate 128MB workspace buffer
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
prefill_wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper(
    workspace_buffer, "NHD"
)
batch_size = 7
nnz_kv = 100
nnz_qo = 100
qo_indptr = torch.tensor(
    [0, 33, 44, 55, 66, 77, 88, nnz_qo], dtype=torch.int32, device="cuda:0"
)
kv_indptr = qo_indptr.clone()
q_at_layer = torch.randn(num_layers, nnz_qo, num_qo_heads, head_dim).half().to("cuda:0")
k_at_layer = torch.randn(num_layers, nnz_kv, num_kv_heads, head_dim).half().to("cuda:0")
v_at_layer = torch.randn(num_layers, nnz_kv, num_kv_heads, head_dim).half().to("cuda:0")
# create auxiliary data structures for batch prefill attention
prefill_wrapper.plan(
    qo_indptr,
    kv_indptr,
    num_qo_heads,
    num_kv_heads,
    head_dim,
    causal=True,
)
outputs = []
for i in range(num_layers):
    q = q_at_layer[i]
    k = k_at_layer[i]
    v = v_at_layer[i]
    # compute batch prefill attention, reuse auxiliary data structures
    o = prefill_wrapper.run(q, k, v)
    outputs.append(o)

outputs[0].shape

# below is another example of creating custom mask for batch prefill attention
mask_arr = []
qo_len = (qo_indptr[1:] - qo_indptr[:-1]).cpu().tolist()
kv_len = (kv_indptr[1:] - kv_indptr[:-1]).cpu().tolist()
for i in range(batch_size):
    mask_i = torch.tril(
        torch.full((qo_len[i], kv_len[i]), True, device="cuda:0"),
        diagonal=(kv_len[i] - qo_len[i]),
    )
    mask_arr.append(mask_i.flatten())

mask = torch.cat(mask_arr, dim=0)
prefill_wrapper.plan(
    qo_indptr,
    kv_indptr,
    num_qo_heads,
    num_kv_heads,
    head_dim,
    custom_mask=mask
)
outputs_custom_mask = []
for i in range(num_layers):
    q = q_at_layer[i]
    k = k_at_layer[i]
    v = v_at_layer[i]
    # compute batch prefill attention, reuse auxiliary data structures
    o_custom = prefill_wrapper.run(q, k, v)
    assert torch.allclose(o_custom, outputs[i], rtol=1e-3, atol=1e-3)

outputs_custom_mask[0].shape
```

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

### Test Environment Setup

Source: https://docs.flashinfer.ai/api/testing.html

Utilities for setting up a reproducible test environment.

```APIDOC
## Test Environment Setup

### Description
Utilities for setting up a reproducible test environment.

### Functions
- `set_seed(random_seed)`: Set random seed for reproducibility during testing.
- `sleep_after_kernel_run(execution_time)`: Sleep after kernel run.
```

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

### Install FlashInfer with Optional Packages

Source: https://docs.flashinfer.ai/_sources/installation.rst.txt

Install the core package along with optional packages for faster initialization and offline usage. The JIT cache package requires specifying your CUDA version.

```bash
pip install flashinfer-python flashinfer-cubin
# JIT cache package (replace cu129 with your CUDA version: cu128, cu129, or cu130)
pip install flashinfer-jit-cache --index-url https://flashinfer.ai/whl/cu129
```

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

### Build and Install FlashInfer-JIT-Cache from Source

Source: https://docs.flashinfer.ai/_sources/installation.rst.txt

Build the flashinfer-jit-cache package from its source directory. Customize FLASHINFER_CUDA_ARCH_LIST for your target GPUs before building.

```bash
export FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 10.3a 11.0a 12.0f"
cd flashinfer-jit-cache
python -m build --no-isolation --wheel
python -m pip install dist/*.whl
```

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

### Install FlashInfer Nightly Builds

Source: https://docs.flashinfer.ai/installation.html

Install nightly builds for FlashInfer. This includes the core and cubin packages. Dependencies are installed from PyPI.

```bash
pip install -U --pre flashinfer-python --index-url https://flashinfer.ai/whl/nightly/ --no-deps
pip install flashinfer-python
pip install -U --pre flashinfer-cubin --index-url https://flashinfer.ai/whl/nightly/
```

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

### Install FlashInfer with Optional Packages

Source: https://docs.flashinfer.ai/installation.html

Install FlashInfer with optional packages for faster initialization and offline usage. Includes the core package and pre-compiled kernel binaries.

```bash
pip install flashinfer-python flashinfer-cubin
```

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

### Create a LogitsPipe with TopK and Sample Processors

Source: https://docs.flashinfer.ai/generated/flashinfer.logits_processor.LogitsProcessor.html

This example demonstrates how to create a LogitsPipe with TopK and Sample processors. It assumes the input is probabilities and shows the resulting compiled operations.

```python
import torch
from flashinfer.logits_processor import LogitsPipe, TopK, Sample, TensorType
torch.manual_seed(42)

# Create a pipeline that legalizes to a fused op.
pipe = LogitsPipe([
    TopK(),         # Top-k filtering on logits
    Sample()        # Sample from the filtered distribution
], input_type=TensorType.PROBS)  # assume the input is probabilities

pipe
```

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

### allreduce_fusion - kARResidualRMSNormFP8Quant

Source: https://docs.flashinfer.ai/generated/flashinfer.comm.allreduce_fusion.html

This example shows how to use `allreduce_fusion` with FP8 quantization for the `kARResidualRMSNormFP8Quant` pattern, enabling memory and computation efficiency.

```APIDOC
## POST /flashinfer/allreduce_fusion

### Description
Performs AllReduce fusion with Residual, RMS Normalization, and FP8 quantization.

### Method
POST

### Endpoint
/flashinfer/allreduce_fusion

### Parameters
#### Request Body
- **input** (Tensor) - Required - Input tensor.
- **workspace** (Tensor) - Required - Workspace tensor.
- **pattern** (Enum) - Required - Fusion pattern, set to `AllReduceFusionPattern.kARResidualRMSNormFP8Quant`.
- **norm_out** (Tensor) - Optional - Output tensor for normalized result.
- **quant_out** (Tensor) - Optional - Output tensor for quantized result.
- **scale_out** (Tensor) - Optional - Output tensor for quantization scales.
- **residual_in** (Tensor) - Optional - Input tensor for residual connection.
- **rms_gamma** (Tensor) - Optional - Gamma parameter for RMS Normalization.
- **scale_factor** (Tensor) - Optional - Scaling factor for quantization.

### Request Example
```json
{
  "input": "hidden_states",
  "workspace": "workspace",
  "pattern": "kARResidualRMSNormFP8Quant",
  "quant_out": "quant",
  "scale_out": "scales",
  "residual_in": "residual",
  "rms_gamma": "norm_weight",
  "scale_factor": "scale_tensor"
}
```

### Response
#### Success Response (200)
- **output** (Tensor) - The final result.

#### Response Example
```json
{
  "output": "quant"
}
```
```

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

### Test Environment Setup

Source: https://docs.flashinfer.ai/_sources/api/testing.rst.txt

Utilities for setting up the testing environment, including seeding random number generators and controlling execution flow.

```APIDOC
## Test Environment Setup

### Description
Utilities for setting up the testing environment.

### Functions
- **set_seed(seed: int)**: Sets the random seed for reproducibility.
- **sleep_after_kernel_run(seconds: float)**: Pauses execution for a specified duration after a kernel runs.
```

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

### Install FlashInfer Core Package

Source: https://docs.flashinfer.ai/_sources/installation.rst.txt

Use this command to install the core FlashInfer Python package. It compiles or downloads kernels on first use.

```bash
pip install flashinfer-python
```

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

### Top-K Page Table Transform Example

Source: https://docs.flashinfer.ai/generated/flashinfer.top_k_page_table_transform.html

Demonstrates the usage of flashinfer.top_k_page_table_transform with sample tensors. Ensure CUDA is available and tensors are on the correct device.

```python
import torch
import flashinfer
num_rows = 8
max_len = 4096
k = 256
scores = torch.randn(num_rows, max_len, device="cuda", dtype=torch.float16)
src_page_table = torch.randint(0, 1000, (num_rows, max_len), device="cuda", dtype=torch.int32)
lengths = torch.full((num_rows,), max_len, device="cuda", dtype=torch.int32)
output = flashinfer.top_k_page_table_transform(scores, src_page_table, lengths, k)
output.shape
```

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

### Install FlashInfer Nightly JIT Cache

Source: https://docs.flashinfer.ai/installation.html

Install the nightly JIT cache package for FlashInfer. Replace 'cu129' with your CUDA version.

```bash
pip install -U --pre flashinfer-jit-cache --index-url https://flashinfer.ai/whl/nightly/cu129
```

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

### Upgrade Pip and Setuptools

Source: https://docs.flashinfer.ai/_sources/installation.rst.txt

If you encounter build isolation errors when installing from source in editable mode, upgrade pip and setuptools to their latest versions.

```bash
python -m pip install --upgrade pip setuptools
```

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

### Sampling from Logits with FlashInfer

Source: https://docs.flashinfer.ai/generated/flashinfer.logits_processor.Sample.html

Demonstrates how to use the Sample processor with logits as input. Ensure the input_type is set to TensorType.LOGITS. This example uses top_k=1 for sampling.

```python
>>> import torch
>>> from flashinfer.logits_processor import LogitsPipe, Sample, TensorType
>>> torch.manual_seed(42)
>>> 
>>> # Sampling from logits
>>> pipe = LogitsPipe([Sample(deterministic=True)], input_type=TensorType.LOGITS)
>>> logits = torch.randn(2, 5, device="cuda")
>>> logits
tensor([[ 0.1940,  2.1614, -0.1721,  0.8491, -1.9244],
        [ 0.6530, -0.6494, -0.8175,  0.5280, -1.2753]], device='cuda:0')
>>> tokens = pipe(logits, top_k=1)
>>> tokens
tensor([0, 1], device='cuda:0')
```

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

### allreduce_fusion - kARResidualRMSNorm

Source: https://docs.flashinfer.ai/generated/flashinfer.comm.allreduce_fusion.html

This example demonstrates the basic usage of `allreduce_fusion` for the `kARResidualRMSNorm` pattern, which combines AllReduce with Residual and RMS Normalization.

```APIDOC
## POST /flashinfer/allreduce_fusion

### Description
Performs AllReduce fusion with Residual and RMS Normalization.

### Method
POST

### Endpoint
/flashinfer/allreduce_fusion

### Parameters
#### Request Body
- **input** (Tensor) - Required - Input tensor.
- **workspace** (Tensor) - Required - Workspace tensor.
- **pattern** (Enum) - Required - Fusion pattern, set to `AllReduceFusionPattern.kARResidualRMSNorm`.
- **launch_with_pdl** (Boolean) - Optional - Whether to launch with PDL.
- **residual_out** (Tensor) - Optional - Output tensor for residual connection.
- **norm_out** (Tensor) - Optional - Output tensor for normalized result.
- **residual_in** (Tensor) - Optional - Input tensor for residual connection.
- **rms_gamma** (Tensor) - Optional - Gamma parameter for RMS Normalization.

### Request Example
```json
{
  "input": "hidden_states",
  "workspace": "workspace",
  "pattern": "kARResidualRMSNorm",
  "launch_with_pdl": true,
  "residual_out": "prenorm",
  "norm_out": "normed",
  "residual_in": "residual",
  "rms_gamma": "norm_weight"
}
```

### Response
#### Success Response (200)
- **output** (Tensor) - The final result, typically equivalent to `norm_out`.

#### Response Example
```json
{
  "output": "normed"
}
```
```

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

### LogitsPipe Pipeline Construction

Source: https://docs.flashinfer.ai/api/logits_processor.html

Demonstrates how to create and use a LogitsPipe for processing LLM outputs. The example shows initializing a pipeline with several processors and then applying it to sample logits.

```APIDOC
## LogitsPipe Pipeline Construction

### Description
Use `LogitsPipe` to create processing pipelines for LLM outputs. This example shows a pipeline with Temperature, Softmax, TopP, and Sample processors.

### Method
```python
pipe = LogitsPipe([
    Temperature(),      # Scale logits by temperature
    Softmax(),          # Convert logits to probabilities
    TopP(),             # Apply top-p filtering
    Sample()            # Sample from the distribution
])

# Apply the pipeline
batch_size = 4
vocab_size = 5
logits = torch.randn(batch_size, vocab_size, device="cuda")
output_ids = pipe(logits, temperature=0.7, top_p=0.9)
```

### Parameters
This section is not applicable for the pipeline construction example itself, as it focuses on the instantiation and usage of `LogitsPipe`.
```

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

### BatchPrefillWithRaggedKVCacheWrapper Initialization and Usage

Source: https://docs.flashinfer.ai/api/attention.html

Demonstrates how to initialize and use the BatchPrefillWithRaggedKVCacheWrapper for batch prefill attention with a ragged KV cache. Includes examples for both standard causal attention and attention with custom masks.

```APIDOC
## BatchPrefillWithRaggedKVCacheWrapper

### Description
Wrapper class for prefill/append attention with ragged (tensor) kv-cache for batch of requests. Check our tutorial for ragged kv-cache layout.

### Constructor
```python
_flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper(_float_workspace_buffer : Tensor_, _kv_layout : str = 'NHD'_, _use_cuda_graph : bool = False_, _qo_indptr_buf : Tensor | None = None_, _kv_indptr_buf : Tensor | None = None_, _custom_mask_buf : Tensor | None = None_, _mask_indptr_buf : Tensor | None = None_, _backend : str = 'auto'_, _jit_args : List[Any] | None = None_, _jit_kwargs : Dict[str, Any] | None = None_)
```

### Parameters
#### Constructor Parameters
- **_float_workspace_buffer** (torch.Tensor) - The user reserved float workspace buffer used to store intermediate attention results in the split-k algorithm. The recommended size is 128MB, the device of the workspace buffer should be the same as the device of the input tensors.
- **_kv_layout** (str, optional) - Layout of the KV cache. Defaults to 'NHD'.
- **_use_cuda_graph** (bool, optional) - Whether to use CUDA graph. Defaults to False.
- **_qo_indptr_buf** (Tensor | None, optional) - Pointer buffer for query/output.
- **_kv_indptr_buf** (Tensor | None, optional) - Pointer buffer for key/value.
- **_custom_mask_buf** (Tensor | None, optional) - Buffer for custom mask.
- **_mask_indptr_buf** (Tensor | None, optional) - Pointer buffer for custom mask.
- **_backend** (str, optional) - Backend to use. Defaults to 'auto'.
- **_jit_args** (List[Any] | None, optional) - JIT arguments.
- **_jit_kwargs** (Dict[str, Any] | None, optional) - JIT keyword arguments.

### Methods
#### plan
Plans the attention computation. This method should be called before `run`.

Parameters:
- **qo_indptr** (torch.Tensor) - Pointer array for query/output sequences.
- **kv_indptr** (torch.Tensor) - Pointer array for key/value sequences.
- **num_qo_heads** (int) - Number of query/output heads.
- **num_kv_heads** (int) - Number of key/value heads.
- **head_dim** (int) - Dimension of each head.
- **causal** (bool, optional) - Whether to apply causal masking. Defaults to False.
- **custom_mask** (torch.Tensor, optional) - Custom mask tensor.

#### run
Runs the prefill/append attention computation.

Parameters:
- **q** (torch.Tensor) - Query tensor.
- **k** (torch.Tensor) - Key tensor.
- **v** (torch.Tensor) - Value tensor.

Returns:
- **torch.Tensor** - The attention output.
- **Tuple[torch.Tensor, torch.Tensor]** - If `return_lse` is True, returns a tuple of (attention output, logsumexp).

### Example
```python
import torch
import flashinfer

num_layers = 32
num_qo_heads = 64
num_kv_heads = 16
head_dim = 128

# allocate 128MB workspace buffer
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
prefill_wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper(
    workspace_buffer, "NHD"
)

batch_size = 7
nnz_kv = 100
nnz_qo = 100

qo_indptr = torch.tensor(
    [0, 33, 44, 55, 66, 77, 88, nnz_qo], dtype=torch.int32, device="cuda:0"
)
kv_indptr = qo_indptr.clone()

q_at_layer = torch.randn(num_layers, nnz_qo, num_qo_heads, head_dim).half().to("cuda:0")
k_at_layer = torch.randn(num_layers, nnz_kv, num_kv_heads, head_dim).half().to("cuda:0")
v_at_layer = torch.randn(num_layers, nnz_kv, num_kv_heads, head_dim).half().to("cuda:0")

# create auxiliary data structures for batch prefill attention
prefill_wrapper.plan(
    qo_indptr,
    kv_indptr,
    num_qo_heads,
    num_kv_heads,
    head_dim,
    causal=True,
)

outputs = []
for i in range(num_layers):
    q = q_at_layer[i]
    k = k_at_layer[i]
    v = v_at_layer[i]
    # compute batch prefill attention, reuse auxiliary data structures
    o = prefill_wrapper.run(q, k, v)
    outputs.append(o)

print(outputs[0].shape)

# below is another example of creating custom mask for batch prefill attention
mask_arr = []
qo_len = (qo_indptr[1:] - qo_indptr[:-1]).cpu().tolist()
kv_len = (kv_indptr[1:] - kv_indptr[:-1]).cpu().tolist()
for i in range(batch_size):
    mask_i = torch.tril(
        torch.full((qo_len[i], kv_len[i]), True, device="cuda:0"),
        diagonal=(kv_len[i] - qo_len[i]),
    )
    mask_arr.append(mask_i.flatten())

mask = torch.cat(mask_arr, dim=0)
prefill_wrapper.plan(
    qo_indptr,
    kv_indptr,
    num_qo_heads,
    num_kv_heads,
    head_dim,
    custom_mask=mask
)

outputs_custom_mask = []
for i in range(num_layers):
    q = q_at_layer[i]
    k = k_at_layer[i]
    v = v_at_layer[i]
    # compute batch prefill attention, reuse auxiliary data structures
    o_custom = prefill_wrapper.run(q, k, v)
    assert torch.allclose(o_custom, outputs[i], rtol=1e-3, atol=1e-3)

print(outputs_custom_mask[0].shape)
```
```

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

### BlockSparseAttentionWrapper Example

Source: https://docs.flashinfer.ai/api/sparse.html

Demonstrates how to use BlockSparseAttentionWrapper for block sparse attention. It includes setting up the workspace, planning the sparse mask, and running the attention computation, followed by a comparison with a dense implementation.

```python
>>> import torch
>>> import flashinfer
>>> num_qo_heads = 32
>>> num_kv_heads = 8
>>> head_dim = 128
>>> # allocate 128MB workspace buffer
>>> workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
>>> bsr_wrapper = flashinfer.BlockSparseAttentionWrapper(workspace_buffer)
>>> # sparse mask: [[0, 0, 1], [1, 0, 1], [0, 1, 1]]
>>> M = 3
>>> N = 3
>>> indptr = torch.tensor([0, 1, 3, 5], dtype=torch.int32, device="cuda:0")
>>> indices = torch.tensor([2, 0, 2, 1, 2], dtype=torch.int32, device="cuda:0")
>>> bsr_wrapper.plan(
...     indptr,
...     indices,
...     M,
...     N,
...     1, # R(block_rows)=1
...     1, # C(block_columns)=1
...     num_qo_heads,
...     num_kv_heads,
...     head_dim,
... )
>>> q = torch.randn((M, num_qo_heads, head_dim), dtype=torch.float16, device="cuda:0")
>>> k = torch.randn((N, num_kv_heads, head_dim), dtype=torch.float16, device="cuda:0")
>>> v = torch.randn((N, num_kv_heads, head_dim), dtype=torch.float16, device="cuda:0")
>>> o = bsr_wrapper.run(q, k, v)
>>> # use dense implementation with attention mask for comparison
>>> mask = torch.tensor([[0, 0, 1], [1, 0, 1], [0, 1, 1]], dtype=torch.bool, device="cuda:0")
>>> o_ref = flashinfer.single_prefill_with_kv_cache(q, k, v, custom_mask=mask)
>>> torch.allclose(o, o_ref)
True
```

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

### Perform Top-K Selection with FlashInfer

Source: https://docs.flashinfer.ai/generated/flashinfer.top_k.html

This example demonstrates the basic usage of flashinfer.top_k to select the top-k elements from a tensor. Ensure torch and flashinfer are imported and a CUDA-enabled device is available. The default behavior returns unsorted results for performance.

```python
import torch
import flashinfer
torch.manual_seed(42)
batch_size = 4
vocab_size = 32000
k = 256
logits = torch.randn(batch_size, vocab_size, device="cuda")
values, indices = flashinfer.top_k(logits, k)
values.shape, indices.shape
```

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

### Check PyTorch and CUDA Versions

Source: https://docs.flashinfer.ai/_sources/installation.rst.txt

Verify your PyTorch installation and its associated CUDA version before installing FlashInfer from source.

```python
python -c "import torch; print(torch.__version__, torch.version.cuda)"
```

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

### Initialize and Use MoeAlltoAll

Source: https://docs.flashinfer.ai/generated/flashinfer.comm.MoeAlltoAll.html

Instantiate MoeAlltoAll for MoE operations. Use dispatch to send data and combine to receive processed data. Ensure correct parameters like mapping, max_num_tokens, top_k, and num_experts are provided.

```python
>>> moe_a2a = MoeAlltoAll(mapping, max_num_tokens=2048, top_k=2, num_experts=8)
>>> recv = moe_a2a.dispatch(experts, [hidden, ids, scales], batch_size)
>>> output = moe_a2a.combine(processed, batch_size)

```

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

### BatchMLAPagedAttentionWrapper Example

Source: https://docs.flashinfer.ai/api/attention.html

Example demonstrating the usage of BatchMLAPagedAttentionWrapper for MLA PagedAttention, including planning and running the attention mechanism.

```APIDOC
## BatchMLAPagedAttentionWrapper Example

### Description
This example demonstrates how to initialize and use the `BatchMLAPagedAttentionWrapper` for MLA PagedAttention. It covers setting up input tensors, planning the attention operation, and running the attention computation.

### Method
run

### Endpoint
N/A (This is a Python class method)

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None (This is a Python class method)

### Request Example
```python
>>> import torch
>>> import flashinfer
>>> num_local_heads = 128
>>> batch_size = 114
>>> head_dim_ckv = 512
>>> head_dim_kpe = 64
>>> page_size = 1
>>> mla_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(
...     torch.empty(128 * 1024 * 1024, dtype=torch.int8).to(0),
...     backend="fa2"
... )
>>> q_indptr = torch.arange(0, batch_size + 1).to(0).int() # for decode, each query length is 1
>>> kv_lens = torch.full((batch_size,), 999, dtype=torch.int32).to(0)
>>> kv_indptr = torch.arange(0, batch_size + 1).to(0).int() * 999
>>> kv_indices = torch.arange(0, batch_size * 999).to(0).int()
>>> q_nope = torch.randn(
...     batch_size * 1, num_local_heads, head_dim_ckv, dtype=torch.bfloat16, device="cuda"
... )
>>> q_pe = torch.zeros(
...     batch_size * 1, num_local_heads, head_dim_kpe, dtype=torch.bfloat16, device="cuda"
... )
>>> ckv = torch.randn(
...     batch_size * 999, 1, head_dim_ckv, dtype=torch.bfloat16, device="cuda"
... )
>>> kpe = torch.zeros(
...     batch_size * 999, 1, head_dim_kpe, dtype=torch.bfloat16, device="cuda"
... )
>>> sm_scale = 1.0 / ((128 + 64) ** 0.5)  # use head dimension before matrix absorption
>>> mla_wrapper.plan(
...     q_indptr,
...     kv_indptr,
...     kv_indices,
...     kv_lens,
...     num_local_heads,
...     head_dim_ckv,
...     head_dim_kpe,
...     page_size,
...     False,  # causal
...     sm_scale,
...     q_nope.dtype,
...     ckv.dtype,
... )
>>> o = mla_wrapper.run(q_nope, q_pe, ckv, kpe, return_lse=False)
>>> o.shape
torch.Size([114, 128, 512])
```

### Response
#### Success Response (200)
- **o** (_torch.Tensor_) - The output tensor after attention computation.

#### Response Example
```json
{
  "example": "torch.Size([114, 128, 512])"
}
```
```

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

### MultiLevelCascadeAttentionWrapper Example

Source: https://docs.flashinfer.ai/api/cascade.html

Example demonstrating how to initialize and use the MultiLevelCascadeAttentionWrapper for batch decode attention with multiple KV cache levels.

```APIDOC
## Example Usage

```python
>>> import torch
>>> import flashinfer
>>> num_layers = 32
>>> num_qo_heads = 64
>>> num_kv_heads = 8
>>> head_dim = 128
>>> page_size = 16
>>> # allocate 128MB workspace buffer
>>> workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
>>> wrapper = flashinfer.MultiLevelCascadeAttentionWrapper(
...     2, workspace_buffer, "NHD"
... )
>>> batch_size = 7
>>> shared_kv_num_pages = 512
>>> unique_kv_num_pages = 128
>>> total_num_pages = shared_kv_num_pages + unique_kv_num_pages
>>> shared_kv_page_indices = torch.arange(shared_kv_num_pages).int().to("cuda:0")
>>> shared_kv_page_indptr = torch.tensor([0, shared_kv_num_pages], dtype=torch.int32, device="cuda:0")
>>> unique_kv_page_indices = torch.arange(shared_kv_num_pages, total_num_pages).int().to("cuda:0")
>>> unique_kv_page_indptr = torch.tensor(
...     [0, 17, 29, 44, 48, 66, 100, 128], dtype=torch.int32, device="cuda:0"
... )
>>> shared_kv_last_page_len = torch.tensor([page_size], dtype=torch.int32, device="cuda:0")
>>> # 1 <= kv_last_page_len <= page_size
>>> unique_kv_last_page_len = torch.tensor(
...     [1, 7, 14, 4, 3, 1, 16], dtype=torch.int32, device="cuda:0"
... )
>>> kv_cache_at_layer = [
...     torch.randn(
...         total_num_pages, 2, page_size, num_kv_heads, head_dim, dtype=torch.float16, device="cuda:0"
...     ) for _ in range(num_layers)
... ]
>>> qo_indptr_arr = [
...     torch.tensor([0, batch_size], dtype=torch.int32, device="cuda:0"),  # top-level for shared KV-Cache
...     torch.arange(batch_size + 1, dtype=torch.int32, device="cuda:0")    # bottom-level for unique KV-Cache
... ]
>>> # create auxiliary data structures for batch decode attention
>>> wrapper.plan(
...     qo_indptr_arr,
...     [shared_kv_page_indptr, unique_kv_page_indptr],
...     [shared_kv_page_indices, unique_kv_page_indices],
...     [shared_kv_last_page_len, unique_kv_last_page_len],
...     num_qo_heads,
...     num_kv_heads,
...     head_dim,
...     page_size,
... )
>>> outputs = []
>>> for i in range(num_layers):
...     q = torch.randn(batch_size, num_qo_heads, head_dim).half().to("cuda:0")
...     # compute batch decode attention, reuse auxiliary data structures for all layers
...     o = wrapper.run(q, kv_cache_at_layer[i])
...     outputs.append(o)
... 
>>> outputs[0].shape
torch.Size([7, 64, 128])
```
```

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

### Initialize and Use BatchDecodeWithSharedPrefixPagedKVCacheWrapper

Source: https://docs.flashinfer.ai/api/cascade.html

Demonstrates the complete workflow of initializing the wrapper, preparing auxiliary data structures, and performing batch decode attention for multiple layers. Auxiliary data structures can be reused across layers.

```python
>>> import torch
>>> import flashinfer
>>> num_layers = 32
>>> num_qo_heads = 64
>>> num_kv_heads = 8
>>> head_dim = 128
>>> max_num_pages = 128
>>> page_size = 16
>>> # allocate 128MB workspace buffer
>>> workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
>>> wrapper = flashinfer.BatchDecodeWithSharedPrefixPagedKVCacheWrapper(
...     workspace_buffer, "NHD"
... )
>>> batch_size = 7
>>> shared_prefix_len = 8192
>>> unique_kv_page_indices = torch.arange(max_num_pages).int().to("cuda:0")
>>> unique_kv_page_indptr = torch.tensor(
...     [0, 17, 29, 44, 48, 66, 100, 128], dtype=torch.int32, device="cuda:0"
... )
>>> # 1 <= kv_last_page_len <= page_size
>>> unique_kv_last_page_len = torch.tensor(
...     [1, 7, 14, 4, 3, 1, 16], dtype=torch.int32, device="cuda:0"
... )
>>> unique_kv_cache_at_layer = [
...     torch.randn(
...         max_num_pages, 2, page_size, num_kv_heads, head_dim, dtype=torch.float16, device="cuda:0"
...     ) for _ in range(num_layers)
... ]
>>> shared_k_data_at_layer = [
...     torch.randn(
...         shared_prefix_len, num_kv_heads, head_dim, dtype=torch.float16, device="cuda:0"
...     ) for _ in range(num_layers)
... ]
>>> shared_v_data_at_layer = [
...     torch.randn(
...         shared_prefix_len, num_kv_heads, head_dim, dtype=torch.float16, device="cuda:0"
...     ) for _ in range(num_layers)
... ]
>>> # create auxiliary data structures for batch decode attention
>>> wrapper.begin_forward(
...     unique_kv_page_indptr,
...     unique_kv_page_indices,
...     unique_kv_last_page_len,
...     num_qo_heads,
...     num_kv_heads,
...     head_dim,
...     page_size,
...     data_type=torch.float16
... )
>>> outputs = []
>>> for i in range(num_layers):
...     q = torch.randn(batch_size, num_qo_heads, head_dim).half().to("cuda:0")
...     k_shared = shared_k_data_at_layer[i]
...     v_shared = shared_v_data_at_layer[i]
...     unique_kv_cache = unique_kv_cache_at_layer[i]
...     # compute batch decode attention, reuse auxiliary data structures for all layers
...     o = wrapper.forward(q, k_shared, v_shared, unique_kv_cache)
...     outputs.append(o)
... 
>>> outputs[0].shape
torch.Size([7, 64, 128])
```

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

### Install FlashInfer JIT Cache Package

Source: https://docs.flashinfer.ai/installation.html

Install the JIT cache package for FlashInfer. Replace 'cu129' with your specific CUDA version (e.g., cu128, cu129, cu130).

```bash
pip install flashinfer-jit-cache --index-url https://flashinfer.ai/whl/cu129
```

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

### Initialize and Run VariableBlockSparseAttentionWrapper

Source: https://docs.flashinfer.ai/api/sparse.html

Demonstrates the typical workflow of initializing the wrapper, planning the attention parameters, and running the attention computation with query, key, and value tensors.

```python
>>> wrapper = flashinfer.VariableBlockSparseAttentionWrapper(workspace_buffer)
>>> block_mask_map = torch.tensor([[[0, 0, 1], [1, 0, 1], [0, 1, 1]]], dtype=torch.bool, device="cuda:0")
>>> block_row_sz = torch.tensor([[1, 2, 3]], dtype=torch.int32, device="cuda:0")
>>> block_col_sz = torch.tensor([[3, 1, 2]], dtype=torch.int32, device="cuda:0")
>>> wrapper.plan(
...     block_mask_map,
...     block_row_sz,
...     block_col_sz,
...     num_qo_heads,
...     num_kv_heads,
...     head_dim,
... )
>>> q = torch.randn((num_qo_heads, seq_len, head_dim), dtype=torch.float16, device="cuda:0")
>>> k = torch.randn((num_kv_heads, seq_len, head_dim), dtype=torch.float16, device="cuda:0")
>>> v = torch.randn((num_kv_heads, seq_len, head_dim), dtype=torch.float16, device="cuda:0")
>>> o = wrapper.run(q, k, v)
```

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

### Install FlashInfer from Source (Editable Mode)

Source: https://docs.flashinfer.ai/_sources/installation.rst.txt

For development purposes, install FlashInfer in editable mode. This allows changes to the source code to be reflected immediately without reinstallation.

```bash
python -m pip install --no-build-isolation -e . -v
```

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

### flashinfer.comm.trtllm_lamport_initialize_all

Source: https://docs.flashinfer.ai/_sources/generated/flashinfer.comm.trtllm_lamport_initialize_all.rst.txt

Initializes TRTLLM Lamport for all available devices.

```APIDOC
## flashinfer.comm.trtllm_lamport_initialize_all

### Description
Initializes TRTLLM Lamport for all available devices.

### Method
N/A (This is a function call, not an HTTP endpoint)

### Endpoint
N/A

### Parameters
N/A

### Request Example
```python
import flashinfer

flashinfer.comm.trtllm_lamport_initialize_all()
```

### Response
N/A (This function does not return a value, it performs an initialization action.)
```

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

### Initialize and Run MLA PagedAttention Wrapper

Source: https://docs.flashinfer.ai/api/attention.html

Demonstrates initializing the BatchMLAPagedAttentionWrapper and running the attention mechanism with sample tensors. Ensure the workspace buffer is correctly sized and on the same device as input tensors. The `plan` method must be called before `run`.

```python
>>> import torch
>>> import flashinfer
>>> num_local_heads = 128
>>> batch_size = 114
>>> head_dim_ckv = 512
>>> head_dim_kpe = 64
>>> page_size = 1
>>> mla_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(
...     torch.empty(128 * 1024 * 1024, dtype=torch.int8).to(0),
...     backend="fa2"
... )
>>> q_indptr = torch.arange(0, batch_size + 1).to(0).int() # for decode, each query length is 1
>>> kv_lens = torch.full((batch_size,), 999, dtype=torch.int32).to(0)
>>> kv_indptr = torch.arange(0, batch_size + 1).to(0).int() * 999
>>> kv_indices = torch.arange(0, batch_size * 999).to(0).int()
>>> q_nope = torch.randn(
...     batch_size * 1, num_local_heads, head_dim_ckv, dtype=torch.bfloat16, device="cuda"
... )
>>> q_pe = torch.zeros(
...     batch_size * 1, num_local_heads, head_dim_kpe, dtype=torch.bfloat16, device="cuda"
... )
>>> ckv = torch.randn(
...     batch_size * 999, 1, head_dim_ckv, dtype=torch.bfloat16, device="cuda"
... )
>>> kpe = torch.zeros(
...     batch_size * 999, 1, head_dim_kpe, dtype=torch.bfloat16, device="cuda"
... )
>>> sm_scale = 1.0 / ((128 + 64) ** 0.5)  # use head dimension before matrix absorption
>>> mla_wrapper.plan(
...     q_indptr,
...     kv_indptr,
...     kv_indices,
...     kv_lens,
...     num_local_heads,
...     head_dim_ckv,
...     head_dim_kpe,
...     page_size,
...     False,  # causal
...     sm_scale,
...     q_nope.dtype,
...     ckv.dtype,
... )
>>> o = mla_wrapper.run(q_nope, q_pe, ckv, kpe, return_lse=False)
>>> o.shape
torch.Size([114, 128, 512])
```

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

### flashinfer.comm.vllm_init_custom_ar

Source: https://docs.flashinfer.ai/generated/flashinfer.comm.vllm_init_custom_ar.html

Initializes custom attention routing for vLLM. This function is used for distributed inference setups.

```APIDOC
## flashinfer.comm.vllm_init_custom_ar

### Description
Initializes custom attention routing for vLLM. This function is crucial for setting up distributed inference environments where attention mechanisms need custom routing logic.

### Method
Python Function Call

### Endpoint
N/A (Python function)

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **_ipc_tensors** (List[int]) - Required - A list of integers representing IPC tensor identifiers.
- **_rank_data** (Tensor) - Required - A tensor containing rank-specific data.
- **_rank** (int) - Required - The current rank of the process.
- **_full_nvlink** (bool) - Required - A boolean indicating whether full NVLink is utilized.

### Request Example
```python
import flashinfer
import torch

# Assuming ipc_tensors, rank_data, rank, and full_nvlink are defined
# Example placeholder values:
ipc_tensors = [1, 2, 3]
rank_data = torch.randn(10)
rank = 0
full_nvlink = True

return_value = flashinfer.comm.vllm_init_custom_ar(ipc_tensors, rank_data, rank, full_nvlink)
print(f"vllm_init_custom_ar returned: {return_value}")
```

### Response
#### Success Response (int)
- **return_value** (int) - An integer representing the status or result of the initialization.

#### Response Example
```json
{
  "return_value": 0
}
```
```

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

### Create and Execute a Logits Processing Pipeline

Source: https://docs.flashinfer.ai/generated/flashinfer.logits_processor.LogitsPipe.html

Demonstrates creating a pipeline with Temperature, Softmax, TopK, and Sample processors, then executing it with random logits. Runtime parameters like temperature and top_k are passed during the call.

```python
import torch
from flashinfer.logits_processor import LogitsPipe, Temperature, Softmax, TopK, Sample
torch.manual_seed(42)

# Basic pipeline with temperature, top-k, and sampling
pipe = LogitsPipe([
    Temperature(),
    Softmax(),
    TopK(),
    Sample(deterministic=True)
])

# Execute the pipeline
logits = torch.randn(4, 32000, device="cuda")  # [batch_size, vocab_size]
token_ids = pipe(logits, temperature=0.9, top_k=40)
print(token_ids)
```

```text
tensor([15806,  8154, 13923, 20311], device='cuda:0')
```

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

### Initialize and Use TopP Logits Processor

Source: https://docs.flashinfer.ai/generated/flashinfer.logits_processor.TopP.html

Demonstrates initializing a LogitsPipe with the TopP processor and applying it to normalized probabilities. Ensure probabilities are normalized before passing to the processor.

```python
import torch
from flashinfer.logits_processor import LogitsPipe, Softmax, TopP, Sample
torch.manual_seed(42)
pipe = LogitsPipe([TopP()])
probs = torch.randn(2, 2, device="cuda")
probs_normed = probs / probs.sum(dim=-1, keepdim=True)
print(probs_normed)
toapp_probs = pipe(probs_normed, top_p=0.9)
print(toapp_probs)
```

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

### Manage Pre-compiled CUDA Binaries (Cubins)

Source: https://docs.flashinfer.ai/_sources/cli.rst.txt

Handle pre-compiled CUDA binaries by downloading them, listing the downloaded cubins, or clearing the downloaded cubins.

```bash
flashinfer download-cubin
```

```bash
flashinfer list-cubins
```

```bash
flashinfer clear-cubin
```

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

### Clone FlashInfer Repository

Source: https://docs.flashinfer.ai/_sources/installation.rst.txt

Clone the FlashInfer repository from GitHub to install from source. Ensure you use the --recursive flag to include submodules.

```bash
git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
```

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

### Incremental Tuning with Multiple Sessions

Source: https://docs.flashinfer.ai/autotuning.html

Demonstrates accumulating configurations from multiple tuning sessions into a single cache file. Each session uses `flashinfer.autotune` with the same cache path.

```python
# Session 1: tune with batch_size=1
with flashinfer.autotune(True, cache="configs.json"):
    run_model(batch_size=1)

# Session 2: tune with batch_size=32 (configs.json now has both)
with flashinfer.autotune(True, cache="configs.json"):
    run_model(batch_size=32)

```
Tokens:
Raw