DeepGEMM如何支持MoE模型？

2025-08-30

1.3 K

DeepGEMM 专门为混合专家模型（MoE）提供了分组 GEMM 支持，特别针对专家共享相同形状的场景进行优化。具体使用方法如下：

导入分组 GEMM 函数::
from deep_gemm import m_grouped_gemm_fp8_fp8_bf16_nt_contiguous
准备连续布局的输入数据::
A = torch.randn(4096, 512, dtype=torch.float8_e4m3fn).cuda() # 多个专家的输入拼接 B = torch.randn(512, 1024, dtype=torch.float8_e4m3fn).cuda() group_sizes = [1024, 1024, 1024, 1024] # 每个专家的 token 数
执行分组 GEMM::
C = m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(A, B, group_sizes) print(C)

Advertências:

Ferramenta de IA da estação de consulta rápida