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 both the existing 8-bit paths and the 0.2.4 mixed-dtype W4A8 path. The current mixed-dtype support is mixed_dtype="s4fp8" / GemmMixedDType.S4FP8 with backend="mubin", a_quant_recipe=(1, -1), and b_quant_recipe=(1, 128).

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'] | 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).

  • input_b (Tuple[Tensor, Union[Tensor, Tuple[Tensor, Tensor]]]) – For S4FP8, scales is one BF16 tensor. 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 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.

  • 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:

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_fp16(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype, out: torch.Tensor | None = None, backend: str = 'auto', c: torch.Tensor | None = None)[source]
mate.gemm.bmm_fp8(a: torch.Tensor, b: torch.Tensor, a_scale: torch.Tensor, b_scale: torch.Tensor, out_dtype: torch.dtype, out: torch.Tensor | None = None, backend: str = 'auto', scale_granularity_mnk: Tuple[int, int, int] | None = None, output_scale: torch.Tensor | None = None, c: torch.Tensor | None = None, major_a_mode: Literal['K', 'M'] = 'K', major_b_mode: Literal['N', 'K'] = 'K')[source]

Perform batched matrix multiplication with FP8 quantized tensors.

This function computes the batched matrix multiplication of two FP8 quantized tensors, applying scaling factors to produce a result in the specified output data type.

Parameters:

  • a (Tensor) – Input tensor A in FP8 format (e4m3/e5m2). Shape is (batch, m, k) when major_a_mode="K" and (batch, k, m) when major_a_mode="M". The declared major matrix dimension must have stride 1.

  • b (Tensor) – Input tensor B in FP8 format (e4m3/e5m2). Shape is (batch, n, k) by default with major_b_mode="K" and (batch, k, n) when major_b_mode="N". The declared major matrix dimension must have stride 1.

  • a_scale (Tensor) – Scaling factors for tensor A with shape depending on scale_granularity. Should be of fp32 type.

  • b_scale (Tensor) – Scaling factors for tensor B with shape depending on scale_granularity. Should be of fp32 type.

  • out_dtype (torch.dtype) – Data type for the output tensor. torch.bfloat16, torch.float16 and torch.float32 are supported.

  • out (Optional[Tensor]) – Pre-allocated output tensor with shape (batch, m, n). Default is None. If None, a new tensor will be allocated.

  • backend (str) – Backend to use for the operation. Current support backends are “mudnn” and “auto”. Default is “auto”.

  • scale_granularity_mnk (Optional[Tuple[int, int, int]]) – Granularity of scaling for batch, m, and n dimensions respectively. (-1, -1, -1), (1, -1, -1), (1, 128, 128) and (1, 1, 128) are supported. If None, defaults to (-1, -1, -1).

  • c (Optional[Tensor]) – Optional FP32 accumulation tensor with shape (batch, m, n).

  • major_a_mode (str) – "K" treats A as (batch, m, k); "M" treats A as (batch, k, m) and asks MatMulLt to transpose A.

  • major_b_mode (str) – "K" treats B as (batch, n, k) and asks MatMulLt to transpose B. "N" treats B as (batch, k, n). Default is "K".

Returns:

Result tensor with shape (batch, m, n) in the specified output data type.

Return type:

Tensor

mate.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 | None = None, scale_granularity_mnk: Tuple[int, int, int] | None = None, out: torch.Tensor | None = None, out_dtype: torch.dtype | None = None, backend: str = 'auto', output_scale: torch.Tensor | None = None)[source]

Perform groupwise FP8 GEMM operation with scaling.

This function computes the matrix multiplication of two FP8 quantized tensors, applying scaling factors to produce a result in the specified output data type. It supports groupwise quantization with configurable scale granularity.

Parameters:

  • a (Tensor) – Input tensor A with shape (m, k) in FP8 format (e4m3/e5m2). Tensor must be contiguous.

  • b (Tensor) – Input tensor B with shape (n, k) in FP8 format (e4m3/e5m2). Tensor must be contiguous.

  • a_scale (Tensor) – Scaling factors for tensor A. Shape depends on scale_granularity_mnk parameter. Should be of fp32 type. Must be contiguous.

  • b_scale (Tensor) – Scaling factors for tensor B. Shape depends on scale_granularity_mnk parameter. Should be of fp32 type. Must be contiguous.

  • scale_major_mode (str) – Scale major mode “MN” or “K” for groupwise operations. Default is “K”.

  • mma_sm (Optional[int]) – MMA SM configuration. Currently only supports 1. Default is 1.

  • scale_granularity_mnk (Optional[Tuple[int, int, int]]) – Granularity of scaling for m, n, and k dimensions respectively. Default is (1, 128, 128).

  • out (Optional[Tensor]) – Pre-allocated output tensor with shape (m, n). Should be bf16/fp16 when output_scale is None, or fp8_e4m3 when output_scale is provided. If None, a new tensor will be allocated.

  • out_dtype (Optional[torch.dtype]) – Data type for the output tensor when out is None. If out is provided, out.dtype is validated instead. Defaults to torch.bfloat16 without output_scale and fp8_e4m3 with output_scale.

  • backend (str) – Backend to use for the operation. Use "mudnn" when output_scale is None and "mubin" when output_scale is provided. "auto" selects the supported backend for the selected output path.

  • output_scale (Optional[torch.Tensor]) – Quantization scale tensor for FP8 output. If provided, the operation uses the mubin FP8-output path. If None, output is not quantized. Default is None.

Returns:

Result tensor with shape (m, n) in the specified output data type.

Return type:

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): - 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 shape: - (seq_len, seq_len_kv) if max_seqlen_k == 0 - (seq_len, max_seqlen_k) if max_seqlen_k > 0 and dtype torch.float32.

Return type:

torch.Tensor

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