Run a FlashInfer example

2026-10-30

1 min read time

Applies to Linux

AITER supports single and batch prefill, along with other operations listed in the kernel support matrix. The auto backend calls the most optimized kernels from AITER or the in-tree HIP kernels, based on the benchmarking results. To select AITER explicitly for an operation that exposes a backend argument, pass backend="aiter".

1. Save the following code snippet to a Python script named flashinfer_example.py.

   import torch
   import flashinfer
   
   # Configuration
   seq_len = 1024        # Prompt length
   num_qo_heads = 32     # Number of query/output heads
   num_kv_heads = 8      # Number of KV heads (GQA with 4:1 ratio)
   head_dim = 128
   
   # Create Q, K, V tensors (NHD layout: sequence, heads, dimension)
   q = torch.randn(seq_len, num_qo_heads, head_dim, dtype=torch.float16, device="cuda")
   k = torch.randn(seq_len, num_kv_heads, head_dim, dtype=torch.float16, device="cuda")
   v = torch.randn(seq_len, num_kv_heads, head_dim, dtype=torch.float16, device="cuda")
   
   # Run single prefill attention with causal masking
   output = flashinfer.single_prefill_with_kv_cache(q, k, v, causal=True, backend="aiter")
   

2. Run the script to use FlashInfer.

   python flashinfer_example.py