Skip to main content

MATE

MATE 是针对生成式 AI 工作负载的集中化库，为 MUSA 平台提供高性能的 Attention 和 GEMM 算子，以及面向 CUDA 的 Python API 兼容性封装。

概述​

什么是 MATE​

MATE（MUSA AI Tensor Engine，MUSA AI 张量引擎）是面向生成式 AI 工作负载的算子库，提供 Attention、GEMM 等高性能算子，以及面向 CUDA 的 Python API 兼容性封装。

关键特性​

MATE 基于 TVM-FFI 构建，专注于高性能 Transformer/LLM 算子实现:

特性说明
FlashAttention高性能注意力机制实现
MLAMulti-head Latent Attention（多头潜在注意力）支持
FP8 GEMM8 位浮点矩阵乘法
Groupwise GEMM分组矩阵乘法支持
MoE GateMoE 门控函数支持

系统要求​

  • MUSA SDK: 5.1.0
  • GPU: MUSA MP31 架构
  • Python: 3.10
  • torch_musa: 2.9.1。更多详情，参见 torch_musa 安装指南

快速开始​

安装 MATE​


# 从源码安装
git clone https://github.com/MooreThreads/mate.git --recursive
cd mate

# (可选）准备构建环境
python -m mate.aot

# 构建 wheel
python -m build --wheel --no-isolation

