The aim of flashinfer_benchmark.py is to provide a single framework for benchmarking any FlashInfer kernel and replace standalone benchmarking scripts.

bench_recurrent_kda_prefill.py --case-set h12 runs the six Kimi-K3 TP8 H12 public-API cases. reduce_kda_h12.py combines successful SM100a and SM103a result files without producing a cross-shape aggregate. The benchmark defaults to the natural device/shape dispatcher and records its resolved module; use --candidate-route nonpersistent for a B200 direct-family route A/B. bench_recurrent_kda_prefill.py --case-set small_bh runs the four fixed-layout small-BH cases through the same cold-L2 CUPTI path. bench_recurrent_kda_prefill.py --case-set production --backend cake runs the complete 29-shape fixed/packed inference portfolio. Its JSON records the logical route, every physical Cake module used (including BT16 prepare plus chain), and the explicit per-shape dry/repeat iteration budget. Large state shapes reduce the sample count to stay within the rotating-state capacity and set timing_iteration_budget.low_sample_count when fewer than ten measured iterations fit. --dry-run-iters and --repeat-iters request fixed iteration counts; they are not duration targets.

This framework provides tools to:

• Benchmark FlashInfer's Attention, GEMM, MOE, Norm, Quantization, Sampling, RoPE, Mamba, GDN, and KDA API performance from different kernel backends such as FlashAttention2/3, cuDNN, cuBLAS, CUTLASS, PrimTS, CuTe-DSL, TensorRT-LLM, and Triton
• Compare performance across different configurations
• Batch performance test multiple test cases

Currently supports testing attention, gemm, fused MOE, normalization, quantization, sampling, RoPE, Mamba, GDN (Gated Delta Net), and KDA APIs:

• Attention:
  • BatchDecodeWithPagedKVCacheWrapper - Decode attention with paged KV cache.
    • Also supports computationally similar cudnn_batch_decode_with_kv_cache and trtllm_batch_decode_with_kv_cache.
    • Speculative decode is supported by setting --s_qo > 1 (subject to backend limitations noted below).
  • BatchPrefillWithPagedKVCacheWrapper - Prefill attention with paged KV cache.
    • Also supports computationally similar cudnn_batch_prefill_with_kv_cache and trtllm_batch_context_with_kv_cache.
  • BatchPrefillWithRaggedKVCacheWrapper - Prefill attention with ragged KV cache.
    • Also supports computationally similar cudnn_batch_prefill_with_kv_cache (cudnn-native) and trtllm_ragged_attention_deepseek.
  • BatchMLAPagedAttentionWrapper - MLA attention proposed in DeepSeek series of models.
    • Also supports computationally similar trtllm_batch_decode_with_kv_cache_mla (trtllm-native) and CuTe DSL MLA decode kernel (cute-dsl, SM100+).
  • trtllm_batch_decode_sparse_mla_dsv4 - DeepSeek-V4 sparse MLA using the public TRTLLM-GEN API on SM100/SM103. Supports varlen prefill-style query lengths, causal SWA and compressed-cache sparse tables, FP8/BF16 inputs, sampled FP32 reference checking, and hot-path Q-tile selector benchmarks.
  • All four wrapper attention routines above accept --backends prims-ts on SM100/SM103. The standalone trtllm_batch_decode_sparse_mla_dsv4 routine supports only trtllm-gen.
• GEMM:
  • gemm_fp8_nt_groupwise - GEMM with FP8 data types using groupwise scaling.
  • group_gemm_fp8_nt_groupwise - Group GEMM with FP8 data types using groupwise scaling.
  • bmm_fp8 - Batched matrix multiplication with FP8 inputs.
  • mm_mxfp8 - Dense MXFP8 matrix multiplication.
  • mm_fp8 - Matrix multiplication with FP8 inputs using the trtllm-gen low-latency GEMM (Blackwell SM10.0+, small-M optimized, pre-shuffled weights).
  • mm_fp4 - Matrix multiplication with NVFP4 inputs.
  • mm_bf16 - Matrix multiplication with BF16 inputs (Blackwell SM10.0+).
  • bmm_bf16 - Batched matrix multiplication with BF16 inputs (Blackwell SM10.0+).
