FlashInfer Wrapper APIs

This page documents the supported flashinfer compatibility functions. See the FlashInfer wrapper guide for installation, supported workflows, constraints, and the end-to-end FP8 MLA example.

FlashInfer Norm covers RMSNorm, LayerNorm, fused add Norm, and quantized variants.

GEMM

flashinfer.gemm.bmm_bf16(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor | None = None, out_dtype: torch.dtype = torch.bfloat16, backend: Literal['auto', 'mudnn'] = 'auto') → torch.Tensor

Perform FlashInfer-compatible 16-bit BMM on MUSA.

This function computes A @ B for each batch and stores the result in the requested output dtype. FP16 and BF16 inputs are supported.

Parameters:

  • A (torch.Tensor) – Input A with shape (batch, m, k) and FP16 or BF16 dtype.

  • B (torch.Tensor) – Input B with shape (batch, k, n) and the same dtype as A.

  • out (Optional[torch.Tensor]) – Preallocated output with shape (batch, m, n). When provided, the same tensor is returned.

  • out_dtype (torch.dtype) – Output dtype. BF16, FP16, and FP32 are supported. Default is BF16.

  • backend (Literal["auto", "mudnn"]) – Backend selector. "auto" selects muDNN.

Returns:

Output with shape (batch, m, n).

Return type:

torch.Tensor

flashinfer.gemm.bmm_fp8(A: torch.Tensor, B: torch.Tensor, A_scale: torch.Tensor, B_scale: torch.Tensor, dtype: torch.dtype, out: torch.Tensor | None = None, backend: Literal['auto', 'mudnn'] = 'auto') → torch.Tensor

Perform FlashInfer-compatible FP8 BMM with scaling on MUSA.

This function computes A @ B for each batch, applies the FP32 scaling factors for A and B, and stores the result in BF16 or FP16.

Parameters:

  • A (torch.Tensor) – FP8 E4M3 or E5M2 input A with shape (batch, m, k).

  • B (torch.Tensor) – FP8 E4M3 or E5M2 input B with shape (batch, k, n).

  • A_scale (torch.Tensor) – FP32 tensorwise scalar or rank-3 channelwise scaling factors for A.

  • B_scale (torch.Tensor) – FP32 tensorwise scalar or rank-3 channelwise scaling factors for B.

  • dtype (torch.dtype) – Output dtype. BF16 and FP16 are supported.

  • out (Optional[torch.Tensor]) – Preallocated output with shape (batch, m, n). When provided, the same tensor is returned.

  • backend (Literal["auto", "mudnn"]) – Backend selector. "auto" selects muDNN.

Returns:

Output with shape (batch, m, n).

Return type:

torch.Tensor

flashinfer.gemm.gemm_fp8_nt_groupwise(a: torch.Tensor, b: torch.Tensor, a_scale: torch.Tensor, b_scale: torch.Tensor, scale_major_mode: Literal['MN', 'K'] | None = None, mma_sm: int = 1, scale_granularity_mnk: Tuple[int, int, int] = (1, 128, 128), out: torch.Tensor | None = None, out_dtype: torch.dtype | None = None, backend: Literal['auto', 'mudnn', 'mubin'] = 'auto', output_scale: torch.Tensor | None = None) → torch.Tensor

Perform FlashInfer-compatible groupwise FP8 NT GEMM on MUSA.

This function computes the matrix product of two FP8 tensors, applies groupwise scaling, and stores the result in the requested output dtype. It forms batch-size-one views and dispatches through MATE BMM.