# 安装生成的 wheel
pip install dist/*.whl

# 验证安装
python -c "import mate; print(mate.__version__)"

Batched GEMM 操作​

import torch
import mate

# Batched GEMM 操作 (FP8)
def bmm_fp8_example():
    batch, m, n, k = 8, 128, 128, 128
    out_dtype = torch.bfloat16
    
    # 创建 FP8 输入 tensors
    a = torch.rand((batch, m, k), device="musa", dtype=torch.float)
    b = torch.rand((batch, n, k), device="musa", dtype=torch.float)
    
    # 量化为 FP8
    fp8_a, scale_a = tensor_quantize_fp8(a, torch.float8_e4m3fn)
    fp8_b, scale_b = tensor_quantize_fp8(b, torch.float8_e4m3fn)
    
    # 创建输出 tensor
    d = torch.empty((batch, m, n), device="musa", dtype=out_dtype)
    
    # 执行 bmm_fp8
    mate.gemm.bmm_fp8(
        fp8_a,
        fp8_b.transpose(-2, -1),
        scale_a,
        scale_b,
        out_dtype,
        d,
        backend="auto",
        scale_granularity_mnk=(-1, -1, -1),
    )
    return d

Groupwise FP8 GEMM​

import torch
import mate

# Groupwise FP8 GEMM (TN layout)
def gemm_fp8_groupwise():
    m, n, k = 1024, 1024, 512
    
    # 创建 FP8 输入 tensors
    a = torch.randn(m, k, dtype=torch.float8_e4m3fn, device='musa')
    b = torch.randn(n, k, dtype=torch.float8_e4m3fn, device='musa')
    
    # Scale shape: (k // 128, m) for row-wise, (k // 128, n) for column-wise
    a_scale = torch.randn(k // 128, m, dtype=torch.float32, device='musa')
    b_scale = torch.randn(k // 128, n, dtype=torch.float32, device='musa')
    
    out = torch.empty(m, n, dtype=torch.bfloat16, device='musa')
    
    mate.gemm.gemm_fp8_nt_groupwise(
        a, b, a_scale, b_scale,
        scale_major_mode='K',
        mma_sm=1,
        scale_granularity_mnk=(1, 128, 128),
        out=out,
        backend='mudnn',
    )
    return out

DeepGEMM-style GEMM​

import torch
import mate
from mate.testing.utils import group_quantize_fp8

# Grouped FP8 GEMM (DeepGEMM-style)
def group_deepgemm_fp8_example():
    nr_group = 4
    ms_per_group = [256, 512, 384, 448]  # 每组不同的 M 维度
    n, k = 1024, 4096
    
    m_total = sum(ms_per_group)
    aligned_ms = [align(m, 128) for m in ms_per_group]
    
    # 创建输入 tensors
    a = torch.rand(m_total, k, device="musa", dtype=torch.float)
    b = torch.rand(nr_group, n, k, device="musa", dtype=torch.float)
    
    # 量化为 FP8
    scale_granularity_mnk = (1, 128, 128)
    scale_a_shape = (m_total, k // 128)
    scale_b_shape = (nr_group, n // 128, k // 128)
    quant_tile_shape = (128, 128)
    
    fp8_a, scale_a = group_quantize_fp8(
        a, scale_a_shape, quant_tile_shape, torch.float8_e4m3fn, "K"
    )
    fp8_b, scale_b = group_quantize_fp8(
        b, scale_b_shape, quant_tile_shape, torch.float8_e4m3fn, "K"
    )
    
    # m_indices: 指定每行属于哪个 group
    m_indices = torch.cat([
        torch.full((m,), i, dtype=torch.int32, device='musa')
        for i, m in enumerate(ms_per_group)
    ])
    
    out = torch.empty(m_total, n, dtype=torch.bfloat16, device='musa')
    
    mate.gemm.group_deepgemm_fp8_nt_groupwise(
        fp8_a, fp8_b, scale_a, scale_b,
        m_indices,
        scale_granularity_mnk,
        out,
        select_tile_m=128,
    )
    return out

MoE Fused Gate​

MoE (Mixture of Experts) 门控函数，用于计算每个 token 路由到各个专家的权重和索引。

使用示例​

import torch
import torchada

# Example shapes
num_tokens = 4
num_experts = 8

# Allocate gating input: shape (num_tokens, num_experts)
input = torch.randn(num_tokens, num_experts, device="cuda", dtype=torch.bfloat16)

# Allocate bias: shape (num_experts,)
bias = torch.randn(num_experts, device="cuda", dtype=torch.bfloat16)

# Very simple invocation: no expert remapping
topk_weights, topk_indices = moe_fused_gate(
    input=input,
    bias=bias,
    num_expert_group=2,                     # split 8 experts into 2 groups
    topk_group=1,                           # select top 1 group per token
    topk=2,                                 # select top 2 experts per token
    num_fused_shared_experts=0,             # no fused shared experts
    routed_scaling_factor=1.0,
    renormalize=True,
    apply_routed_scaling_factor_on_output=False,
    num_token_non_padded=0,                 # 0 means all tokens are valid
    static_index_map=None,
    dynamic_index_map=None,
    dynamic_index_map_valid=None,
    random_index=None,
    num_physical_experts=0,
    map_policy=0,                           # 0 = no index mapping
)

print("topk_weights.shape =", topk_weights.shape)   # (num_tokens, topk)
print("topk_indices.shape =", topk_indices.shape)   # (num_tokens, topk)
print("topk_weights =\n", topk_weights)
print("topk_indices =\n", topk_indices)

函数签名​

函数返回类型参数
moe_fused_gateTuple[torch.Tensor, torch.Tensor]input: torch.Tensor
bias: torch.Tensor
num_expert_group: int
topk_group: int
topk: int
num_fused_shared_experts: int = 0
routed_scaling_factor: float = 1.0
renormalize: bool = True
apply_routed_scaling_factor_on_output: bool = False
num_token_non_padded: int = 0
static_index_map: Optional[torch.Tensor] = None
dynamic_index_map: Optional[torch.Tensor] = None
dynamic_index_map_valid: Optional[torch.Tensor] = None
random_index: Optional[torch.Tensor] = None
num_physical_experts: int = 0
map_policy: int = 0

参数说明​

参数类型说明
inputtorch.Tensor门控输入张量，形状为 (num_tokens, num_experts)
biastorch.Tensor偏置张量，形状为 (num_experts,)
num_expert_groupint专家组数量，将专家分成若干组。
topk_groupint每个 token 选择的专家组数量。
topkint每个 token 选择的专家数量。
num_fused_shared_expertsint融合共享专家数量，0 表示不使用。
routed_scaling_factorfloat路由缩放因子。
renormalizebool是否对权重进行归一化。
apply_routed_scaling_factor_on_outputbool是否在输出权重上应用路由缩放因子。
num_token_non_paddedint非填充 token 数量，0 表示所有 token 都有效。
static_index_mapOptional[torch.Tensor]静态专家索引映射。
dynamic_index_mapOptional[torch.Tensor]动态专家索引映射。
dynamic_index_map_validOptional[torch.Tensor]动态专家索引映射有效标记。
random_indexOptional[torch.Tensor]随机索引张量。
num_physical_expertsint物理专家数量，0 表示不额外指定。
map_policyint索引映射策略，0 表示无映射。

返回值​

返回值类型说明
topk_weightstorch.Tensor形状为 (num_tokens, topk) 的张量，表示选中专家的权重。
topk_indicestorch.Tensor形状为 (num_tokens, topk) 的张量，表示选中专家的索引。

FlashAttention​

MATE 提供高性能的 FlashAttention 实现，包括 variable-length 和 KV cache 两种模式。

flash_attn_varlen_func​

Variable-length FlashAttention，支持变长序列输入。

使用示例​

import torch
import mate

query_lens = [1024, 523, 4063]
kv_lens = [1024, 523, 4063]
num_query_heads = 16
num_kv_heads = 2
head_size_qk = 128
head_size_v = 128
dtype = torch.float16

# 创建输入 tensors
query = torch.randn(
    (sum(query_lens), num_query_heads, head_size_qk), dtype=dtype, device='musa'
)
key_cache = torch.randn(
    (sum(kv_lens), num_kv_heads, head_size_qk), dtype=dtype, device='musa'
)
value_cache = torch.randn(
    (sum(kv_lens), num_kv_heads, head_size_v), dtype=dtype, device='musa'
)

# 累积序列长度
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(dim=0, dtype=torch.int32)
cu_kv_lens = torch.tensor([0] + kv_lens, dtype=torch.int32).cumsum(dim=0, dtype=torch.int32)

max_query_len = max(query_lens)
max_kv_len = max(kv_lens)

# 调用 flash_attn_varlen_func
output, lse = mate.flash_attn_varlen_func(
    query,
    key_cache,
    value_cache,
    cu_query_lens,
    cu_kv_lens,
    max_query_len,
    max_kv_len,
    seqused_q=None,
    seqused_k=None,
    softmax_scale=None,
    causal=True,
    qv=None,
    q_descale=None,
    k_descale=None,
    v_descale=None,
    window_size=(-1, -1),
    attention_chunk=0,
    softcap=0.0,
    num_splits=1,
    pack_gqa=None,
    deterministic=False,
    sm_margin=0,
    return_attn_probs=True,
    backend='mubin',
)

函数签名​

函数返回类型参数
flash_attn_varlen_funcUnion[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]q: torch.Tensor
k: torch.Tensor
v: torch.Tensor
cu_seqlens_q: Optional[torch.Tensor] = None
cu_seqlens_k: Optional[torch.Tensor] = None
max_seqlen_q: Optional[int] = None
max_seqlen_k: Optional[int] = None
seqused_q: Optional[torch.Tensor] = None
seqused_k: Optional[torch.Tensor] = None
softmax_scale: Optional[float] = None
causal: bool = False
qv: Optional[torch.Tensor] = None
q_descale: Optional[torch.Tensor] = None
k_descale: Optional[torch.Tensor] = None
v_descale: Optional[torch.Tensor] = None
window_size: Tuple[int, int] = (-1, -1)
attention_chunk: int = 0
softcap: float = 0.0
num_splits: int = 1
pack_gqa: Any = None
deterministic: bool = False
sm_margin: int = 0
return_attn_probs: bool = False
backend: str = "auto"

参数说明​

参数类型说明
qtorch.TensorQuery 张量。非变长输入形状为 [batch_size, seqlen, nheads, headdim]；变长输入形状为 [total_q, nheads, headdim]。
ktorch.TensorKey 张量。非变长输入形状为 [batch_size, seqlen_k, nheads_k, headdim]；变长输入形状为 [total_k, nheads_k, headdim]。
vtorch.TensorValue 张量。非变长输入形状为 [batch_size, seqlen_k, nheads_k, headdim]；变长输入形状为 [total_k, nheads_k, headdim_v]。
cu_seqlens_qOptional[torch.Tensor]Query 累积序列长度张量，形状为 [batch_size + 1]。
cu_seqlens_kOptional[torch.Tensor]Key/Value 累积序列长度张量，形状为 [batch_size + 1]。
max_seqlen_qOptional[int]Query 最大序列长度；变长输入时需要提供。
max_seqlen_kOptional[int]Key/Value 最大序列长度。
seqused_qOptional[torch.Tensor]Query 实际使用的序列长度。
seqused_kOptional[torch.Tensor]Key/Value 实际使用的序列长度。
softmax_scaleOptional[float]对 QKTQK^T 应用 softmax 前的缩放因子；默认使用 1 / sqrt(headdim)。
causalbool是否应用因果注意力掩码。
qvOptional[torch.Tensor]可选的 QV 输入张量。
q_descaleOptional[torch.Tensor]Query 反量化缩放张量。
k_descaleOptional[torch.Tensor]Key 反量化缩放张量。
v_descaleOptional[torch.Tensor]Value 反量化缩放张量。
window_sizeTuple[int, int]滑动窗口大小；当前仅支持 (-1, -1)。
attention_chunkintAttention 分块大小，0 表示不启用分块。
softcapfloatSoftcap 参数；当前仅支持 0.0。
num_splitsint计算拆分数量。
pack_gqaAnyGQA 打包相关配置。
deterministicbool是否启用确定性执行。
sm_marginint预留的 SM 数量。
return_attn_probsbool是否返回 attention scores 的 log-sum-exp (LSE)。
backendstr后端选择，建议使用默认值 "auto"。

返回值​

返回值类型说明
outputtorch.TensorAttention 输出张量，形状为 [total_q, nheads, headdim_v]。
lsetorch.Tensor当 return_attn_probs=True 时返回，表示 log-sum-exp (LSE)，形状为 [nheads, total_q]。

flash_attn_with_kvcache​

带 KV cache 的 FlashAttention（paged KV cache 模式）。

使用示例​

import torch
import mate

# Paged KV cache attention
def flash_attn_with_kvcache_example():
    seq_lens = [(2, 1328), (2, 18), (2, 463)]
    num_query_heads = 16
    num_kv_heads = 2
    head_size_qk = 128
    head_size_v = 128
    dtype = torch.bfloat16
    block_size = 64
    num_blocks = 2048
    is_causal = True
    
    num_seqs = len(seq_lens)
    query_lens = [x[0] for x in seq_lens]
    kv_lens = [x[1] for x in seq_lens]
    max_query_len = max(query_lens)
    max_kv_len = max(kv_lens)
    scale = head_size_qk ** -0.5
    
    # 创建输入 tensors
    query = torch.randn(
        num_seqs, max_query_len, num_query_heads, head_size_qk, dtype=dtype, device='musa'
    )
    key_cache = torch.randn(
        num_blocks, block_size, num_kv_heads, head_size_v, dtype=dtype, device='musa'
    ) / 10
    value_cache = torch.randn_like(key_cache) / 10
    
    # KV 序列长度
    kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
    
    # 构建 block tables
    max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
    block_tables = torch.randint(
        0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32, device='musa'
    )
    
    # 调用 flash_attn_with_kvcache
    output, lse, *rest = mate.flash_attn_with_kvcache(
        query,
        key_cache,
        value_cache,
        k=None,
        v=None,
        qv=None,
        rotary_cos=None,
        rotary_sin=None,
        cache_seqlens=kv_lens_tensor,
        cache_batch_idx=None,
        cache_leftpad=None,
        page_table=block_tables,
        cu_seqlens_q=None,
        cu_seqlens_k_new=None,
        max_seqlen_q=max_query_len,
        rotary_seqlens=None,
        q_descale=None,
        k_descale=None,
        v_descale=None,
        softmax_scale=scale,
        causal=is_causal,
        window_size=(-1, -1),
        attention_chunk=0,
        softcap=0.0,
        rotary_interleaved=False,
        scheduler_metadata=None,
        num_splits=1,
        pack_gqa=None,
        sm_margin=0,
        return_softmax_lse=True,
    )
    return output, lse

函数签名​

函数返回类型参数
flash_attn_with_kvcacheUnion[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]q: torch.Tensor
k_cache: torch.Tensor
v_cache: torch.Tensor
k: Optional[torch.Tensor] = None
v: Optional[torch.Tensor] = None
qv: Optional[torch.Tensor] = None
rotary_cos: Optional[torch.Tensor] = None
rotary_sin: Optional[torch.Tensor] = None
cache_seqlens: Optional[Union[int, torch.Tensor]] = None
cache_batch_idx: Optional[torch.Tensor] = None
cache_leftpad: Optional[torch.Tensor] = None
page_table: Optional[torch.Tensor] = None
cu_seqlens_q: Optional[torch.Tensor] = None
cu_seqlens_k_new: Optional[torch.Tensor] = None
max_seqlen_q: Optional[int] = None
rotary_seqlens: Optional[torch.Tensor] = None
q_descale: Optional[torch.Tensor] = None
k_descale: Optional[torch.Tensor] = None
v_descale: Optional[torch.Tensor] = None
softmax_scale: Optional[float] = None
causal: bool = False
window_size: Tuple[int, int] = (-1, -1)
attention_chunk: int = 0
softcap: float = 0.0
rotary_interleaved: bool = True
scheduler_metadata: Optional[torch.Tensor] = None
num_splits: int = 0
pack_gqa: Any = None
sm_margin: int = 0
return_softmax_lse: bool = False

参数说明​

参数类型说明
qtorch.TensorQuery 张量。非变长输入形状为 [batch_size, seqlen, nheads, headdim]；变长输入形状为 [total_q, nheads, headdim]。
k_cachetorch.TensorKey cache 张量。Paged KV cache 模式下形状为 [num_blocks, page_block_size, nheads_k, headdim]，其中 page_block_size 必须为 64。
v_cachetorch.TensorValue cache 张量。Paged KV cache 模式下形状为 [num_blocks, page_block_size, nheads_k, headdim_v]，其中 page_block_size 必须为 64。
kOptional[torch.Tensor]可选的新 Key 张量。
vOptional[torch.Tensor]可选的新 Value 张量。
qvOptional[torch.Tensor]可选的 QV 输入张量。
rotary_cosOptional[torch.Tensor]RoPE (rotary positional embedding) 的 cosine 张量。
rotary_sinOptional[torch.Tensor]RoPE (rotary positional embedding) 的 sine 张量。
cache_seqlensOptional[Union[int, torch.Tensor]]KV cache 序列长度；为张量时形状为 [batch_size]。
cache_batch_idxOptional[torch.Tensor]Cache batch 索引张量。
cache_leftpadOptional[torch.Tensor]Cache 左侧 padding 信息。
page_tableOptional[torch.Tensor]Page table 张量，形状为 [batch_size, max_num_blocks_per_seq]；当前 paged KV cache 模式需要提供。
cu_seqlens_qOptional[torch.Tensor]Query 累积序列长度张量。
cu_seqlens_k_newOptional[torch.Tensor]新 Key/Value 的累积序列长度张量。
max_seqlen_qOptional[int]Query 最大序列长度。
rotary_seqlensOptional[torch.Tensor]RoPE 使用的序列长度。
q_descaleOptional[torch.Tensor]Query 反量化缩放张量。
k_descaleOptional[torch.Tensor]Key 反量化缩放张量。
v_descaleOptional[torch.Tensor]Value 反量化缩放张量。
softmax_scaleOptional[float]对 QKTQK^T 应用 softmax 前的缩放因子；默认使用 1 / sqrt(headdim)。
causalbool是否应用因果注意力掩码。
window_sizeTuple[int, int]滑动窗口大小；当前仅支持 (-1, -1)。
attention_chunkintAttention 分块大小，0 表示不启用分块。
softcapfloatSoftcap 参数；当前仅支持 0.0。
rotary_interleavedboolRoPE 是否使用 interleaved 布局。
scheduler_metadataOptional[torch.Tensor]调度器元数据。
num_splitsint计算拆分数量。
pack_gqaAnyGQA 打包相关配置。
sm_marginint预留的 SM 数量。
return_softmax_lsebool是否返回 attention scores 的 log-sum-exp (LSE)。

返回值​

返回值类型说明
outputtorch.TensorAttention 输出张量。
lsetorch.Tensor当 return_softmax_lse=True 时返回，表示 log-sum-exp (LSE)。

Multi-head Latent Attention (MLA)​

MLA（Multi-head Latent Attention，多头潜在注意力）是用于 DeepSeek 等模型的高效注意力机制，支持压缩 KV cache。

使用示例​

import torch
import mate

# MLA with KV cache (用于 DeepSeek 等模型)
def mla_example():
    batch_size = 16
    num_heads = 16
    head_dim = 576
    kv_dim = 64  # 压缩后的 KV 维度
    
    # q_nope: (batch, num_heads, head_dim - kv_dim) - 不含位置信息的 Q
    q_nope = torch.randn(batch_size, num_heads, 512, dtype=torch.float16, device='musa')
    # q_pe: (batch, num_heads, kv_dim) - 位置编码部分
    q_pe = torch.randn(batch_size, num_heads, 64, dtype=torch.float16, device='musa')
    # ckv: (batch, kv_len, num_heads, head_dim) - compressed KV cache
    ckv = torch.randn(batch_size, 512, num_heads, head_dim, dtype=torch.float16, device='musa')
    # kpe: (batch, kv_len, num_heads, kv_dim) - compressed KV cache 的位置编码部分
    kpe = torch.randn(batch_size, 512, num_heads, 64, dtype=torch.float16, device='musa')
    
    # page_table: [batch] - page table for paged attention
    page_table = torch.arange(batch_size, dtype=torch.int32, device='musa')
    # kv_len: [batch] - KV sequence lengths
    kv_len = torch.tensor([512] * batch_size, dtype=torch.int32, device='musa')
    
    out = mate.mla(
        q_nope, q_pe, ckv, kpe,
        page_table, kv_len,
        sm_scale=1.0,
        is_causal=True
    )
    return out

函数签名​

函数返回类型参数
mlatorch.Tensorq_nope: torch.Tensor
q_pe: torch.Tensor
ckv: torch.Tensor
kpe: torch.Tensor
page_table: torch.Tensor
kv_len: torch.Tensor
sm_scale: float = 1.0
is_causal: bool = True

参数说明​

参数类型说明
q_nopetorch.Tensor不含位置信息的查询张量，形状为 (batch, num_heads, head_dim - kv_dim)。
q_petorch.Tensor位置编码部分的查询张量，形状为 (batch, num_heads, kv_dim)。
ckvtorch.Tensor压缩的 KV cache，形状为 (batch, kv_len, num_heads, head_dim)。
kpetorch.Tensor压缩的 KV cache 的位置编码部分，形状为 (batch, kv_len, num_heads, kv_dim)。
page_tabletorch.Tensor页表张量，形状为 (batch,)。
kv_lentorch.TensorKV 序列长度张量，形状为 (batch,)。
sm_scalefloatsoftmax 缩放因子。
is_causalbool是否应用因果注意力掩码，默认 True。

返回值​

返回值类型说明
outtorch.TensorMLA 输出张量。

性能优化​

量化粒度​

# scale_granularity_mnk = (m, n, k)
# -1 表示 per-tensor
# 正整数表示 per-group

# Per-tensor scaling
scale_granularity = (-1, -1, -1)

# Per-channel scaling (m dimension)
scale_granularity = (1, 128, 128)

result = mate.gemm.bmm_fp8(
    a, b, a_scale, b_scale,
    scale_granularity_mnk=scale_granularity
)

调试与诊断​

启用日志​

# 启用详细日志
export MATE_LOGLEVEL=3

架构说明​

支持的 MUSA 架构​

架构支持的算子
MP31GEMM, FlashAttention, MLA, FP8

数据类型支持​

数据类型说明支持的算子
float16FP16All
bfloat16BF16All
float8_e4m3FP8 (E4M3)GEMM
float8_e5m2FP8 (E5M2)GEMM

版本信息​

版本日期更新
v0.2.02026-04-15MP31 支持, MLA, FP8，FlashAttention 优化等
v0.1.02025-12-19初始版本

相关文档​

版权所有 2024 摩尔线程，本文档受国际版权法保护。

摩尔线程和摩尔线程徽标是摩尔线程智能科技 (北京)有限责任公司的注册商标。

免责声明

本文档提供有关摩尔线程产品的信息。本文档并未授权任何知识产权的许可，并未以明示或暗示，或以禁止反言或其他方式授予任何知识产权许可。