• MOE:
  • trtllm_fp4_block_scale_moe - MOE with FP4 quantized weights and block-wise scaling.
  • trtllm_fp8_block_scale_moe - MOE with FP8 quantized weights and block-wise scaling.
  • trtllm_fp8_per_tensor_scale_moe - MOE with FP8 quantized weights and per-tensor scaling.
  • cutlass_fused_moe - CUTLASS fused MoE (base/fp8/nvfp4 variants with optional TP/EP)
  • unified_moe - Unified MoE API comparison between the CUTLASS and cuTile backends. It supports BF16 and NVFP4 W4A4 with gated SwiGLU, SwiGLU-Step, GeGLU, GeGLU-Tanh, and SiTU or non-gated GELU, ReLU, SiLU, ReLU2, and Identity; filters unsupported backends at runtime; and can autotune each backend independently.
• MOE Communication:
  • moe_a2a_dispatch_combine - MoE All-to-All dispatch + combine benchmark for multi-GPU expert-parallel inference. Requires mpirun for multi-GPU execution. Supports optional quantization (FP8, NVFP4, FP8 block-scale) and real MoE kernel computation.
• AllReduce Communication:
  • allreduce_fusion - AllReduce fusion benchmark for multi-GPU inference. Requires mpirun for multi-GPU execution. Supports TRTLLM and TRTLLM MNNVL backends with multiple fusion patterns (plain allreduce, allreduce + residual + RMSNorm).
• Norm:
  • rmsnorm - Root Mean Square Layer Normalization.
  • fused_add_rmsnorm - Fused residual add + RMSNorm.
  • gemma_rmsnorm - Gemma-style RMSNorm using (weight + 1).
  • gemma_fused_add_rmsnorm - Gemma-style fused residual add + RMSNorm.
  • rmsnorm_quant - RMSNorm with FP8 quantized output.
  • fused_add_rmsnorm_quant - Fused residual add + RMSNorm with FP8 quantized output.
  • rmsnorm_fp4quant - RMSNorm with FP4 quantized output (CuTe-DSL, Blackwell SM10.0+).
  • add_rmsnorm_fp4quant - Fused residual add + RMSNorm with FP4 quantized output (CuTe-DSL, Blackwell SM10.0+).
• Quantization:
  • mxfp8_quantize - Quantize tensor to MxFP8 format (Blackwell SM10.0+).
  • mxfp4_quantize - Quantize tensor to MxFP4 format (Blackwell SM10.0+).
  • nvfp4_quantize - Quantize tensor to NVFP4 format with configurable scale factor layout (Blackwell SM10.0+).
  • nvfp4_batched_quantize - Batched NVFP4 quantization (Blackwell SM10.0+).
• Sampling:
  • softmax - Softmax with optional temperature scaling.
  • sampling_from_probs - Sample token indices from probability distributions.
  • sampling_from_logits - Sample token indices from logits (fused softmax + sampling).
  • top_k_sampling_from_probs - Top-K sampling from probabilities.
  • top_p_sampling_from_probs - Top-P (nucleus) sampling from probabilities.
  • top_k_top_p_sampling_from_probs - Combined Top-K and Top-P sampling from probabilities.
  • top_k_top_p_sampling_from_logits - Combined Top-K and Top-P sampling from logits.
  • min_p_sampling_from_probs - Min-P sampling from probabilities.
  • top_k_renorm_probs - Renormalize probabilities after Top-K filtering.
  • top_p_renorm_probs - Renormalize probabilities after Top-P filtering.
  • top_k_mask_logits - Mask logits outside Top-K values.
  • chain_speculative_sampling - Chain speculative sampling for speculative decoding.
  • top_k - Radix-based Top-K selection.
  • top_k_page_table_transform - Fused Top-K with page table lookup.
  • top_k_ragged_transform - Fused Top-K with ragged index transform.
