Perplexity 上的 CuTeDSL
我们如何使用 CuTeDSL 构建 GPU 内核,使我们的模型在 NVIDIA GPU 上达到峰值性能。
在 Perplexity,我们使用自主研发的运行时优化服务引擎 (ROSE) 来支持从嵌入到万亿参数大语言模型 (LLM) 的各类模型。我们的运行时架构支撑着包括 Sonar、搜索和嵌入在内的所有 API,旨在赋能研究人员探索最适合的模型与架构,并能迅速将最佳候选模型投入生产。为实现这一目标,我们需要具备快速实施、部署并优化前沿架构推理的能力。本篇博文中,我们将探讨 CuTeDSL 内核如何集成到我们的推理引擎中,从而助力我们构建所需的 GPU 内核,使模型在 NVIDIA GPU 上达到峰值性能。
Perplexity AI 的推理技术
由 Perplexity 各个团队构建的定制模型,均在 NVIDIA Hopper 和 Blackwell GPU 上进行内部托管,并运行于由 AI 推理团队开发的 ROSE 推理引擎之上。ROSE 最初是为了更好地服务于用于语言解码和分类的定制 Llama 模型而构建的,其接口兼容 NVIDIA Triton 推理服务器。随着时间的推移,它已发展成为一个成熟的推理引擎,为众多团队提供服务,支持大型 LLM 以及排序、分类、评分和其他基于 Transformer 的模型,从而驱动 Perplexity 的产品。ROSE 的设计目标非常明确:持续投入以达成峰值性能,但首要优先级始终是适应我们模型不断演进的架构和应用场景。
从宏观层面看,ROSE 由一个 引擎 组成,用于将模型适配到客户端可交互的接口。这些引擎具有高度的可复用性,负责处理请求调度、批处理,以及从底层模型检索解码后的标记 (token)、嵌入或分数。设备的初始化、进程间和节点间通信通道以及权重加载均由其处理。对于 LLM,引擎负责批处理、分块和采样;同时还负责全注意力与线性注意力的 KV 存储分配,以及前缀匹配。嵌入引擎则较为简单,仅执行在线批处理。执行过程经过优化,使得在模型于加速器上并行运行时,引擎能够准备下一轮工作。
ROSE 为建模提供了一系列丰富的自定义层,封装了使用 CuTeDSL、Triton、CUDA、CUTLASS 和 cuBLAS 构建的 GPU 内核。对于矩阵乘法和注意力机制的实现,我们主要依赖来自 CUTLASS 和 cuBLAS 的 NVIDIA 内核。流水线的其余部分(包括嵌入、归一化、MoE 和激活函数内核)均属于 ROSE 的内核库。为最大限度地支持推理运行的硬件,ROSE 配备了专有的设备间通信原语,用于 MoE 路由、分发和合并。虽然其他内核相对简单,但必须经过彻底的专业化设计,才能在各种隐藏维度和设备上实现峰值性能。此外,我们依赖编译时专业化来适配各种模型的各项参数和微调(如 RMS 归一化的权重偏差与否、激活函数等)。因此,我们决定使用 DSL 来实现它们:最初我们使用 Triton 构建了 ROSE。然而,自那以后,我们已采用 CuTeDSL 作为 GPU 编程的主要环境。
为何选择 CuTeDSL?
从数学角度来看,推理内核通常并不复杂。因此,在开发内核时,复杂性通常不在于如何表达所需计算,而在于硬件具体应如何执行计算。为了构建最优内核,我们需要一种能够提供针对底层硬件进行适当抽象的语言。合适的语言使我们能够指定执行计算的高级操作,并在广泛的配置空间的所有点上进行积极的专业化和优化。
CuTeDSL 是一种基于 Python 的 DSL,面向 NVIDIA GPU,构建于 CuTe 布局代数和 MLIR 之上。它能够即时 (JIT) 编译为高度优化的 PTX,同时保留对底层硬件原语的控制。它的目标是提供与 CUTLASS C++ 相同的功能和性能,且无需昂贵的编译时间,从而避免阻碍开发进度。在我们调研或使用过的语言(包括 CUDA、cuTile、TVM IR 和 Triton)中,我们认为它是开发内核的最佳选择。
总而言之,虽然 CuTeDSL 在高级抽象和简洁性上不及 Triton 或 cuTile,但我们发现 CuTe 布局代数及其组合(通过 cute.Tensor 暴露)是合理的替代方案。此外,CuTeDSL 提供对硬件原语的无限制访问,为我们在所有边缘情况下实现峰值性能提供了所需的控制级别。使用 CUDA 时,编译时专业化的机会有限,因为每个额外参数的添加都会导致编译时间增加或构建系统复杂性提升,代价高昂。相比之下,借助即时编译的 DSL,我们可以在不牺牲迭代速度的情况下进行更激进的专业化。与使用 CUDA 作为中间步骤的 TVM IR 不同,我们发现基于 MLIR 并降低到 PTX 的流水线比输出 CUDA 作为中间步骤的流水线更加稳健,且更易于调试。此外,通过使用自定义操作封装 PTX 指令,我们发现可以充分利用底层硬件的所有功能。
开发人员体验
CUDA 内核的大部分编译成本通常消耗在内核的模板展开阶段,因为从 PTX 阶段开始的后续优化和编译在几秒钟内即可完成。在存在大量编译时参数的情况下,对于非简单的内核,此编译阶段的时间可能会激增至数分钟。这一瓶颈由于并行编译的机会仅出现在 IR 和 PTX 级别(而 C++ 模板处理是按顺序进行的)而进一步凸显。尽管可以通过将模板实例化移至单独的编译单元来解决此问题,但这样做需要对构建系统进行重大调整并使用繁琐的宏。相比之下,即时 (JIT) 编译的 CuTeDSL 内核不需要复杂的构建规则,也不会阻碍编译时间,因为开发周期可以专注于配置空间中的一个点,而后续细致的单元测试可以覆盖内核的整个设计空间。
尽管 JIT 编译将成本从开发时转移到了运行时,但其开销是完全可以接受的。当模型部署时,它仅激活整个配置空间中的少数点,在推理引擎启动/预热期间只会产生秒级的开销。相比之下,预编译内核的所有配置可能需要数分钟。此外,通过缓存 JIT 编译后的工件,可以进一步降低开销。
CUDA 的高编译成本在涉及模板分发和专业化(如 CuTe 和 CUTLASS 的情况)时,因难以追踪错误源头而进一步加剧。使用 DSL,可以将错误直接指向产生该错误的代码行,从而显著简化调试。此外,通过中间 MLIR 阶段或依赖附加到生成的 PTX 的调试信息,将生成的代码映射回其原始 DSL 语句要简单得多。得益于 DSL 的这些特性,我们在开发内核时效率大幅提升。
大型配置空间
虽然不同模型的内核有很多相似之处,但它们必须配备一系列对性能有影响的配置项。相关参数包括:
隐藏维度/头维度/注意力头数
输入/输出数据类型
量化方案
编译时常量(偏差、激活函数等)
编译时专业化对于性能和调试都非常有用。特别是在流水线处理器(如 GPU)中,不存在的分支优于总是执行或从不执行的分支。此外,针对特定配置项集合编译内核会产生更少的 PTX/SASS 指令,这在调查过程中使用 cuda-gdb 时更易于分析和调试。
对于给定的批处理大小和配置,可以估算 DRAM 的预期流量和算术运算次数。随后,可以在个案基础上对最优网格分配、同步以及某些硬件功能的使用做出最优决策。CuTeDSL 使我们能够定义和复用组件进行深度专业化,而无需承担不可接受的编译时开销。
预填充与解码专业化
通常,推理内核的性能必须针对两种不同的用例进行优化:
预填充 (Prefill): 小批处理大小(少数块、序列),每个序列有多个标记。对于嵌入和 LLM 的输入标记,均需最大化吞吐量。延迟在几十到几百毫秒之间。
解码 (Decode): 小到大批处理大小,每个序列一个(无推测解码)到几个标记(草稿目标推测、多标记预测)。延迟接近内核启动延迟。
对于 GEMM 和注意力机制,预填充阶段在大量标记上运行,性能受限于设备的计算吞吐量。解码则在小批处理大小上涉及相同的操作,其执行时间主要受限于将激活和权重从 DRAM 加载到流式多处理器 (SM) 和张量核心 (tensor core) 的成本,反而受限于设备的内存带宽。尽管批处理大小或批次中的总标记数是一个相对清晰的指标,但预填充和解码之间的切换阈值较为微妙,因为内存带宽也由模型的隐藏维度和激活的数据类型决定。因此,如果算法差异或网格/块分配会对其中任何一方有利,内核将被拆分为预填充和解码两种实现。
在预填充和解码内核之间普遍不同的参数是块中的线程束 (warp) 数量(线程数除以每个线程束 32 个线程的线程束大小)。将此参数设为编译时常量是理想的,因为它有助于寄存器分配。在 CUDA 中,这通常通过指定 __launch_bounds__ 中的线程数并提供线程束数量作为模板参数来完成:
template <size_t NUM_WARPS>
__global__ void
__launch_bounds__(NUM_WARPS * WARP_SIZE, 1)
some_kernel(...) {
}块中所有线程可用的寄存器数量是一个硬件限制,在 Hopper 和 Blackwell 上为 64K。如果预先知道线程数量,编译器就能准确知道可以为每个线程分配多少寄存器,从而成功编译出更积极优化的内核。然而,这会产生编译时成本,因为现在内核必须针对不同的线程束大小进行重新编译。有了 CuTeDSL,这就不再是问题,因为我们不需要预先为整个配置空间支付编译成本。DSL 还期望块大小为一个常量表达式,隐式执行了 CUDA 中 __launch_bounds__ 所隐含的优化。
除了将线程束数量设为常量外,在预填充和解码之间对其进行调整也是理想的。为了充分利用 GPU 的内存带宽,仅仅为每个块分配大量的线程束是不够的,因为只有使用尽可能多的可用流式多处理器 (SM) 才能实现全部带宽。因此,对于给定数量的标记以及已知的数据类型和隐藏维度,我们首先选择一个理想的每线程向量维度(对于 bfloat16 为 8,对于 float32 为 4),以确保生成的代码使用 128 位加载/存储指令(在 Blackwell 上为 256 位)。在知道向量大小后,我们可以计算内核所需的线程束数量。我们首先尝试最大化 SM 的数量:如果线程束数量少于 SM 数量,则选择每个块一个线程束。这通常是内核的解码变体。如果线程束数量充足,我们将逐渐将数量增加到 8、16 或 32,并进行基准测试,找出性能不再提高的上限。
虽然 Triton 在底层执行了部分此类优化,但我们发现,通过谨慎使用 CuTe 布局代数,我们可以在不受高级编译器局限性的情况下实现类似的结果。特别是当涉及线程束内水平归约时,我们发现 CuTeDSL 为我们提供了生成最优代码所需的细粒度控制。
网格同步
在某些情况下,为了将工作尽可能分布到更多的块上,可以使用 cooperative_groups::this_grid().sync() 来同步所有参与的块。但存在限制:由于块一旦开始执行就无法从 SM 中取消调度,为了保证向前推进,块的数量不能超过设备上可用的 SM 数量,以确保在任何时刻所有块都能更新并轮询同步所涉及的全局内存位置。此限制由 cudaLaunchCooperativeKernel 强制执行。在 Hopper 上,网格屏障的成本约为 3us,外加某些情况下通过隐式内存屏障刷新所有先前内存操作的成本。总的来说,这与将内核拆分为两个启动的开销相当,这导致了一个对性能敏感的设计决策。尽管程序化依赖启动 (PDL) 从 Hopper 开始可以降低启动开销,但对于解码内核而言,避免该开销是值得的。
如果某个工作负载需要的块数量超过了 SM 数量,则应拆分内核,依赖内核启动的隐式屏障。这主要适用于预填充。如果所有工作负载的理想块数量适合网格,则首选网格屏障,例如在处理较少标记的解码内核中。
在 CuTeDSL 中,这些决策可以在主机侧的启动代码中做出,内核可以根据需要进行组合并与网格屏障交错,如下例所示。通过这种混合策略,解码内核可以被优化以达到最小延迟,而预填充内核可以被优化以获得最大吞吐量。与基于 CUDA 的解决方案不同,这种逻辑可以更简洁地定义,且编译时开销更小。
@cute.kernel
def dispatch_kernel_decode(...):
route_tokens(...)
grid_barrier(grid_barrier_ptr)
copy_tokens_to_experts(...)
@cute.kernel
def dispatch_kernel_prefill_route(...):
route_tokens(...)
@cute.kernel
def dispatch_kernel_prefill_copy(...):
copy_tokens_to_experts(...)
@cute.jit
def dispatch(...):
num_route_blocks = ...
num_copy_blocks = ...
if num_copy_blocks < NUM_SMS:
dispatch_kernel_decode(...).launch(
grid=(min(max(num_route_blocks, num_copy_blocks), NUM_SMS), 1, 1)
)
else:
dispach_kernel_prefill_route(...).launch(
grid=(num_route_blocks, 1, 1)
)
dispatch_kernel_prefill_copy(...).launch(
grid=(num_copy_blocks, 1, 1)
)在我们的 MoE 全对全 (all-to-all) 内核中,如下详述,通过 CuTeDSL 灵活地构建带有或不带有网格屏障的内核,并在编译时对其进行优化,帮助我们实现了显著的性能增益。
CuTeDSL 内核
QK 归一化
我们在 CuTeDSL 中构建的第一个内核是一个专用的 RMS 归一化内核,针对小维度进行了优化。QK 归一化用于近期 Qwen 模型的自注意力层,在每个头的基础上应用,处理 64、128 或 256 个元素。与在超过数千的隐藏维度上运行、需要通过共享内存和寄存器进行水平归约的常规 RMS 归一化不同,QK 归一化仅需跨线程束进行归约,并有可能将多个头压缩到一个线程束中。这成为一个用例,我们可以在编译时进行积极优化,为所有头维度和数据类型找到最佳的标记到线程束分配方案。
RMS 归一化是一个相当直观的操作,尽管它需要水平求和来计算向量的平方范数:
r = w * (x * ((x * x).sum(dim=-1) / x.shape[-1] + EPS))由于注意力头通常较小,且每线程操作经过高度向量化,单个线程束可以计算平方范数。为了最大化内存带宽,每个线程应处理 128 位输入,这相当于 8 个 bfloat16 值。因此,半个线程束足以处理一个头,且平方范数可以使用由 shuffle_sync_bfly 执行的水平归约来确定。从头维度和数据类型出发,我们可以在编译时确定每个线程束的头数,并指定参数和归约轮数。
然后,我们可以依靠自动向量化的加载/存储,并使用更底层的架构特定操作进行线程束归约,以相对较少的代码行构建专用内核。该内核确定要处理哪个头(Q 或 K),并识别要应用的相应权重。它从该头加载整个向量并首先在本地求和。随后,基于每个头的线程数,计算平方范数。最后,确定自动向量化的 RMS 归一化结果并写回内存。模型特定参数(如 WEIGHT_BIAS 和 EPS)进行常量折叠。
@cute.kernel
def _qk_norm_kernel(
q: cute.Tensor,
k: cute.Tensor,
q_weight: cute.Tensor,
k_weight: cute.Tensor,
WEIGHT_BIAS: cutlass.Constexpr[float],
EPS: cutlass.Constexpr[float],
) -> None:
# Find the block and thread indices.
bidx, _, _ = cute.arch.block_idx()
bdim, _, _ = cute.arch.block_dim()
tidx, _, _ = cute.arch.thread_idx()
_, num_q_heads, (threads_per_head, vecs_per_thread, vec_size) = q.shape
_, num_k_heads, _ = k.shape
head_dim = threads_per_head * vecs_per_thread * vec_size
# Determine the thread and the vector to work on.
vec_id = (bidx * bdim + tidx) // threads_per_head
thread_id = tidx % threads_per_head
# Determine the token and head indices.
token = vec_id // (num_q_heads + num_k_heads)
head_offset = vec_id % (num_q_heads + num_k_heads)
is_q = head_offset < num_q_heads
# Determine which tensor to work with.
thread_coord = (thread_id, None, None)
head_coord = (token, head_offset if is_q else head_offset - num_q_heads, thread_coord)
# Load the head into registers.
x_ref = q[head_coord] if is_q else k[head_coord]
w = (q_weight if is_q else k_weight)[thread_coord].load()
# Compute the squared sum locally.
x = x_ref.load()
x_f32 = x.to(cute.Float32)
x_square = x_f32 * x_f32
x_square_sum = x_square.reduce(cute.ReductionOp.ADD, init_val=0.0)
# Sum up across the warp.
mask = activemask()
start = cute.arch.WARP_SIZE // (cute.arch.WARP_SIZE // threads_per_head * 2)
for offset in [start >> i for i in range(5) if (start >> i) > 0]:
x_square_sum += cute.arch.shuffle_sync_bfly(
x_square_sum,
offset=offset,
mask=mask,
)
# RMS norm.
rms = cute.math.rsqrt(x_square_sum / head_dim + EPS)
x_norm = ((w + WEIGHT_BIAS) * (x * rms)).to(x.element_type)
x_ref.store(x_norm)与常规 RMS 归一化内核相比,这种专用变体在预填充和解码方面均提供了 2-3 倍的吞吐量和延迟提升。
混合专家 (MoE) 分发/合并
MoE 分发和合并将输入标记重新洗牌至路由到的专家,并在 MoE GEMM 内核执行相关矩阵乘法后执行加权累积。性能是此处关注的重点,因为各种产品依赖于不同规模的 MoE 模型。基于部署类型,我们区分了不同的优化内核:
TP=1:单 GPU 的本地内核。分发将标记洗牌至相关专家,合并进行累积。主要用于 H100 上的 Qwen3 30B 部署。
DP=1, TP ≠ 1:张量并行,单节点内核。输入被复制,分发选择当前秩 (rank) 分配到的标记,合并跨所有秩进行累积。通过 NVLink 或 NVLS 进行跨秩累积。用于大型预填充器。
DP ≠ 1, NVLink:具有可选张量并行的数据并行内核。每个秩向承载特定专家的每个其他秩发送标记。路由信息被聚合,标记跨 NVIDIA NVLink 传输。用于数据并行 H200 或 GB200 部署。
DP ≠ 1, InfiniBand:跨 NVLink 和 NVIDIA Quantum InfiniBand 运行的对等数据并行和张量并行内核。用于多节点 Hopper 部署。
本地
本地内核基于路由表计算分配给每个专家的标记数量,并为属于每个专家的标记列表中的每个标记分配一个唯一索引。路由步骤后,标记被复制到可选的填充缓冲区中,以传递给 MoE GEMM 内核。利用标记位置,合并内核随后在寄存器中执行累积,将最终标记写入输出缓冲区。
分发阶段在预填充和解码之间有所不同:对于少量标记,几乎没有路由信息需要读取,因此我们在单个块中处理它并向网格的其余部分发出信号,由各个块复制标记块。标记分配通过原子增加共享内存中的计数器来确定。由于预填充必须处理更多的路由信息,因此在单独的内核中将路由拆分到多个块中。首先使用共享内存中的计数器在本地分配索引,然后在全局内存中跨所有块聚合,以确定每个标记的位置。为了能够使用网格级屏障,此内核的网格被限制为 SM 的数量。后续内核执行复制,在大量块和线程束上启动,以充分利用可用的内存带宽。
合并内核相对简单,因为它执行向量化累积,依赖于路由信息和分发计算的偏移量。CuTeDSL 在此证明了其价值,因为计算可以通过简单的 CuTe 张量来表达,这些张量可以为所有数据类型和隐藏维度编译成高效的代码。
单节点
在张量并行部署中,路由信息和输入标记被复制。本地分发内核可以完全复用,只需稍作修改即可仅路由和选择分配给当前秩上专家的标记。借助一些参数和编译时专业化,本地分发内核得以复用。
由于秩的数量通常少于标记路由到的专家数量,在一般情况下可以假设每个标记至少有一个副本路由到每个秩。在这种情况下,可以复用本地合并内核来聚合标记的本地副本,随后进行全归约 (all-reduce) 以跨参与秩求和。对于预填充内核,我们启动独立的内核,但解码将两个操作融合到单个启动中。
对于全归约,如果只有两个秩参与,我们使用普通的 NVLink;当涉及 4 个或更多秩时,切换到 NVLS 以实现 10-20% 的吞吐量提升。得益于 DSL,我们可以在 MoE 和常规全归约内核之间共享大部分归约逻辑和对等缓冲区交换。在这种情况下,对硬件的严密控制是有益的,因为 NVLS 需要通过 multimem.ld_reduce 和 multimem.st 进行 DSL 操作。归约本身的实现主要遵循 NCCL 测试中的 all_reduce 示例。
数据并行
数据并行内核依赖 NVLink 为多个设备聚合标记。在分发阶段,在交换标记之前,对等方首先在本地路由,以识别它们发送的每个标记在每个秩上的目标偏移量。随后,它们开始将标记复制到为每个发送方分配的接收缓冲区,以避免冲突。每个标记的发布-获取存储 (release-store) 写入用于发出完成信号。在接收端,路由信息被聚合,以便能够将所有标记浓缩到单个缓冲区中以供 MoE GEMM 使用。接收方通过获取-加载 (acquire-load) 轮询对等方写入的标志,以等待写入完成。
在合并侧,标记通过 NVLink 复制到源秩,并带有指示接收的标记标志。合并内核类似于本地和单节点变体,但在寻址和一些额外的同步逻辑上略有差异。
我们构建了一个 PTX 包装器库,以暴露系统级和设备级的发布-获取加载和存储,从而实现此内核。通过提高同步标志的粒度并分配更多的块以通过 NVLink 发送和接收标记,我们实现了比之前的 CUDA 内核 2 倍的提升。
对等 (Peer-to-Peer)
我们的对等内核之前在配备 ConnectX-7 InfiniBand 适配器的 Hopper GPU 上实现了业界领先的性能。自那以后,我们将合并内核移植到 CuTeDSL 以进一步提升性能。由于 CUDA 版本已经严重依赖编译时专业化,它们成为了我们仓库中编译时间的瓶颈,迫使我们省略了一些参数。通过迁移到 JIT 编译的 DSL,我们针对更多参数进行了专业化并微调了网格分配,在预填充和解码之间引入了进一步的差异,使延迟提升了约 3%。
Blackwell 支持
从 ROSE 内核的角度来看,Blackwell 支持分为两类:架构无关和架构相关内核。依赖张量核心的架构相关内核需要明确的工作来适配。然而,得益于 DSL 的使用,架构无关内核可以轻松适配。
Blackwell GPU 最多拥有 160 个 SM,相比之下 Hopper 为 132 个。对于大多数内核,这是一个编译时常量,主机侧启动代码中的逻辑将其考虑在内,以确定运行的最优(预填充与解码)内核。此外,Blackwell 还引入了 256 位加载指令,可以通过增加向量宽度(CuTe 张量的最后一个维度)轻松发出,从而提高内存带宽利用率。
Blackwell 引入了对大量量化方案(MXFP8、MXFP4、NVFP4)的硬件支持。对此的支持主要是通过调整现有内核以实现 FP8 块级缩放来实施的。JIT 支持在这里至关重要,因为可以在不使编译时配置空间爆炸的情况下支持新模式。过去,我们发现 Triton 在表达跨少数元素(在线程束的四分之一或八分之一内)的水平归约时显得笨拙。在为 NVFP4 和 MXFP8/MXFP4 分别计算 16 或 32 个元素的绝对最大值时,对归约进行更细粒度的控制对于实现最优性能至关重要。
AI 辅助内核生成
作为一种相对较新的语言,大语言模型 (LLM) 在其预训练期间接触 CuTeDSL 代码的机会很少。此外,我们当前代码的风格(内核针对激活和归一化进行了优化)与现有的 CUTLASS 或 Quack 示例略有不同,后者针对 TMA 和架构特定矩阵乘法指令进行了高度优化。达到大语言模型可以协助内核编写的地步是一个循序渐进的过程。
在能够依赖代理进行内核转换之前,我们手动构建了一个初始的内核和实用函数库,遵循一致的模式。随着库的完善,我们可以执行从 CUDA 到 CuTeDSL 的粗略转换,随后进行大量的手动重构。特别地,我们发现模型在指针到张量的转换和向量化方面感到吃力,在试图编译测试时经常陷入修复类型错误的循环。它们对内核的网格特性表现出很少的理解,混淆了线程和网格同步。最令人沮丧的是,它们倾向于臆造不存在的 CUDA 特定库翻译,而不是基于我们自己的 CuTeDSL 替换件构建。随着时间的推移,随着 CuTeDSL 内核库的壮大,我们发现大语言模型可以构建大部分骨架代码并填充空白,从而使我们能够专注于高级设计。
结论与未来工作
与 CUDA 相比,我们发现 CuTeDSL 提供了更好的编译时性能,使我们能够针对更多参数进行专业化并实现性能提升。虽然由于缺乏 Triton 中的某些抽象,代码稍微冗长,但我们认为这种复杂性是合理的,因为对硬件的细粒度控制使我们能够构建更快的内核。
在撰写本文时,CUTLASS C++ 和 CuTeDSL 内核之间的性能差距极小,这为将高性能 GEMM 内核与设备间通信、归一化或激活函数组合和融合开启了令人兴奋的机会。通过致力于这一 DSL,我们将能够从设备中榨取更多的性能,以改善我们从嵌入到 LLM 等所有工作负载的吞吐量。
参考资料
https://github.com/Dao-AILab/quack
https://github.com/NVIDIA/cutlass
https://github.com/NVIDIA/nccl-tests
https://docs.nvidia.com/cuda/parallel-thread-execution/
引用
@online{perplexity2026CutedslAtPerplexity,
author = {{Perplexity} and Nandor Licker and Bex Brooker and Jesse Dwyer and Phi Hoang},
title = {{CuTeDSL at Perplexity}},
date = {2026-05-06},
url = {https://www.perplexity.ai/hub/blog/cutedsl-at-perplexity}
}