Parameters:

  • a (torch.Tensor) – Row-major FP8 input A with shape (m, k).

  • b (torch.Tensor) – Row-major FP8 input B with shape (n, k). The operation computes a @ b.T.

  • a_scale (torch.Tensor) – FP32 scaling factors for A. Shape is (m, k_blocks) for K-major scales or (k_blocks, m) for MN-major scales. It must be contiguous.

  • b_scale (torch.Tensor) – FP32 scaling factors for B. Shape is (n_blocks, k_blocks) for K-major scales or (k_blocks, n_blocks) for MN-major scales. A non-scalar tensor must be contiguous.

  • scale_major_mode (Optional[Literal["MN", "K"]]) – Common scale layout. None defaults to "K".

  • mma_sm (int) – MMA configuration. MUSA currently supports 1. Default is 1.

  • scale_granularity_mnk (Tuple[int, int, int]) – Scale granularity (m, n, k). Default is (1, 128, 128).

  • out (Optional[torch.Tensor]) – Preallocated output with shape (m, n). Its dtype takes precedence over out_dtype and the same tensor is returned.

  • out_dtype (Optional[torch.dtype]) – Output dtype when out is omitted. Without output_scale, BF16 and FP16 are supported and the default is BF16. With output_scale, FP8 E4M3 is required and selected by default.

  • backend (Literal["auto", "mudnn", "mubin"]) – Backend selector. Without output_scale, "auto" selects muDNN. With output_scale, "auto" selects MUBIN.

  • output_scale (Optional[torch.Tensor]) – FP32 scales for FP8 E4M3 output. Providing this MUSA extension selects the MUBIN output path.

Returns:

Output with shape (m, n).

Return type:

torch.Tensor

flashinfer.gemm.group_deepgemm_fp8_nt_groupwise(a: torch.Tensor, b: torch.Tensor, a_scale: torch.Tensor, b_scale: torch.Tensor, m_indices: torch.Tensor, scale_granularity_mnk: Tuple[int, int, int] = (1, 128, 128), out: torch.Tensor | None = None, out_dtype: torch.dtype | None = None) → torch.Tensor

Perform contiguous grouped FP8 NT GEMM.

For each valid row i, this function computes out[i] = a[i] @ b[m_indices[i]].T. Rows for each expert must occupy a contiguous, 128-row-aligned region. Entries with m_indices == -1 are padding and their output values are unspecified.

