GEMM

For framework integrations that already target DeepGEMM style Python APIs, prefer the deep-gemm wrapper first. Use the MATE APIs below when wrapper coverage is not enough.

MoE GEMM

The direct MoE GEMM entrypoints below cover 16-bit, 8-bit, and mixed-dtype W4A8 paths. GemmMixedDType.S4FP8 uses b_quant_recipe=(1, 128). For masked_moe_gemm_mixed_dtype, compatible S4FP8 decode workloads can use backend="mutlass". backend="auto" selects MUTLASS for supported non-overlap inputs, generally when the dispatch M is at most 32. Grouped-A a_quant_recipe=(1, 128) requires MUTLASS. GemmMixedDType.FP4FP8 uses E2M1 weights, E8M0 residual scales, an FP32 epilogue scale, a_quant_recipe=(1, -1), and b_quant_recipe=(1, 32) with the MUBIN backend.

mate.gemm.ragged_m_moe_gemm_8bit(input_a: Tuple[torch.Tensor, torch.Tensor], input_b: Tuple[torch.Tensor, torch.Tensor], ragged_tokens_info: torch.Tensor, out: torch.Tensor, gemm_mode: Literal['per_token', 'psum_expert', 'per_expert'] | None = 'per_token', major_a_mode: Literal['M', 'K'] | None = 'K', major_b_mode: Literal['N', 'K'] | None = 'K', scale_granularity_mnk: Tuple[int, int, int] | None = None, num_mp: int | None = None, alignment_m: int | None = None, backend: Literal['auto', 'mubin', 'mutlass'] | None = 'auto')[source]

Perform 8-bit GEMM operation for MoE (Mixture of Experts) with ragged tensor inputs.

This function computes matrix multiplication between 8-bit quantized tensors for MoE models where different experts may have variable numbers of tokens assigned to them.

