Sparse MLA

Use the FlashInfer wrapper when your framework expects the flashinfer.rope and flashinfer.decode package surface. The native APIs below expose MATE’s fused MLA RoPE quantization and FP8 sparse decode path directly.

RoPE and FP8 quantization

mla_rope_quantize_fp8 applies rotary embedding to the 64-value RoPE tail, quantizes the 512-value latent component and the rotated tail to FP8 E4M3, and supports writing both components into slices of caller-owned 576-byte rows.

mate.sparse_mla_interface.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][source]

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]

FP8 sparse decode

Use mate.sparse_mla_interface.get_batch_decode_metadata_mla to prepare scheduler metadata once while the query shape, sparse lengths, and top-k remain unchanged, then pass it to subsequent decode calls.

mate.sparse_mla_interface.sparse_mla_fp8_decode(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)[source]

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]