Parameters:

  • a (torch.Tensor) – K-major E4M3 input with shape (m, k).

  • b (torch.Tensor) – K-major E4M3 expert weights with shape (num_groups, n, k).

  • a_scale (torch.Tensor) – FP32 K-major scales for a with shape (m, k // 128).

  • b_scale (torch.Tensor) – FP32 K-major scales for b with shape (num_groups, n // 128, k // 128).

  • m_indices (torch.Tensor) – Contiguous int32 expert indices with shape (m,). Use -1 for padded rows.

  • scale_granularity_mnk (Tuple[int, int, int]) – Scale granularity (m, n, k). Defaults to (1, 128, 128).

  • out (Optional[torch.Tensor]) – Preallocated BF16 output with shape (m, n). When provided, this exact tensor is returned.

  • out_dtype (Optional[torch.dtype]) – Output dtype used only when out is omitted. The supported and default dtype is BF16. Ignored when out is provided.

Returns:

BF16 output with shape (m, n).

Return type:

torch.Tensor

flashinfer.gemm.batch_deepgemm_fp8_nt_groupwise(a: torch.Tensor, b: torch.Tensor, a_scale: torch.Tensor, b_scale: torch.Tensor, masked_m: torch.Tensor, expected_m: int, scale_granularity_mnk: Tuple[int, int, int] = (1, 128, 128), out: torch.Tensor | None = None, out_dtype: torch.dtype | None = None) → torch.Tensor

Perform masked batched FP8 NT GEMM.

For each expert g, this function computes out[g, :masked_m[g]] = a[g, :masked_m[g]] @ b[g].T. Output rows at or beyond masked_m[g] are unspecified.

Parameters:

  • a (torch.Tensor) – K-major E4M3 input with shape (num_groups, max_m, k).

  • b (torch.Tensor) – K-major E4M3 expert weights with shape (num_groups, n, k).

  • a_scale (torch.Tensor) – FP32 K-major scales for a with shape (num_groups, max_m, k // 128).

  • b_scale (torch.Tensor) – FP32 K-major scales for b with shape (num_groups, n // 128, k // 128).

  • masked_m (torch.Tensor) – Contiguous int32 valid-row counts with shape (num_groups,).

  • expected_m (int) – Host-side expected row-count hint used for kernel selection.

  • scale_granularity_mnk (Tuple[int, int, int]) – Scale granularity (m, n, k). Defaults to (1, 128, 128).

  • out (Optional[torch.Tensor]) – Preallocated BF16 output with shape (num_groups, max_m, n). When provided, this exact tensor is returned.

  • out_dtype (Optional[torch.dtype]) – Output dtype used only when out is omitted. The supported and default dtype is BF16. Ignored when out is provided.

Returns:

BF16 output with shape (num_groups, max_m, n).

Return type:

torch.Tensor

Norm

flashinfer.norm.rmsnorm(input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-06, out: torch.Tensor | None = None, enable_pdl: bool | None = None) → torch.Tensor[source]

Apply FlashInfer-compatible RMSNorm.

flashinfer.norm.gemma_rmsnorm(input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-06, out: torch.Tensor | None = None, enable_pdl: bool | None = None) → torch.Tensor[source]

Apply FlashInfer-compatible Gemma RMSNorm.

flashinfer.norm.rmsnorm_quant(out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor, scale: float | torch.Tensor, eps: float = 1e-06, enable_pdl: bool | None = None) → None[source]

Apply RMSNorm and write the quantized result to out.

flashinfer.norm.fused_add_rmsnorm(input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 1e-06, enable_pdl: bool | None = None) → None[source]

Add input to residual and write RMSNorm back to input.

flashinfer.norm.gemma_fused_add_rmsnorm(input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 1e-06, enable_pdl: bool | None = None) → None[source]

Apply FlashInfer-compatible Gemma fused add RMSNorm.

flashinfer.norm.fused_add_rmsnorm_fp8_block_quant(out: torch.Tensor, block_scale: torch.Tensor, normed_out: torch.Tensor, input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 1e-06, enable_pdl: bool | None = None) → None[source]

Apply fused add RMSNorm and row-major 1x128 FP8 quantization.

flashinfer.norm.fused_add_rmsnorm_quant(out: torch.Tensor, input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, scale: float | torch.Tensor, eps: float = 1e-06, enable_pdl: bool | None = None) → None[source]

Apply fused add RMSNorm and write the quantized result to out.

flashinfer.norm.layernorm(input: torch.Tensor, gemma: torch.Tensor, beta: torch.Tensor, eps: float = 1e-06) → torch.Tensor[source]

Apply FlashInfer-compatible LayerNorm.

flashinfer.norm.layernorm_quant(out: torch.Tensor, input: torch.Tensor, gemma: torch.Tensor, beta: torch.Tensor, scale: float | torch.Tensor, eps: float = 1e-06) → None[source]

Apply LayerNorm and write the quantized result to out.

flashinfer.norm.fused_rmsnorm_silu(input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-06, out: torch.Tensor | None = None, block_scale: torch.Tensor | None = None) → torch.Tensor | tuple[torch.Tensor, torch.Tensor][source]

Apply fused RMSNorm and SiLU activation.

flashinfer.norm.fused_dit_gate_residual_layernorm_gamma_beta(input: torch.Tensor, residual: torch.Tensor, gate: torch.Tensor, gamma: torch.Tensor, beta: torch.Tensor, *, gate_bias: torch.Tensor | None = None, epsilon: float = 1e-06, use_nvfp4: bool = False, use_mxfp8: bool = False, global_scaling_factor: torch.Tensor | None = None, input_global_scaling_factor: torch.Tensor | None = None, residual_out: torch.Tensor | None = None, norm_out: torch.Tensor | None = None, sf_out: torch.Tensor | None = None) → tuple[torch.Tensor, torch.Tensor]
flashinfer.norm.fused_dit_gate_residual_layernorm_scale_shift(input: torch.Tensor, residual: torch.Tensor, gate: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, *, gate_bias: torch.Tensor | None = None, scale_bias: torch.Tensor | None = None, shift_bias: torch.Tensor | None = None, epsilon: float = 1e-06, use_nvfp4: bool = False, use_mxfp8: bool = False, global_scaling_factor: torch.Tensor | None = None, input_global_scaling_factor: torch.Tensor | None = None, residual_out: torch.Tensor | None = None, norm_out: torch.Tensor | None = None, sf_out: torch.Tensor | None = None) → tuple[torch.Tensor, torch.Tensor]
flashinfer.norm.fused_dit_residual_layernorm_scale_shift(input: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, *, residual: torch.Tensor | None = None, scale_bias: torch.Tensor | None = None, shift_bias: torch.Tensor | None = None, epsilon: float = 1e-06, use_nvfp4: bool = False, use_mxfp8: bool = False, global_scaling_factor: torch.Tensor | None = None, input_global_scaling_factor: torch.Tensor | None = None, residual_out: torch.Tensor | None = None, norm_out: torch.Tensor | None = None, sf_out: torch.Tensor | None = None) → tuple[torch.Tensor, torch.Tensor]
flashinfer.norm.fused_qk_rmsnorm_rope(qkv: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, *, ppf: int, pph: int, ppw: int, num_frame_channels: int, num_height_channels: int, num_width_channels: int, num_heads_q: int, num_heads_k: int, num_heads_v: int, head_dim: int, eps: float = 1e-06, base: float = 10000.0, interleave: bool = True, factor: float = 1.0, low: float = 0.0, high: float = 0.0, attention_factor: float = 1.0, is_qk_norm: bool = True, output_fp8: bool = False, output_quant_scale: float = 1.0, v_quant_scale: float = 1.0, q_out: torch.Tensor | None = None, k_out: torch.Tensor | None = None, v_out: torch.Tensor | None = None) → tuple[torch.Tensor, torch.Tensor, torch.Tensor]

FP8 MLA RoPE quantization

flashinfer.rope.mla_rope_quantize_fp8(q_rope: torch.Tensor, k_rope: torch.Tensor, q_nope: torch.Tensor, k_nope: torch.Tensor, cos_sin_cache: torch.Tensor, pos_ids: torch.Tensor, is_neox: bool = True, quantize_dtype: torch.dtype | None = None, quant_scale_q: float = 1.0, quant_scale_kv: float = 1.0, q_rope_out: torch.Tensor | None = None, k_rope_out: torch.Tensor | None = None, q_nope_out: torch.Tensor | None = None, k_nope_out: torch.Tensor | None = None, enable_pdl: bool = False) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

Apply MLA RoPE and per-tensor Q/K quantization to FP8 E4M3.

quant_scale_q and quant_scale_kv are independent scalar multipliers: Q_fp8 = cast(Q * quant_scale_q) and K_fp8 = cast(K * quant_scale_kv). cos_sin_cache may be float32 or bfloat16; the kernel performs the rotary arithmetic in float32. All tensor arguments must be on the same MUSA device.

Parameters:

  • q_rope (torch.Tensor) – BF16 query RoPE values with shape [nnz, num_heads, 64]. The last dimension must be contiguous.

  • k_rope (torch.Tensor) – BF16 key RoPE values with shape [nnz, 64]. The last dimension must be contiguous.

  • q_nope (torch.Tensor) – BF16 query latent values with shape [nnz, num_heads, 512]. The nnz and num_heads dimensions must match q_rope and the last dimension must be contiguous.

  • k_nope (torch.Tensor) – BF16 key latent values with shape [nnz, 512]. Its nnz dimension must match k_rope and the last dimension must be contiguous.

  • cos_sin_cache (torch.Tensor) – Contiguous FP32 or BF16 rotary cache with shape [max_seq_len, 64]. Columns [:32] contain cosine values and columns [32:] contain sine values.

  • pos_ids (torch.Tensor) – Contiguous int32 or int64 positions with shape [nnz]. Every value must be in [0, max_seq_len).

  • is_neox (bool) – Select the rotary layout. True rotates the two contiguous 32-value halves; False rotates adjacent even/odd value pairs.

  • quantize_dtype (Optional[torch.dtype]) – Output quantization dtype. Only torch.float8_e4m3fn is supported. When omitted, it is inferred from the first supplied output buffer and otherwise defaults to torch.float8_e4m3fn.

  • quant_scale_q (float) – Host scalar multiplied into both query components before the FP8 cast.

  • quant_scale_kv (float) – Host scalar multiplied into both key components before the FP8 cast.

  • q_rope_out (Optional[torch.Tensor]) – Optional FP8 output buffer with shape [nnz, num_heads, 64]. Strided leading dimensions are supported; the last dimension must be contiguous. A new buffer is allocated when omitted.

  • k_rope_out (Optional[torch.Tensor]) – Optional FP8 output buffer with shape [nnz, 64]. The last dimension must be contiguous. A new buffer is allocated when omitted.

  • q_nope_out (Optional[torch.Tensor]) – Optional FP8 output buffer with shape [nnz, num_heads, 512]. Strided leading dimensions are supported; the last dimension must be contiguous. A new buffer is allocated when omitted.

  • k_nope_out (Optional[torch.Tensor]) – Optional FP8 output buffer with shape [nnz, 512]. The last dimension must be contiguous. A new buffer is allocated when omitted.

  • enable_pdl (bool) – Unused CUDA PDL compatibility argument. It has no effect on MUSA.

Returns:

(q_rope_out, k_rope_out, q_nope_out, k_nope_out). Supplied output buffers are returned directly; omitted buffers are newly allocated on the corresponding input device.

Return type:

tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

Sparse MLA decode metadata

Prepare metadata once and reuse it only while the query shape, sparse lengths, and top-k remain unchanged.

flashinfer.decode.get_batch_decode_metadata_mla(query: torch.Tensor, seq_lens: torch.Tensor, sparse_mla_top_k: int) → _DecodeMetadata

Prepare reusable sparse MLA scheduler metadata.

Pass the returned object to decode through metadata. Prepare a new object after changing seq_lens or the query configuration.

query must be a MUSA tensor with shape [batch, q_len, num_heads, 576]. seq_lens must be a contiguous int32 tensor with shape [batch] (or [batch, 1]) on the same device. sparse_mla_top_k must be a positive multiple of 64.

The return value is opaque scheduler metadata. Pass it to trtllm_batch_decode_with_kv_cache_mla through metadata. Prepare new metadata after changing the query shape or device, the seq_lens values, or the sparse top-k, and do not share one metadata object between overlapping decode calls.

FP8 sparse MLA decode

flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(query: torch.Tensor, kv_cache: torch.Tensor, workspace_buffer: object, qk_nope_head_dim: int, kv_lora_rank: int, qk_rope_head_dim: int, block_tables: torch.Tensor, seq_lens: torch.Tensor | None, max_seq_len: int, sparse_mla_top_k: int = 0, out: torch.Tensor | None = None, bmm1_scale: float | torch.Tensor = 1.0, bmm2_scale: float | torch.Tensor = 1.0, sinks: list[torch.Tensor] | None = None, skip_softmax_threshold_scale_factor: float | None = None, enable_pdl: bool | None = None, backend: str = 'auto', is_var_seq: bool = True, uses_shared_paged_kv_idx: bool = True, lse: torch.Tensor | None = None, return_lse: bool = False, cute_dsl_impl: str = 'auto', kv_scale_format: str = 'auto', cum_seq_lens_q: torch.Tensor | None = None, max_q_len: int | None = None, multi_ctas_kv_counter_buffer: torch.Tensor | None = None, metadata: _DecodeMetadata | None = None)

Run FP8 DeepSeek V3.2 sparse MLA decode through the MATE backend.

For per-tensor FP8 inputs, the caller folds Q/K descales and the attention scale into bmm1_scale. The V descale and any output scale are folded into bmm2_scale, matching the FlashInfer contract.

Parameters:

  • query (torch.Tensor) – Contiguous FP8 E4M3 query with shape [batch, q_len, heads, 576].

  • kv_cache (torch.Tensor) – Contiguous FP8 E4M3 cache with shape [pages, page_size, 576] or [pages, 1, page_size, 576].

  • workspace_buffer (object) – Unused compatibility placeholder for the FlashInfer API.

  • qk_nope_head_dim (int) – Unused compatibility argument. The kernel uses a fixed 512-dimensional latent vector plus a 64-dimensional RoPE tail.

  • kv_lora_rank (int) – Latent-vector width. Only 512 is supported.

  • qk_rope_head_dim (int) – RoPE width. Only 64 is supported.

  • block_tables (torch.Tensor) – Contiguous int32 physical-token indices with shape [batch, q_len, sparse_mla_top_k].

  • seq_lens (Optional[torch.Tensor]) – Contiguous int32 valid sparse lengths with shape [batch].

  • max_seq_len (int) – Unused compatibility argument; runtime work is bounded by sparse_mla_top_k and seq_lens.

  • sparse_mla_top_k (int) – Sparse index capacity. Must be a positive multiple of 64.

  • out (Optional[torch.Tensor]) – Optional contiguous BF16 output buffer with shape [batch, q_len, heads, 512].

  • bmm1_scale (float | torch.Tensor) – QK/softmax scale. Only a host scalar is supported.

  • bmm2_scale (float | torch.Tensor) – Output scale. Only a host scalar is supported.

  • sinks (Optional[list[torch.Tensor]]) – Unused FlashInfer compatibility argument.

  • skip_softmax_threshold_scale_factor (Optional[float]) – Unused FlashInfer compatibility argument.

  • enable_pdl (Optional[bool]) – Unused FlashInfer CUDA-launch compatibility argument.

  • backend (str) – Unused FlashInfer compatibility argument. MATE always uses its FP8 sparse-MLA kernel regardless of this value.

  • is_var_seq (bool) – Unused compatibility argument. Variable sparse lengths are always read from seq_lens.

  • uses_shared_paged_kv_idx (bool) – Unused FlashInfer compatibility argument. The MATE kernel always treats the supplied sparse indices as shared by all query heads.

  • lse (Optional[torch.Tensor]) – Optional contiguous FP32 LSE buffer with shape [batch * q_len, heads] or [batch, q_len, heads]. The kernel writes directly into this buffer.

  • return_lse (bool) – Return (out, lse) instead of only out.

  • cute_dsl_impl (str) – Unused FlashInfer compatibility argument. MATE has no CuTeDSL decode branch, so this value does not refine or override backend.

  • kv_scale_format (str) – Unused compatibility argument. FP8 scales are supplied through bmm1_scale and bmm2_scale.

  • cum_seq_lens_q (Optional[torch.Tensor]) – Unused FlashInfer ragged-query compatibility argument.

  • max_q_len (Optional[int]) – Unused FlashInfer ragged-query compatibility argument.

  • multi_ctas_kv_counter_buffer (Optional[torch.Tensor]) – Unused FlashInfer multi-CTA compatibility argument.

  • metadata (Optional[_DecodeMetadata]) – Reusable metadata returned by get_batch_decode_metadata_mla(). When omitted, metadata is prepared for this call. Prepare a new object whenever seq_lens values, query shape, or top-k change.

Returns:

The BF16 output tensor, or (out, lse) when return_lse is true.

Return type:

torch.Tensor | tuple[torch.Tensor, torch.Tensor]