Parameters:

  • input_a (Tuple[Tensor, Tensor]) – Tuple containing (fp8_tensor, scale_tensor) for input A. fp8_tensor has shape (total_tokens, hidden_size) and should be of fp8 (e4m3/e5m2) type. scale_tensor has shape (total_tokens, hidden_size // scale_granularity_m) and should be of fp32 type.

  • input_b (Tuple[Tensor, Tensor]) – Tuple containing (fp8_tensor, scale_tensor) for input B. fp8_tensor has shape (num_expert, out_hidden_size, hidden_size) and should be of fp8 (e4m3/e5m2) type. scale_tensor has shape (num_expert, out_hidden_size // scale_granularity_n, hidden_size // scale_granularity_k) and should be of fp32 type.

  • ragged_tokens_info (Tensor) – Metadata tensor whose meaning depends on gemm_mode. For per_token, it has shape (total_tokens,) and stores the expert index for each token, with -1 for unused positions. For psum_expert, it has shape (num_expert,) and stores how many tokens the leading experts have in prefix-sum form. For per_expert, it has shape (num_expert,) and stores the token count for each expert.

  • out (Tensor) – Output tensor with shape (total_tokens, out_hidden_size).

  • major_a_mode (Optional[str]) – Indicating major stride of A. Default to K.

  • major_b_mode (Optional[str]) – Indicating major stride of B. Default to K.

  • gemm_mode (Optional[str],) – Indicating different meaning of ragged_tokens_info.

  • scale_granularity_mnk (Optional[Tuple[int, int, int]]) – Quantization granularity for total_tokens, out_hidden_size, hidden_size (m, n, k) dimensions respectively. Default is (1, 128, 128).

  • alignment_m (Optional[int]) – Alignment requirement for total_tokens (m) dimension. Must be 128 or 256. Default is 128.

  • num_mp (Optional[int]) – Suggest mp number. If None, will be get from device info.

Returns:

Result tensor with shape (total_tokens, out_hidden_size) containing the GEMM output in fp16 or bf16 data type.

Return type:

Tensor

mate.gemm.ragged_m_moe_gemm_16bit(input_a: torch.Tensor, input_b: torch.Tensor, ragged_tokens_info: torch.Tensor, out: torch.Tensor, gemm_mode: Literal['per_token', 'psum_expert', 'per_expert'] | None = 'per_token', major_a_mode: Literal['M', 'K'] | None = 'K', major_b_mode: Literal['N', 'K'] | None = 'K', num_mp: int | None = None, alignment_m: int | None = None, backend: Literal['auto', 'mubin', 'mutlass'] | None = 'auto')[source]

Perform 16-bit GEMM operation for MoE (Mixture of Experts) with ragged tensor inputs.

This function computes matrix multiplication between 16-bit quantized tensors for MoE models where different experts may have variable numbers of tokens assigned to them.

Parameters:

  • input_a (Tensor) – Input tensor A with shape (total_tokens, hidden_size) in fp16/bf16 format.

  • input_b (Tensor) – Input tensor B with shape (num_expert, out_hidden_size, hidden_size) in fp16/bf16 format.

  • ragged_tokens_info (Tensor) –

    If gemm_mode is per_token:

    Tensor indicating which expert each token belongs to, with shape (total_tokens,). Values represent expert indices, with -1 for unused positions.

    If gemm_mode is psum_expert

    Tensor with shape (num_expert, ), indicating how many tokens that first few experts have.

    If gemm_mode is per_expert

    Tensor with shape (num_expert, ), indicating how many tokens that every expert has.

  • out (Tensor) – Output tensor with shape (total_tokens, out_hidden_size).

  • major_a_mode (Optional[str]) – Indicating major stride of A. Default to K.

  • major_b_mode (Optional[str]) – Indicating major stride of B. Default to K.

  • gemm_mode (Optional[str],) – Indicating different meaning of ragged_tokens_info.

  • alignment_m (Optional[int]) – Alignment requirement for total_tokens (m) dimension. Must be 128 or 256. Default is 128.

  • num_mp (Optional[int]) – Suggest mp number. If None, will be get from device info.

Returns:

Result tensor with shape (total_tokens, out_hidden_size) containing the GEMM output in fp16 or bf16 data type.

Return type:

Tensor

mate.gemm.ragged_k_moe_gemm_8bit(input_a: Tuple[torch.Tensor, torch.Tensor], input_b: Tuple[torch.Tensor, torch.Tensor], ragged_tokens_info: torch.Tensor, out: torch.Tensor, gemm_mode: Literal['per_expert'] | None = 'per_expert', major_a_mode: Literal['M', 'K'] | None = 'M', major_b_mode: Literal['N', 'K'] | None = 'N', scale_granularity_mnk: Tuple[int, int, int] | None = None, num_mp: int | None = None)[source]

Perform 8-bit GEMM operation for MoE (Mixture of Experts) with token of each expert.

This function computes matrix multiplication between 8-bit quantized tensors for MoE models where different experts may have variable numbers of tokens.

Parameters:

  • input_a (Tuple[Tensor, Tensor]) – Tuple containing (fp8_tensor, scale_tensor) for input A. fp8_tensor has shape (k, m) and should be of fp8 (e4m3/e5m2) type. scale_tensor has shape (k // scale_granularity_k, m) and should be of fp32 type.

  • input_b (Tuple[Tensor, Tensor]) – Tuple containing (fp8_tensor, scale_tensor) for input B. fp8_tensor has shape (k, n) and should be of fp8 (e4m3/e5m2) type. scale_tensor has shape (k // scale_granularity_k, n) and should be of fp32 type.

  • ragged_tokens_info (Tensor) – Tensor indicating the actual number of tokens for each expert, with shape (num_expert,). Values represent token counts for each expert.

  • out (Tensor) – Output tensor with shape (num_expert, max_tokens, out_hidden_size). Should be of float type. Should not be None.

  • gemm_mode (Optional[str],) – Indicating different meaning of ragged_tokens_info.

  • major_a_mode (Optional[str]) – Major mode of A, defult to M. Only support TN m_grouped_gemm on MP31.

  • major_b_mode (Optional[str]) – Major mode of B, defult to N.

  • scale_granularity_mnk (Optional[Tuple[int, int, int]]) – Quantization granularity for max_tokens, out_hidden_size, hidden_size (m, n, k) dimensions respectively. Kgroupgemm only support 1D1D scale, should be (1, 1, 128).

  • num_mp (Optional[int]) – Suggest mp number. If None, will be get from device info.

Returns:

  • Result tensor with shape (num_experts, total_tokens, out_hidden_size) containing the GEMM output in float data type,

  • Representing D = D + A * B for each expert

mate.gemm.ragged_k_moe_gemm_16bit(input_a: torch.Tensor, input_b: torch.Tensor, ragged_tokens_info: torch.Tensor, out: torch.Tensor, gemm_mode: Literal['per_expert'] | None = 'per_expert', major_a_mode: Literal['M', 'K'] | None = 'M', major_b_mode: Literal['N', 'K'] | None = 'N', num_mp: int | None = None)[source]

Perform 16-bit GEMM operation for MoE (Mixture of Experts) with token of each expert.

This function computes matrix multiplication between 16-bit quantized tensors for MoE models where different experts may have variable numbers of tokens.

Parameters:

  • input_a (Tensor) – Input tensor A with shape (sum(ks), m) in FP16 or BF16 format.

  • input_b (Tensor) – Input tensor B with shape (sum(ks), n) and the same dtype as A.

  • ragged_tokens_info (Tensor) – Per-expert K lengths with shape (num_expert,).

  • out (Tensor) – Output tensor with shape (num_expert, m, n). FP32 is supported for FP16/BF16 inputs; BF16 output requires BF16 inputs.

  • gemm_mode (Optional[str],) – Indicating different meaning of ragged_tokens_info.

  • major_a_mode (Optional[str]) – Major mode of A, defult to M. Only support TN m_grouped_gemm on MP31.

  • major_b_mode (Optional[str]) – Major mode of B, defult to N.

  • num_mp (Optional[int]) – Suggest mp number. If None, will be get from device info.

Returns:

  • Result tensor with shape (num_expert, m, n) containing the GEMM output in FP32 or BF16,

  • Representing D = D + A * B for each expert

mate.gemm.masked_moe_gemm_8bit(input_a: Tuple[torch.Tensor, torch.Tensor], input_b: Tuple[torch.Tensor, torch.Tensor], masked_tokens_info: torch.Tensor, out: torch.Tensor, scale_granularity_mnk: Tuple[int, int, int] | None = None, expect_tokens: int | None = None, enable_overlap: bool = False, signal: torch.Tensor | None = None, backend: Literal['auto', 'mubin', 'mutlass'] | None = 'auto')[source]

Perform 8-bit GEMM operation for MoE (Mixture of Experts) with masked tensor inputs.

This function computes matrix multiplication between 8-bit quantized tensors for MoE models where different experts may have variable numbers of tokens, using a mask to indicate the actual number of tokens per expert.

Parameters:

  • input_a (Tuple[Tensor, Tensor]) – Tuple containing (fp8_tensor, scale_tensor) for input A. fp8_tensor has shape (num_expert, max_tokens, hidden_size) and should be of fp8 (e4m3/e5m2) type. scale_tensor has shape (num_expert, max_tokens, hidden_size // scale_granularity_k) and should be of fp32 type.

  • input_b (Tuple[Tensor, Tensor]) – Tuple containing (fp8_tensor, scale_tensor) for input B. fp8_tensor has shape (num_expert, out_hidden_size, hidden_size) and should be of fp8 (e4m3/e5m2) type. scale_tensor has shape (num_expert, out_hidden_size // scale_granularity_n, hidden_size // scale_granularity_k) and should be of fp32 type.

  • masked_tokens_info (Tensor) – Tensor indicating the actual number of tokens for each expert, with shape (num_expert,). Values represent token counts for each expert.

  • out (Tensor) – Output tensor with shape (num_expert, max_tokens, out_hidden_size). Should be of fp16 or bf16 type. If None, a new tensor will be created.

  • scale_granularity_mnk (Optional[Tuple[int, int, int]]) – Quantization granularity for max_tokens, out_hidden_size, hidden_size (m, n, k) dimensions respectively. Default is (1, 128, 128).

  • expect_tokens (Optional[int]) – Expected number of tokens. If None, defaults to 0.

  • enable_overlap (Optional[bool]) – Whether to enable Single-Batch Overlap (SBO). Default is False.

  • signal (Optional[Tensor]) – Signal tensor with shape (num_expert * ceil_div(max_m, 64)) for SBO. Required if enable_overlap is True. If None, a new tensor is created when needed.

Returns:

If enable_overlap is False, returns result tensor with shape (num_expert, max_tokens, out_hidden_size). If enable_overlap is True, returns a tuple containing:

  • result tensor with shape (num_expert, max_tokens, out_hidden_size)

  • signal tensor

  • block_m int

  • threshold int

Return type:

Union[Tensor, Tuple[Tensor, Tensor, int, int]]

mate.gemm.masked_moe_gemm_16bit(a: torch.Tensor, b: torch.Tensor, masked_tokens_info: torch.Tensor, out: torch.Tensor, expect_tokens: int | None = None, enable_overlap: bool = False, signal: torch.Tensor | None = None, backend: Literal['auto', 'mubin', 'mutlass'] | None = 'auto')[source]

Perform 16-bit GEMM operation for MoE (Mixture of Experts) with masked tensor inputs.

This function computes matrix multiplication between 16-bit quantized tensors for MoE models where different experts may have variable numbers of tokens, using a mask to indicate the actual number of tokens per expert.

Parameters:

  • a (Tensor) – Input tensor A with shape (num_expert, max_tokens, hidden_size) in fp16/bf16 format.

  • b (Tensor) – Input tensor B with shape (num_expert, out_hidden_size, hidden_size) in fp16/bf16 format.

  • masked_tokens_info (Tensor) – Tensor indicating the actual number of tokens for each expert, with shape (num_expert,). Values represent token counts for each expert.

  • out (Tensor) – Output tensor with shape (num_expert, max_tokens, out_hidden_size). Should be of fp16 or bf16 type. If None, a new tensor will be created.

  • expect_tokens (Optional[int]) – Expected number of tokens. If None, defaults to 0.

  • enable_overlap (Optional[bool]) – Whether to enable Single-Batch Overlap (SBO). Default is False.

  • signal (Optional[Tensor]) – Signal tensor with shape (num_expert * ceil_div(max_m, 64)) for SBO. Required if enable_overlap is True. If None, a new tensor will be created if needed.

Returns:

If enable_overlap is False, returns result tensor with shape (num_expert, max_tokens, out_hidden_size). If enable_overlap is True, returns a tuple containing:

  • result tensor with shape (num_expert, max_tokens, out_hidden_size)

  • signal tensor

  • block_m int

  • threshold int

Return type:

Union[Tensor, Tuple[Tensor, Tensor, int, int]]

mate.gemm.ragged_moe_gemm_mixed_dtype(input_a: Tuple[torch.Tensor, torch.Tensor], input_b: Tuple[torch.Tensor, torch.Tensor | Tuple[torch.Tensor, torch.Tensor]], ragged_tokens_info: torch.Tensor, out: torch.Tensor, alignment_m: int | None = None, *, mixed_dtype: GemmMixedDType | str, backend: Literal['auto', 'mubin'] | None = 'auto', a_quant_recipe: Tuple[int, int], b_quant_recipe: Tuple[int, int])[source]

Perform mixed-dtype GEMM operation for MoE (Mixture of Experts) with ragged tensor inputs.

This function computes matrix multiplication between mixed-dtype tensors for MoE models where different experts may have variable numbers of tokens assigned to them. GemmMixedDType.S4FP8 selects signed int4 weights, while GemmMixedDType.FP4FP8 selects E2M1 FP4 weights. Both use FP8 activations for input A. The quantization recipes describe quantization block sizes for input A and input B separately, and must be provided explicitly.

Parameters:

  • input_a (Tuple[Tensor, Tensor]) – Tuple containing (activation_tensor, scale_tensor) for input A. activation_tensor has shape (total_tokens, hidden_size). Its dtype is selected by mixed_dtype. S4FP8 accepts E4M3 or E5M2; FP4FP8 requires E4M3. scale_tensor shape and dtype are selected by a_quant_recipe. For a_quant_recipe=(1, -1), it has shape (total_tokens, 1).

  • input_b (Tuple[Tensor, Union[Tensor, Tuple[Tensor, Tensor]]]) – For S4FP8, scales is one BF16 tensor with shape (num_expert, out_hidden_size, ceil_div(hidden_size, 128)). For FP4FP8, scales is (residual_e8m0, epilogue_fp32) with shapes (num_expert, out_hidden_size, ceil_div(hidden_size, 32)) and (num_expert, out_hidden_size).

  • ragged_tokens_info (Tensor) – Tensor indicating which expert each token belongs to, with shape (total_tokens,). Values represent expert indices, with -1 for unused positions.

  • out (Tensor) – Output tensor with shape (total_tokens, out_hidden_size). S4FP8 supports FP16 or BF16; FP4FP8 requires BF16.

  • alignment_m (Optional[int]) – S4FP8 accepts 128 or 256 and defaults to 128. FP4FP8 requires 256 and defaults to 256.

  • mixed_dtype (GemmMixedDType or str) – Mixed dtype selector for input A and input B. Must be provided explicitly. GemmMixedDType.S4FP8 and "s4fp8" mean signed int4 weights for input B and fp8 activations for input A. GemmMixedDType.FP4FP8 and "fp4fp8" mean E2M1 FP4 weights and E4M3 activations.

  • backend (Optional[str]) – Backend selector. Only "auto" and "mubin" are supported.

  • a_quant_recipe (Tuple[int, int]) – Quantization block-size recipe for input A. The tuple is interpreted as (m, k). -1 means the corresponding axis is not split into smaller quantization blocks. Currently, only (1, -1) is supported.

  • b_quant_recipe (Tuple[int, int]) – Quantization block-size recipe for input B. The tuple is interpreted as (n, k). -1 means the corresponding axis is not split into smaller quantization blocks. S4FP8 uses (1, 128) and FP4FP8 uses (1, 32).

Returns:

Result tensor with shape (total_tokens, out_hidden_size) containing the GEMM output in fp16 or bf16 data type.

Return type:

Tensor

mate.gemm.masked_moe_gemm_mixed_dtype(input_a: Tuple[torch.Tensor, torch.Tensor], input_b: Tuple[torch.Tensor, torch.Tensor | Tuple[torch.Tensor, torch.Tensor]], masked_tokens_info: torch.Tensor, out: torch.Tensor, expect_tokens: int | None = None, enable_overlap: bool = False, signal: torch.Tensor | None = None, *, mixed_dtype: GemmMixedDType | str, backend: Literal['auto', 'mubin', 'mutlass'] | None = 'auto', a_quant_recipe: Tuple[int, int], b_quant_recipe: Tuple[int, int])[source]

Perform mixed-dtype GEMM operation for MoE (Mixture of Experts) with masked tensor inputs.

This function computes matrix multiplication between mixed-dtype tensors for MoE models where different experts may have variable numbers of tokens, using a mask to indicate the actual number of tokens per expert. GemmMixedDType.S4FP8 selects signed int4 weights, while GemmMixedDType.FP4FP8 selects E2M1 FP4 weights. The quantization recipes describe quantization block sizes for input A and input B separately, and must be provided explicitly.

Parameters:

  • input_a (Tuple[Tensor, Tensor]) – Tuple containing (activation_tensor, scale_tensor) for input A. activation_tensor has shape (num_expert, max_tokens, hidden_size). Its dtype is selected by mixed_dtype. S4FP8 accepts E4M3 or E5M2; FP4FP8 requires E4M3. scale_tensor shape and dtype are selected by a_quant_recipe. For a_quant_recipe=(1, -1), it has shape (num_expert, max_tokens, 1). For a_quant_recipe=(1, 128), it has shape (num_expert, max_tokens, ceil_div(hidden_size, 128)).

  • input_b (Tuple[Tensor, Union[Tensor, Tuple[Tensor, Tensor]]]) – For S4FP8, scales is one BF16 tensor with shape (num_expert, out_hidden_size, ceil_div(hidden_size, 128)). For FP4FP8, scales is (residual_e8m0, epilogue_fp32) with shapes (num_expert, out_hidden_size, ceil_div(hidden_size, 32)) and (num_expert, out_hidden_size).

  • masked_tokens_info (Tensor) – Tensor indicating the actual number of tokens for each expert, with shape (num_expert,). Values represent token counts for each expert.

  • out (Tensor) – Output tensor with shape (num_expert, max_tokens, out_hidden_size). S4FP8 supports FP16 or BF16; FP4FP8 requires BF16.

  • expect_tokens (Optional[int]) – Expected typical number of tokens per expert. A positive value participates in automatic backend selection. If None or 0, the tensor capacity is used.

  • enable_overlap (Optional[bool]) – Whether to enable Single-Batch Overlap (SBO). Default is False.

  • signal (Optional[Tensor]) – Signal tensor with shape (num_expert * ceil_div(max_m, 64)) for SBO. Required if enable_overlap is True. If None, a new tensor is created when needed.

  • mixed_dtype (GemmMixedDType or str) – Mixed dtype selector for input A and input B. Must be provided explicitly. GemmMixedDType.S4FP8 and "s4fp8" mean signed int4 weights for input B and fp8 activations for input A. GemmMixedDType.FP4FP8 and "fp4fp8" mean E2M1 FP4 weights and E4M3 activations.

  • backend (Optional[str]) – Backend selector. "mutlass" uses the JIT/AOT MP31 kernel. "auto" selects it for compatible non-overlap inputs with min(max_tokens, expect_tokens) <= 32 when expect_tokens is positive, and otherwise uses max_tokens as the threshold input. Grouped-A a_quant_recipe=(1, 128) always selects MUTLASS when its tensor contract is compatible because MUBIN grouped-A support is not available.

  • a_quant_recipe (Tuple[int, int]) – Quantization block-size recipe for input A. The tuple is interpreted as (m, k). -1 means the corresponding axis is not split into smaller quantization blocks. (1, -1) and (1, 128) are supported by the MUTLASS backend for GemmMixedDType.S4FP8. MUBIN supports only (1, -1). Grouped Scale-A may be contiguous in either its M or K-block dimension. FP4FP8 supports only (1, -1).

  • b_quant_recipe (Tuple[int, int]) – Quantization block-size recipe for input B. The tuple is interpreted as (n, k). -1 means the corresponding axis is not split into smaller quantization blocks. S4FP8 uses (1, 128) and FP4FP8 uses (1, 32).

Returns:

If enable_overlap is False, returns result tensor with shape (num_expert, max_tokens, out_hidden_size). If enable_overlap is True, returns a tuple containing:

  • result tensor with shape (num_expert, max_tokens, out_hidden_size)

  • signal tensor

  • block_m int

  • threshold int

Return type:

Union[Tensor, Tuple[Tensor, Tensor, int, int]]

Dense GEMM

mate.gemm.bmm(a: torch.Tensor, b: torch.Tensor, out: torch.Tensor | None = None, *, trans_a: bool = False, trans_b: bool = True, scale_a: torch.Tensor | None = None, scale_b: torch.Tensor | None = None, scale_out: torch.Tensor | None = None, recipe_a: Tuple[int, int] | None = None, recipe_b: Tuple[int, int] | None = None, c: torch.Tensor | None = None, out_dtype: torch.dtype | None = None, fixed_scale_layout: bool | None = None, backend: str = 'auto') → torch.Tensor[source]

Perform FP8, FP16, or BF16 batched matrix multiplication.

This function computes the batched matrix product of A and B, optionally adds C, and stores the result in the requested output dtype. It supports unscaled FP16/BF16 inputs and FP8 inputs with tensorwise, channelwise, or groupwise scaling.

Parameters:

  • a (torch.Tensor) – Input A. Its physical shape is (batch, m, k) when trans_a=False or (batch, k, m) when trans_a=True. Supported dtypes are FP16, BF16, FP8 E4M3, and FP8 E5M2. The final physical dimension must have stride 1.

  • b (torch.Tensor) – Input B. Its physical shape is (batch, n, k) when trans_b=True or (batch, k, n) when trans_b=False. It must use the same 16-bit dtype as A for unscaled BMM, or an FP8 dtype for scaled BMM. The final physical dimension must have stride 1.

  • out (Optional[torch.Tensor]) – Preallocated output D with shape (batch, m, n). When omitted, a tensor is allocated using out_dtype. The default dtype is the input dtype for 16-bit BMM and BF16 for FP8 BMM without scale_out.

  • trans_a (bool) – Whether to transpose the final two dimensions of physical A before multiplication. Default is False.

  • trans_b (bool) – Whether to transpose the final two dimensions of physical B before multiplication. Default is True.

  • scale_a (Optional[torch.Tensor]) – FP32 scaling factors for A. Required for FP8 inputs and ignored for 16-bit inputs. Its logical granularity is specified by recipe_a.

  • scale_b (Optional[torch.Tensor]) – FP32 scaling factors for B. Required for FP8 inputs and ignored for 16-bit inputs. Its logical granularity is specified by recipe_b.

  • scale_out (Optional[torch.Tensor]) – FP32 output scales for FP8 E4M3 output. Providing this tensor selects the MUBIN backend; its shape is (batch, m, ceil(n / 128)). Ignored for 16-bit inputs.

  • recipe_a (Optional[Tuple[int, int]]) – Required FP8 quantization recipe (m_granularity, k_granularity) for A. (-1, -1), (1, -1), and (1, 128) represent tensorwise, channelwise, and K-grouped scaling, respectively.

  • recipe_b (Optional[Tuple[int, int]]) – Required FP8 quantization recipe (n_granularity, k_granularity) for B. In addition to tensorwise, channelwise, and grouped scaling, (128, 128) represents block scaling. Its K granularity must match recipe_a.

  • c (Optional[torch.Tensor]) – Optional accumulation tensor with shape (batch, m, n). It must match the output dtype. FP8 BMM with C requires FP32 output. The MUBIN and MUTLASS backends do not support C.

  • out_dtype (Optional[torch.dtype]) – Output dtype used only when out is omitted. The muDNN backend accepts the 16-bit input dtype or FP32 for unscaled BMM, and FP16, BF16, or FP32 for FP8 BMM. MUBIN requires FP8 E4M3 output. MUTLASS requires BF16 output.

  • fixed_scale_layout (Optional[bool]) – Common packed layout for non-scalar FP8 scales. False selects K-major, True selects MN-major, and None selects K-major for NT or MN-major for NN, TN, and TT. Ignored for 16-bit inputs.

  • backend (str) – Backend selector. "auto" uses muDNN unless scale_out selects MUBIN. Explicitly supported backends are "mudnn", "mubin", and "mutlass". MUTLASS supports NT BF16 or group/block FP8 E4M3 BMM with BF16 output.

Returns:

Output D with shape (batch, m, n). If out is provided, the same tensor is returned.

Return type:

torch.Tensor

DeepGemm Lighting Indexer

mate.deep_gemm.fp8_einsum(expr: str, a: Tuple[torch.Tensor, torch.Tensor], b: Tuple[torch.Tensor, torch.Tensor], d: torch.Tensor, c: torch.Tensor | None = None, recipe: Tuple[int, int, int] = (1, 128, 128)) → None[source]

DeepGEMM-compatible FP8 einsum.

Supported expressions are "bhr,hdr->bhd", "bhd,hdr->bhr", and "bhd,bhr->hdr". The quantization recipe must be (1, 128, 128) or (1, 1, 128).

mate.deep_gemm.tf32_hc_prenorm_gemm(a: torch.Tensor, b: torch.Tensor, d: torch.Tensor, sqr_sum: torch.Tensor, num_splits: int | None = None) → None[source]

TF32 HyperConnection prenorm GEMM.

Parameters:

  • a (torch.Tensor) – Input tensor with shape (M, K) and dtype torch.bfloat16.

  • b (torch.Tensor) – Weight tensor with shape (N, K) and dtype torch.float32.

  • d (torch.Tensor) – Output GEMM tensor with dtype torch.float32. Shape is (M, N) when num_splits is None or <= 1; otherwise (num_splits, M, N).

  • sqr_sum (torch.Tensor) – Output row-wise squared-sum tensor with dtype torch.float32. Shape is (M,) when num_splits is None or <= 1; otherwise (num_splits, M).

  • num_splits (int, default=None) – Optional split-K factor. When greater than 1, the kernel writes per-split partial outputs and callers should reduce them along dim 0.

Return type:

None

mate.deep_gemm.fp8_mqa_logits(q: torch.Tensor, kv: tuple[torch.Tensor, torch.Tensor], weights: torch.Tensor, cu_seq_len_k_start: torch.Tensor, cu_seq_len_k_end: torch.Tensor, clean_logits: bool = False, max_seqlen_k: int = 0) → torch.Tensor[source]

FP8 MQA logits.

This operator computes MQA (multi-query attention) logits for a query sequence against a non-paged KV tensor. It supports both full logits and “compressed logits” mode (when max_seqlen_k > 0), where the output width is limited to a window size.

Parameters:

  • q (torch.Tensor) – FP8 query tensor with shape (seq_len, heads, head_dim) and dtype torch.float8_e4m3fn.

  • kv (tuple[torch.Tensor, torch.Tensor]) –

    A tuple (kv_fp8, kv_scale) containing:

    • kv_fp8: FP8 KV tensor with shape (seq_len_kv, head_dim) and dtype torch.float8_e4m3fn. (MQA uses a single KV head.)

    • kv_scale: FP32 scale tensor with shape (seq_len_kv,) and dtype torch.float32.

  • weights (torch.Tensor) – FP32 weight tensor with shape (seq_len, heads) and dtype torch.float32.

  • cu_seq_len_k_start (torch.Tensor) – Per-row valid KV start offsets (inclusive) for each query row, with shape (seq_len,) and dtype torch.int32.

  • cu_seq_len_k_end (torch.Tensor) – Per-row valid KV end offsets (exclusive) for each query row, with shape (seq_len,) and dtype torch.int32.

  • clean_logits (bool, default=False) – Whether to clean logits outside valid KV range. Must be False when max_seqlen_k > 0.

  • max_seqlen_k (int, default=0) – If > 0, enables compressed logits mode. The output width becomes max_seqlen_k (a windowed logits range per row). In this mode, clean_logits must be False.

Returns:

FP32 logits tensor with dtype torch.float32 and shape:

  • (seq_len, seq_len_kv) if max_seqlen_k == 0

  • (seq_len, max_seqlen_k) if max_seqlen_k > 0

Return type:

torch.Tensor

mate.deep_gemm.fp8_gemm_nt_skip_head_mid(a: Tuple[torch.Tensor, torch.Tensor], b: Tuple[torch.Tensor, torch.Tensor], d: torch.Tensor, head_splits: Tuple[int, int, int], recipe: Tuple[int, int, int] | None = None, compiled_dims: str = 'nk', disable_ue8m0_cast: bool = True)[source]
mate.deep_gemm.get_paged_mqa_logits_metadata(context_lens: torch.Tensor, block_kv: int, num_mps: int = 0) → torch.Tensor[source]

Get metadata for paged MQA logits

Parameters:

  • context_lens (Tensor) – Context lengths of each query, shape (batch_size)

  • block_kv (Tensor) – Block size of kv cache, must be 64 now.

  • num_mps (int) – Number of MP to execute. 0 means use all MPs of the current device

Returns:

Schedule metadata, shape (num_mps + 1, 2)

Return type:

Tensor

mate.deep_gemm.fp8_paged_mqa_logits(q: torch.Tensor, fused_kv_cache: torch.Tensor, weights: torch.Tensor, context_lens: torch.Tensor, block_table: torch.Tensor, schedule_meta: torch.Tensor, max_context_len: int, clean_logits: bool) → torch.Tensor[source]

FP8 Paged MQA logits

Parameters:

  • q (Tensor) – The FP8 query tensor with shape (batch_size, next_n, heads, index_dim)

  • fused_kv_cache (Tensor) – The FP8 kv cache with fp32 scale, shape (num_blocks, block_size, 1, index_dim + 4)

  • weights (Tensor) – The FP32 weight tensor for each query, shape (batch_size * next_n, heads)

  • context_lens (Tensor) –

    Context lengths tensor, supports two layouts:

    • 1D (batch_size,) — all next_n draft tokens of request i share the same context length context_lens[i]. The visible KV range for draft token j is implicitly [0, context_lens[i] - next_n + j].

    • 2D (batch_size, next_n) — each draft token has an independent context length context_lens[i, j], with visible KV range [0, context_lens[i, j] - 1]. Useful for tree-based speculative decoding (e.g. Medusa / EAGLE) where tokens on different branches see different KV prefixes.

    The shape is auto-detected; get_paged_mqa_logits_metadata must be called with the same context_lens tensor.

  • block_table (Tensor) – Block table tensor with shape (batch_size, max_blocks)

  • schedule_meta (Tensor) – Schedule metadata tensor with shape (num_mps + 1, 2), produced by get_paged_mqa_logits_metadata()

  • max_context_len (int) – Maximum context length

  • clean_logits (bool) – Whether to zero-fill logit positions that are out of the valid KV range

Returns:

FP32 logits, shape (batch_size * next_n, max_context_len)

Return type:

Tensor