• RoPE (Rotary Positional Embeddings):
  • apply_rope - Apply RoPE with indptr/offsets.
  • apply_rope_pos_ids - Apply RoPE with position IDs.
  • apply_llama31_rope - Apply Llama 3.1 style RoPE with indptr/offsets.
  • apply_llama31_rope_pos_ids - Apply Llama 3.1 style RoPE with position IDs.
  • apply_rope_with_cos_sin_cache - Apply RoPE with precomputed cos/sin cache.
  • mla_rope_quantize_fp8 - MLA RoPE with FP8 quantization (SM8.9+).
  • rope_quantize_fp8 - RoPE with FP8 quantization (SM8.9+).
  • rope_quantize_fp8_append_paged_kv_cache - RoPE with FP8 quantization and paged KV cache append (SM8.9+).
• Mamba (Selective State Space Models):
  • selective_state_update - Selective state update for Mamba layers (generation phase). Supports both single-token prediction (STP) and multi-token prediction (MTP) via --cache_steps. Backends: flashinfer (CUDA, architecture-specific kernels for base/SM90/SM100+) and triton (reference).
• GDN (Gated Delta Net linear attention, SM90+):
  • gated_delta_rule_decode - Single-token (T=1) gated delta rule decode. --state_layout selects between gated_delta_rule_decode_pretranspose ([B, HV, V, K] state, default) and gated_delta_rule_decode ([B, HV, K, V] state). --state_dtype bfloat16 selects the BF16 state kernels (head_size=128, pretranspose only). Backends: flashinfer (CuTe-DSL) and triton (reference).
  • gated_delta_rule_mtp - Multi-token (T>=2) gated delta rule for speculative-decoding verification, with a state pool + indices. --state_dtype float32 uses gated_delta_rule_mtp; --state_dtype bfloat16 uses the BF16 MTP kernel via gated_delta_rule_decode_pretranspose. Backends: flashinfer, triton.
  • chunk_gated_delta_rule - Chunked GDN prefill over varlen sequences (uniform per-sequence length --s_qo). Backends: flashinfer (SM90 C++ / SM100 CuTe-DSL) and fla (flash-linear-attention Triton baseline, perf-only).
• KDA (SM120a):
  • recurrent_kda_prefill - Ordinary multi-token recurrent KDA prefill with fixed or packed inputs. Backends: flashinfer (automatic variant policy), flashinfer-decomp, flashinfer-fused, and optional external cutekda / flash-kda baselines.

A test case is generally invoked as python3 flashinfer_benchmark.py --routine <routine_name> <flags>.

The unified MoE comparison runs both backends from the same routing, activation, and weight inputs. This example uses the Nemotron-3.5-Lightning MoE shape:

python3 flashinfer_benchmark.py --routine unified_moe --backends cutlass cutile --quant-variant bf16 --num_tokens 128 --hidden_size 2688 --intermediate_size 1856 --num_experts 128 --top_k 6 --activation-type Relu2 --input_dtype bfloat16 --autotune

CUDA graph timing is enabled by default and captures one MoE invocation per graph replay with cold-L2 benchmarking enabled; pass --no_cuda_graph for eager timing. Without --autotune, results are named cutlass and cutile; autotuned results use cutlass_autotune and cutile_autotune.

Representative Qwen3.6 and Nemotron cases are in samples/sample_testlist.txt.

See samples in samples/sample_testlist.txt for various example test flags. Example commands and outputs areas follows

# bmm_fp8
$ python3 flashinfer_benchmark.py --routine bmm_fp8 --batch_size 256 --m 1 --n 1024 --k 7168 --input_dtype fp8_e4m3 --mat2_dtype fp8_e4m3 --out_dtype bfloat16 --backends cudnn cublas cutlass --refcheck -vv --generate_repro_command
[INFO] args = Namespace(routine='bmm_fp8', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag=None, generate_repro_command=True, repro_command='', batch_size=256, m=1, n=1024, k=7168, tile_size=128, group_size=1, scale_major_mode='MN', input_dtype='fp8_e4m3', mat2_dtype='fp8_e4m3', out_dtype='bfloat16', mma_sm=1, backends=['cudnn', 'cublas', 'cutlass'], use_128x4_sf_layout=False, use_nvfp4=False, autotune=False)
[INFO] Running testBmmFp8
[INFO] FlashInfer version: 0.3.1
[VVERBOSE] gpu_name = 'NVIDIA_B200'
[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine bmm_fp8 --batch_size 256 --m 1 --n 1024 --k 7168 --input_dtype fp8_e4m3 --mat2_dtype fp8_e4m3 --out_dtype bfloat16 --backends cudnn cublas cutlass --refcheck -vv --generate_repro_command
[VVERBOSE] input_fp8.shape = torch.Size([256, 1, 7168])
[VVERBOSE] input_fp8.dtype = torch.float8_e4m3fn
[VVERBOSE] mat2_fp8.shape = torch.Size([256, 7168, 1024])
[VVERBOSE] mat2_fp8.dtype = torch.float8_e4m3fn
[VVERBOSE] input_inv_s = tensor(0.0109, device='cuda:0')
[VVERBOSE] input_inv_s.dtype = torch.float32
[VVERBOSE] mat2_inv_s = tensor(0.0135, device='cuda:0')
[VVERBOSE] mat2_inv_s.dtype = torch.float32
[PERF] cudnn          :: median time 0.285 ms; std 0.000 ms; achieved tflops 13.180 TFLOPs/sec; achieved tb_per_sec 0.026 TB/sec
[PERF] cublas         :: median time 0.286 ms; std 0.000 ms; achieved tflops 13.159 TFLOPs/sec; achieved tb_per_sec 0.026 TB/sec
[PERF] cutlass        :: median time 0.266 ms; std 0.001 ms; achieved tflops 14.137 TFLOPs/sec; achieved tb_per_sec 0.028 TB/sec

# non-paged (ragged) prefill
$ python3 flashinfer_benchmark.py --routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 fa3 cutlass cudnn --batch_size 16 --s_qo 1024 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_qk 192 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag "DeepSeek-R1"
[INFO] args = Namespace(routine='BatchPrefillWithRaggedKVCacheWrapper', no_cuda_graph=False, use_cupti=False, refcheck=True, allow_output_mismatch=True, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='DeepSeek-R1', generate_repro_command=True, repro_command='', backends=['fa2', 'fa3', 'cutlass', 'cudnn'], page_size=0, batch_size=16, s_qo=1024, s_kv=1024, num_qo_heads=128, num_kv_heads=128, head_dim_qk=192, head_dim_vo=128, head_dim_ckv=None, head_dim_kpe=None, q_dtype='bfloat16', kv_dtype='bfloat16', causal=True, random_actual_seq_len=True)
[INFO] Running testBatchPrefillWithRaggedKVCacheWrapper
[INFO] FlashInfer version: 0.3.1
[VVERBOSE] gpu_name = 'NVIDIA_B200'
[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 fa3 cutlass cudnn --batch_size 16 --s_qo 1024 --s_kv 1024 --num_qo_heads 128 --num_kv_heads 128 --head_dim_qk 192 --head_dim_vo 128 --random_actual_seq_len -vv --refcheck --causal --q_dtype bfloat16 --kv_dtype bfloat16 --allow_output_mismatch --generate_repro_command --case_tag DeepSeek-R1
[WARNING] fa3 for routine BatchPrefillWithRaggedKVCacheWrapper is not supported on compute capability 10.0. Skipping.
[VVERBOSE] s_qo == s_kv, making actual_seq_lens_kv the same as actual_seq_lens_q
[VERBOSE] Average actual qo seq len: 327
[VERBOSE] Average actual kv seq len: 327
[VVERBOSE] actual_seq_lens_q.flatten() = tensor([103, 436, 861, 271, 107,  72, 701,  21, 615, 122, 467, 215, 331, 459,
         88, 373], dtype=torch.int32)
[VVERBOSE] actual_seq_lens_kv.flatten() = tensor([103, 436, 861, 271, 107,  72, 701,  21, 615, 122, 467, 215, 331, 459,
         88, 373], dtype=torch.int32)
[VVERBOSE] q.shape = torch.Size([5242, 128, 192])
[VVERBOSE] k.shape = torch.Size([5242, 128, 192])
[VVERBOSE] v.shape = torch.Size([5242, 128, 128])
[VVERBOSE] qo_indptr.shape = torch.Size([17])
[VVERBOSE] kv_indptr.shape = torch.Size([17])
[VVERBOSE] scale = 0.07216878364870323
[PERF] fa2            :: median time 0.495 ms; std 0.006 ms; achieved tflops 219.336 TFLOPs/sec; achieved tb_per_sec 1.736 TB/sec
[PERF] cutlass        :: median time 0.530 ms; std 0.002 ms; achieved tflops 204.674 TFLOPs/sec; achieved tb_per_sec 1.620 TB/sec
[PERF] cudnn          :: median time 0.313 ms; std 0.000 ms; achieved tflops 346.715 TFLOPs/sec; achieved tb_per_sec 2.745 TB/sec

# RMSNorm with FP8 quantized output
$ python3 flashinfer_benchmark.py --routine rmsnorm_quant --batch_size 32 --hidden_size 4096 --input_dtype bfloat16 --out_dtype fp8_e4m3 --scale 1.0 --refcheck -vv --generate_repro_command --case_tag "rmsnorm_quant_fp8_e4m3"
[INFO] Running testRmsnormQuant
[INFO] FlashInfer version: 0.6.1
[VVERBOSE] gpu_name = 'NVIDIA_B300_SXM6_AC'
[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine rmsnorm_quant --batch_size 32 --hidden_size 4096 --input_dtype bfloat16 --out_dtype fp8_e4m3 --scale 1.0 --refcheck -vv --generate_repro_command --case_tag rmsnorm_quant_fp8_e4m3
[VVERBOSE] input_tensor.shape = torch.Size([32, 4096])
[VVERBOSE] input_tensor.dtype = torch.bfloat16
[VVERBOSE] weight.shape = torch.Size([4096])
[VVERBOSE] out_tensor.dtype = torch.float8_e4m3fn
[VVERBOSE] scale = 1.0
[PERF] cuda           :: median time 0.003 ms; std 0.000 ms; achieved tflops 0.229 TFLOPs/sec; achieved tb_per_sec 0.140 TB/sec

# MxFP8 Quantization (Blackwell SM10.0+ only)
$ python3 flashinfer_benchmark.py --routine mxfp8_quantize --m 2048 --k 8192 --input_dtype bfloat16 --refcheck -vv --generate_repro_command --case_tag "mxfp8_quantize"
[INFO] args = Namespace(routine='mxfp8_quantize', no_cuda_graph=False, use_cupti=False, use_cuda_events=False, refcheck=True, allow_output_mismatch=False, random_seed=42, verbose=2, output_path=None, num_iters=30, dry_run_iters=5, case_tag='mxfp8_quantize', generate_repro_command=True, repro_command='', m=2048, k=8192, input_dtype='bfloat16', is_sf_swizzled_layout=True, no_sf_swizzled_layout=False, alignment=32, enable_pdl=False, backends=['cuda'], batch_size=None, global_scale=1.0, sf_layout='128x4', do_shuffle=False, sf_vec_size=16)
[INFO] Running testMxfp8Quantize
[INFO] FlashInfer version: 0.6.1
[VVERBOSE] gpu_name = 'NVIDIA_B300_SXM6_AC'
[INFO] To reproduce this test case, run the following command: python3 flashinfer_benchmark.py --routine mxfp8_quantize --m 2048 --k 8192 --input_dtype bfloat16 --refcheck -vv --generate_repro_command --case_tag mxfp8_quantize
[VVERBOSE] input_tensor.shape = torch.Size([2048, 8192])
[VVERBOSE] input_tensor.dtype = torch.bfloat16
[VVERBOSE] is_sf_swizzled_layout = True
[VVERBOSE] alignment = 32
[VVERBOSE] enable_pdl = False
[VVERBOSE] Backend cuda: x_q.shape = torch.Size([2048, 8192]), x_q.dtype = torch.float8_e4m3fn, sf.shape = torch.Size([524288]), sf.dtype = torch.uint8
[VVERBOSE] Round-trip error: 0/16777216 (0.00%) elements differ
[PERF] cuda           :: median time 0.016 ms; std 0.000 ms; achieved tflops 3.118 TFLOPs/sec; achieved tb_per_sec 3.150 TB/sec

Run multiple tests from a file and save results:

python3 flashinfer_benchmark.py --testlist samples/sample_testlist.txt --output_path samples/sample_testlist_output.csv

See samples/sample_testlist.txt for an example stdout output from the above command; samples/sample_testlist_output.csv for csv output from the same run.

The output CSV will contain detailed metrics including:

• Median execution time
• Standard deviation
• TFLOPS/sec
• Memory throughput (TB/sec)
• Input flags
• Reproducer commands if --generate_repro_command is provided

Notes:

• Group parameters (--n_group, --topk_group) are ONLY used with DeepSeekV3 routing method. Using them with other routing methods will cause the error: "Routing kernel with groups implies DeepSeekV3 routing method."
• Different MOE kernel implementations have different top_k constraints. FP8 MOE kernels (both Block Scale and Per-Tensor) have stricter limits than FP4 for non-DeepSeekV3 routing methods.
• FP8 MOE kernels require integer values for group parameters, while FP4 MOE kernels accept optional values.
• CUTLASS fused MoE (cutlass_fused_moe) ignores --routing_method, --n_group, and --topk_group; it computes routing via softmax+top-k internally from the provided logits.

The moe_a2a_dispatch_combine routine benchmarks MoE All-to-All communication for multi-GPU expert-parallel inference. It must be launched with mpirun.

Launch Examples:

# Basic (no quantization)
mpirun -np 8 python benchmarks/flashinfer_benchmark.py \
    --routine moe_a2a_dispatch_combine \
    --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8

# With FP8 quantization
mpirun -np 8 python benchmarks/flashinfer_benchmark.py \
    --routine moe_a2a_dispatch_combine \
    --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8 \
    --quant_dtype fp8

# With NVFP4 quantization and real MoE kernel
mpirun -np 8 python benchmarks/flashinfer_benchmark.py \
    --routine moe_a2a_dispatch_combine \
    --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8 \
    --quant_dtype nvfp4 --real_math --intermediate_size 18432

# With validation and per-phase timing
mpirun -np 8 python benchmarks/flashinfer_benchmark.py \
    --routine moe_a2a_dispatch_combine \
    --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8 \
    --validate --per_phase_timing

# Multi-tenant LoRA: carry per-token adapter ID through dispatch
mpirun -np 8 python benchmarks/flashinfer_benchmark.py \
    --routine moe_a2a_dispatch_combine \
    --num_tokens 2048 --hidden_size 7168 --num_experts 256 --top_k 8 \
    --use_lora --validate

The allreduce_fusion routine benchmarks AllReduce fusion operations for multi-GPU inference. It must be launched with mpirun. Both oneshot and twoshot strategies are benchmarked automatically and reported side by side.

Launch Examples:

# Basic allreduce with auto backend
mpirun -np 8 python benchmarks/flashinfer_benchmark.py \
    --routine allreduce_fusion \
    --num_tokens 64 --hidden_size 4096

# With specific backend
mpirun -np 8 python benchmarks/flashinfer_benchmark.py \
    --routine allreduce_fusion \
    --num_tokens 64 --hidden_size 4096 \
    --ar_backend mnnvl

# AllReduce + Residual + RMSNorm fusion
mpirun -np 8 python benchmarks/flashinfer_benchmark.py \
    --routine allreduce_fusion \
    --num_tokens 64 --hidden_size 4096 \
    --pattern ar_residual_rmsnorm

# With validation
mpirun -np 8 python benchmarks/flashinfer_benchmark.py \
    --routine allreduce_fusion \
    --num_tokens 64 --hidden_size 4096 \
    --validate

Applies to gated_delta_rule_decode, gated_delta_rule_mtp, and chunk_gated_delta_rule (SM90+).

Notes:

• Refcheck compares against the torch reference in tests/gdn/reference_delta_rule.py.
• Prefill pre-L2-normalizes k and calls the kernel with use_qk_l2norm_in_kernel=False so the kernel and reference see identical inputs.

The following table summarizes the support surface of each routine & backend's on various CUDA Compute Capabilities.

Each column represents a compute capability. Backends inside cells represent supported backends. A blank cell means no backend is supported for that routine at that compute capability.

Backend Legend:

• fa2: FlashAttention2
• fa2_tc: FlashAttention2 (with Tensor Cores for BatchDecodeWithPagedKVCacheWrapper)
• fa3: FlashAttention-3
• cublas: cuBLAS
• cudnn: cuDNN (via wrapper API)
• cudnn-native: cuDNN (direct API call)
• cutlass: CUTLASS
• tgv: TGV
• trtllm: TensorRT-LLM
• trtllm-gen: TensorRT-LLM
• trtllm-native: TensorRT-LLM (out-of-wrapper)
• prims-ts: Experimental task-scheduled attention kernels (Blackwell SM100/SM103)
• cuda: FlashInfer CUDA kernels
• cute-dsl: FlashInfer CuTe-DSL kernels (Blackwell SM10.0+)
• cute-dsl-prims: SM120 PRIMS FP8 batch-prefill kernels. Ragged inputs use packed NHD storage, while paged K/V uses HND storage. Paged attention accepts the standard combined cache or separate K/V pools without copying. The backend supports FP32 log2 LSE and requires cutlass.experimental.

SM120 PRIMS examples:

python benchmarks/flashinfer_benchmark.py \
  --routine BatchPrefillWithRaggedKVCacheWrapper \
  --backends cute-dsl-prims --batch_size 16 --s_qo 256 --s_kv 2048 \
  --num_qo_heads 32 --num_kv_heads 8 \
  --head_dim_qk 128 --head_dim_vo 128 \
  --q_dtype fp8_e4m3 --kv_dtype fp8_e4m3 \
  --out_dtype bfloat16 --causal --refcheck -vv

python benchmarks/flashinfer_benchmark.py \
  --routine BatchPrefillWithPagedKVCacheWrapper \
  --backends cute-dsl-prims --page_size 64 \
  --batch_size 16 --s_qo 256 --s_kv 2048 \
  --num_qo_heads 32 --num_kv_heads 8 \
  --head_dim_qk 128 --head_dim_vo 128 \
  --q_dtype fp8_e4m3 --kv_dtype fp8_e4m3 \
  --out_dtype bfloat16 --causal --refcheck -vv

• moe_a2a: MoE All-to-All communication (requires mpirun, Blackwell SM10.0+ with MNNVL)
• allreduce: AllReduce fusion communication (requires mpirun, Blackwell SM10.0+ with MNNVL)
• triton: Triton reference kernels (used for Mamba selective_state_update and GDN decode/MTP)
• fla: flash-linear-attention Triton kernels (GDN prefill baseline)
• flashinfer-decomp / flashinfer-fused: pinned SM120 KDA prefill variants
• cutekda / flash-kda: optional external SM120 KDA prefill baselines