Helion炸裂！AI推理提速153%，代码减半，震惊业界！

人工智能技术飞速发展，大模型（LLM）的应用日益普及，其背后推理效率的提升，无疑是业界关注的焦点。在这一背景下，如何让核心算法在不同硬件平台之间实现高性能且便捷的迁移，成为了工程师们面临的一大挑战。最近，一个令人振奋的消息传来，PyTorch社区推出了一项名为Helion的新技术，旨在为开发者提供一个基于PyTorch的特定领域语言，以更轻松地开发高性能且具备良好可移植性的核心算法。

新媒网跨境获悉，Helion内置了强大的自动调优功能，它所展现出的潜力，有望在性能可移植性方面超越目前广泛使用的Triton，开辟新的局面。为了验证Helion的这一承诺并深入学习其特性，我们的团队选择了一项在人工智能领域中性能至关重要的核心任务——分页注意力（Paged Attention）机制进行实践。分页注意力是vLLM推理框架的核心基石，其效率直接影响到大模型推理的整体表现。过去一年里，我们成功为vLLM贡献了一套完全由Triton编写的性能和平台可移植的注意力后端，该后端不依赖任何外部组件，并且能够在英伟达、AMD和英特尔等不同品牌的图形处理器（GPU）上高效运行。考虑到我们此前在PyTorch大会上分享的经验，我们决定在vLLM中实现其中一个核心算法（即unified_attention_2d）的Helion版本，作为一项实验性后端，其相关代码更新（PR#27293）已经提交。

大模型推理框架的幕后英雄：vLLM、Triton与Helion的崛起

深入探讨Helion的实践之前，我们有必要先了解一下与其紧密相关的几个关键技术。

vLLM：驱动大模型高效推理的引擎

vLLM作为PyTorch基金会的一部分，在大模型推理领域被广泛应用。随着AI应用对实时性、高吞吐量的需求日益增长，vLLM正逐渐成为生产环境中的标准选择。它不仅能够兼容英伟达、AMD、英特尔等主流GPU，甚至还能支持谷歌的TPU、华为的昇腾NPU、AWS Inferentia或IBM Spyre等定制化加速器。vLLM之所以能为几乎所有大模型提供高效能的推理服务，得益于其精心设计的软件架构以及与torch.compile组件的深度融合，这使得模型能够在多种硬件上发挥出最大潜力。

Triton：性能优异但挑战犹存的领域特定语言

Triton是一种基于Python编写的领域特定语言（DSL），能够实现针对AMD、英特达和英伟达GPU的即时（JIT）编译。通过Triton编写的核心算法，已经展现出卓越的性能，并具备一定的可移植性。例如，我们团队此前用Triton实现了分页注意力机制，同一份核心代码在英伟达H100和AMD MI300上均取得了业内领先的性能表现。这项成果的细节可以在我们发布的论文或相关博客文章中查阅。尽管Triton也提供了自动调优功能，但在实际生产环境中，其应用受到了严重限制，无法充分发挥其对性能可移植性的积极影响。因此，在我们的Triton注意力后端中，目前仍主要依赖简单的if-else语句作为启发式规则。此外，Triton也是PyTorch Inductor的输出语言，后者正是torch.compile的核心编译模块。

Helion：展望未来，迈向更智能的自动调优

Helion是另一个新近崛起的领域特定语言，于去年十月底进入测试阶段。Helion自称为“分块式PyTorch”，它主要有两个核心目标：首先，将分块（tiling）的概念引入PyTorch，让开发者能够使用PyTorch的API编写分块程序；其次，通过更全面的自动调优来增强代码的可移植性。与Triton不同，Helion的自动调优器不仅拥有更实用的缓存机制，其调优的自由度也大大提高。这种更大的自由度体现在，Helion的自动调优器不仅能够调整底层编译参数，比如warp数量或流水线深度，还能改变算法层面的实现细节。此外，它还集成了先进的搜索算法，这正是我们之前在Triton语境下也曾探索过的方向。

分块PyTorch下的分页注意力：实现细节揭秘

启动网格与并行化策略

在最初的实践中，我们选择重新实现我们统一注意力Triton核心算法中相对简单的“2D”版本。之所以称之为“2D”，是因为这个核心算法采用了二维的启动网格。我们选择这个版本，是因为初期认为并行的分块softmax实现会过于复杂。

然而，由于Helion处理启动网格的方式与Triton有所不同，我们并没有完全照搬2D方法，而是围绕“Q块”的核心概念构建了Helion核心算法。下图清晰地展示了“Q块”的概念：
Q块的概念

图1：Helion核心算法中“Q块”的概念。图中展示了一个请求需要计算的三个维度。注意力核心算法需要遍历到query_length（底部轴）为止的所有查询令牌。在我们的核心算法中，我们同时获取多个查询令牌，这个分块大小（TILE_Q）是可调优的。接下来，对于每个令牌，有多个查询头和KV头（左侧轴）。我们重新实现了QGA优化，以便一次性获取一个KV头的所有查询头。每个KV头的查询头数量（QpKV）是这个方向的分块大小，称为TILE_M。最后，我们需要遍历这个查询对应的KV缓存，直到当前的上下文长度，以可调优的TILE_N大小的块（对角线轴）进行。在这个内循环中，实际的注意力计算，包括矩阵乘法（hl.dot），通过在线softmax实现。在核心算法中，还有一个围绕这一切的额外循环，用于遍历批次中的所有请求（图中未显示）。

然而，vLLM处理的输入张量，其第一个维度是序列数量与查询长度的组合（通常称为“扁平化变长”布局）。因此，vLLM会额外提供一个张量作为索引，用于标识哪个令牌属于哪个序列。

经过一番实验，我们最终确定了上述的四层循环方法：

# 注册可调优的块大小
q_block_size = hl.register_block_size(1, q_block_padded_size)
num_pages_at_once = hl.register_block_size(1, 32)
# 外层循环 -> 成为启动网格
for seq_tile, tile_m in hl.tile([num_seqs, num_query_heads],
                                 block_size=[1, num_queries_per_kv],):
    seq_len = t_seq_lens[seq_tile]
    query_start = t_query_start_lens[seq_tile]
    query_end = t_query_start_lens[seq_tile + 1]
    query_len = query_end - query_start
    # 循环处理一个请求的查询
    for tile_q in hl.tile(query_len, block_size=q_block_size):
        ...
        # 循环处理KV缓存
        for tile_n in hl.tile(num_blocks,
                                 block_size=num_pages_at_once):
            ...

可以看到，最外层循环是一个融合循环，包含两个维度：批次中的序列和QpKV。这个外层循环将成为Triton中的启动网格（Helion推荐使用hl.tile而非hl.grid，即使是对于外层循环）。由于在循环之前，例如边界计算中也需要调优后的块大小，我们明确地预先注册了这些块大小。此外，为了简化启动网格，我们调整了实现中循环的顺序，使其与上述描述相反，让外层循环遍历查询头。

接下来，第二个循环是遍历所选序列的查询长度，并使用一个可调优的分块大小。但需要注意的是，我们对这个分块大小的上限进行了填充（参见q_block_padded_size），这样即时编译器（JIT）和自动调优器就不会因为所有可能的查询长度组合而被频繁触发。相反，我们只提供2的幂次的填充长度，从而减少了运行时的JIT/自动调优开销。最内层循环则遍历所选序列中KV缓存页的数量。因此，相应注册块大小的上限也表示32页KV缓存内存（每页例如包含16个令牌）。

由此生成的Triton代码可能看起来像这样：

# src[helion_unified_attention.py:129]: for seq_tile, tile_m in hl.tile(
# src[helion_unified_attention.py:130]:     [num_seqs, num_query_heads],
# src[helion_unified_attention.py:131]:     block_size=[1, num_queries_per_kv],
# src[helion_unified_attention.py:129-132]: ...
num_pid_m = num_seqs
num_pid_n = tl.cdiv(32, _BLOCK_SIZE_3)
inner_2d_pid = tl.program_id(0)
num_pid_in_group = 4 * num_pid_n
group_id = inner_2d_pid // num_pid_in_group
first_pid_m = group_id * 4
group_size_m = min(num_pid_m - first_pid_m, 4)
pid_0 = first_pid_m + inner_2d_pid % num_pid_in_group % group_size_m
pid_1 = inner_2d_pid % num_pid_in_group // group_size_m
offset_2 = pid_0
offset_3 = pid_1 * _BLOCK_SIZE_3
...
# src[helion_unified_attention.py:141]: for tile_q in hl.tile(query_len, block_size=q_block_size):
# src[helion_unified_attention.py:141-252]: ...
for offset_9 in tl.range(0, v_0.to(tl.int64), _BLOCK_SIZE_0, loop_unroll_factor=2, num_stages=2, disallow_acc_multi_buffer=False, flatten=False):
    indices_9 = offset_9 + tl.arange(0, _BLOCK_SIZE_0).to(tl.int64)
    # src[helion_unified_attention.py:174]: for tile_n in hl.tile(num_blocks, block_size=num_pages_at_once):
    # src[helion_unified_attention.py:174-244]: ...
    for offset_10 in tl.range(0, v_19.to(tl.int64), _BLOCK_SIZE_1, loop_unroll_factor=1, num_stages=1, disallow_acc_multi_buffer=True):
        indices_10 = offset_10 + tl.arange(0, _BLOCK_SIZE_1).to(tl.int64)
        mask_1 = indices_10 < v_19

如上所示，这段代码采用了Helion自动调优器确定的pid_type="flat"程序启动方式。在这种程序类型中，核心算法只有一个“真实”的PID（tl.program_id(0)），并由此派生出所有其他的“局部”ID。

Helion中的分块艺术

通常来说，Helion要求分块操作需要是通用的，这意味着我们不能假定它们总是与vLLM的KV缓存块大小相对应，因此需要相应地编写程序。Helion中的分块功能非常强大，由hl.tile生成的块会自动适应不完整的块。例如，如果分块大小为32，张量形状为63，那么第二个块将只有31个元素，并且所有的掩码都会自动生成。

然而，在分页注意力机制中，存在额外的约束：我们必须始终从KV缓存中加载完整的页面，因为在编译时我们无法确定是否需要完整的页面。这对Helion来说并非问题，但这意味着我们需要创建自己的掩码。

分块的动态性也给我们带来了另一个挑战：想象一个批次，其中查询长度分别为7、2、1，如下图所示：
访问扁平化“变长”张量

图2：使用固定分块大小访问扁平化“变长”张量需要手动掩码。因此，如果我们总是使用相同的分块大小（本例中为4）进行循环，第二个分块中将包含第二个请求的令牌，以及第一个请求的最后三个令牌！同样，下一个分块会混合第二个和第三个请求的令牌。然而，我们不能为批次中的每个序列改变分块大小，因为每个编译的核心算法只能有一个分块大小（请记住，这是一个编译时常量）。一个解决方案是这里只使用块大小为1，但从我们开发vLLM中Triton注意力后端了解到，这通常表现非常差。因此，唯一的其他选择是调整分块的索引并使用hl.load应用手动掩码。

adjusted_tile_q_index = query_start + tile_q.begin + hl.arange(q_block_size)
query_head_offset = tile_m.begin + hl.arange(num_queries_per_kv)
q_load_mask = adjusted_tile_q_index[:, None, None] < query_end # (tile_q, tile_m, HEAD_SIZE)
q = hl.load(t_query,
            [adjusted_tile_q_index, query_head_offset, hl.arange(head_size)],
            extra_mask=q_load_mask)

总体来看，Helion的这个核心算法实现仅需133行代码（vLLM格式，含注释），而Triton版本则需要295行。与编写Triton版本相比，尽管Helion的编程模型要求我们进行算法上的调整，但编写Helion核心算法的过程对我们来说非常直接。在Triton中，大量时间和精力（以及代码行数）都用于确保所有块都具有正确的掩码和边界，这要求在所有维度上手动处理张量步长和偏移量。Helion则能自动且正确地完成这些工作，大大节省了开发者的精力，尤其是调试时间。此外，Helion还通过自动调优处理其他底层细节，例如实际的启动网格实现、块大小发现或张量内存分配。

智慧之选：Helion的自动调优机制

自动调优是Helion的显著优势和核心功能之一。它不仅能够探索各种不同的调优参数，还能自行检测这些参数所有可能的有效值。用户只需定义分块或块大小的下限和上限。甚至可以这样定义：

for tile_n in hl.tile(seqlen//page_size, block_size=None):

这与Triton形成了鲜明对比，Triton需要用户列出所有可能的有效配置（通常会用大量嵌套的num_stages=s for s in [1, 2, 3, 4, 5, 6, 8, 42]等方式实现）。除了API不便之外，这也容易导致用户错过在某个平台上可能表现出色的组合。Helion的方法通过要求用户仅在符号层面定义形状，然后推导出所有可能的组合来解决这个问题。如果不是最外层的hl.tile循环（这将成为启动网格），则可以从张量形状中派生出分块的边界。Helion还会根据默认配置，检查自动调优过程中创建的每个核心算法变体的正确性，并丢弃所有数值误差过高的实验。在我们实验初期，这个基线（或Helion称之为“默认配置”）是自动从已发现的搜索空间“中间”某个位置派生出来的。然而，这种自动发现方式给我们带来了一些问题，因为生成的结果对于我们的核心算法来说是一个无效配置，无法与vLLM中16的页面大小配合使用。因此，自动调优完全失效了，我们不得不修补Helion并手动定义默认配置。是的，Helion目前仍处于测试阶段并积极开发中，这个问题后来通过允许用户定义一个外部函数作为自动调优基线（autotune_baseline_fn=callable()）而得到了解决。这个功能解决了我们的问题，我们随后可以将现有的Triton实现定义为基线，我们知道它能提供非常好的性能和正确的结果。这极大地简化并增强了我们的自动调优过程。我们非常欣赏的另一个功能是Helion添加的用户自定义设置“调优工作量”（autotune_effort=），它可以是'none'（无）、'quick'（快速）或'full'（完整）。从我们在vLLM中开发Triton注意力后端时获得的经验，我们了解到由于分页内存和由此产生的有限有效配置数量，自动调优过程是受限的，通常进行数天的自动调优并不会带来太多回报。因此，一旦这个功能可用，我们便将工作量设置为快速。但像往常一样，“快速”和“慢速”是相对的，Helion自动调优器的快速模式仍然需要10小时来为72种不同场景（如批次大小、序列长度、头部大小等）调优我们的核心算法。这比“完整”（或默认）设置所需的25小时要快，但可能不如我们所希望的那么“快速”。自动调优器完成后，它还会打印出推荐的最佳配置：

[2752s] Autotuning complete in 2752.1s after searching 5353 configs.
One can hardcode the best config and skip autotuning with:
@helion.kernel(config=helion.Config(
    block_sizes=[32, 4],
    indexing=['pointer', 'pointer', 'pointer', 'pointer', 'tensor_descriptor', 'pointer', 'tensor_descriptor', 'pointer', 'tensor_descriptor'],
    l2_groupings=[2],
    load_eviction_policies=['', '', '', '', '', 'last', 'last', ''],
    loop_orders=[[1, 2, 0], [1,0]],
    num_stages=6,
    num_warps=8,
    pid_type='flat',
    range_flattens=[None, True, True, True],
    range_multi_buffers=[None, None, None, False],
    range_num_stages=[],
    range_unroll_factors=[0, 1, 2, 1],
    range_warp_specializes=[]),
    static_shapes=False)

这再次展现了Helion能够探索的各种不同配置参数。因此，Helion在不同的自动调优算法上投入大量精力是至关重要且有意义的。目前，默认采用的是遗传算法。在上面显示的例子中，Helion找到了32个令牌（TILE_Q）和4个页面（TILE_N，相当于64个令牌）的分块大小作为最佳配置。它还能够确定如何寻址所有相关的张量，或者循环是否应该被扁平化和重新排序。对所有参数的详细讨论超出了本文的范围。我们对Triton核心算法进行自动调优的经验告诉我们，为广泛的场景（在微基准测试设置中）进行调优，然后选择少数配置用于vLLM，并使用决策树或其他启发式方法在它们之间进行选择，是一种很好的权衡。

然而，当前版本Helion的一个缺点是，它期望为整个核心算法提供一个配置，我们无法像在Triton中那样区分对预填充批次（prefill batch）或解码批次（decode batch）有利的配置。因此，对于本次博客中的实验，我们针对在英伟达GPU上运行的场景，选择了6种配置进行“实时”调优；对于AMD GPU，则选择了7种配置。

性能深度剖析：Helion与Triton的较量

尽管卓越的开发者体验令人欣喜，但如果性能表现不如Triton，我们可能仍会犹豫是否采用Helion。因此，我们对Helion中新实现的分页注意力机制，在英伟达H100和AMD MI300X平台上与我们成熟的Triton核心算法进行了基准测试。对于推理服务器这类应用场景，我们总是需要从两个方面进行考量：核心算法本身的独立性能表现，以及新核心算法对整个系统（即vLLM）的影响。为此，我们首先在微基准测试中单独分析核心算法的性能，然后进行端到端分析。

微基准测试：揭示核心算法的内在实力

为了分析核心算法性能，我们使用了我们专门为开发vLLM Triton注意力后端而构建的微基准测试套件。在这里，我们基于Llama3.1-8B架构（128头部大小，32个查询头和8个KV头）设置测试参数，并根据实际样本，变化序列长度和批次大小。批次中包含的序列长度是可变的，即vLLM默认采用的“变长”模式。所有实验均使用Helion 0.2.4、Triton 3.5.0和PyTorch 2.9。

针对每个GPU平台，我们制作了两张图表，数据相同，但在其中一张图中按批次内解码请求的比例排序，另一张图则按混合预填充-解码变长批次中的最大序列长度排序。请注意，所有测量均在启用CUDA/HIP图的情况下进行，因此我们不评估任何软件开销，如编译时间或启动时间，只关注纯粹的核心算法性能。由于我们侧重于Triton和Helion核心算法之间的比较，因此在每张图中，所有数据都以最左侧的Triton结果进行归一化。每张图显示了三个结果：作为基线的Triton 2D核心算法（vLLM当前版本），具有动态形状（static_shapes=False）的Helion核心算法（使用为每个平台在vLLM中进行小规模“实时”自动调优所选的配置），以及具有静态形状（static_shapes=True）并经过全面自动调优的Helion核心算法。当然，将动态形状与全面自动调优，以及静态形状与快速自动调优结合也是有效的组合，但为了清晰起见，我们在此不予绘制。相反，我们选择了两个“极端”情况：首先，“较少”的每次请求优化，这意味着动态形状和在所选配置上的快速自动调优。其次是具有最多可能优化的场景，这意味着静态形状和针对每个场景的全面调优。
H100微基准测试

图3：H100上的微基准测试。延迟按批次内解码请求的比例排序。批次中包含的序列长度是可变的，中位数是最大序列长度的40%，这可能发生在真实的在线推理场景中。横轴表示批次中令牌的总数。
H100微基准测试，按最大序列长度排序

图4：H100上的微基准测试。设置和结果与上图相同，但此处延迟按批次中的最大序列长度排序，其中预填充、部分预填充和解码组合在一起。

从H100的测试结果可以看出，我们基于Helion的分页注意力实现，在解码任务中已经超越了Triton 2D注意力核心算法，并且在大型批次预填充任务中表现持平。数据显示，Helion核心算法在预填充请求中的性能是Triton的29%到137%，而在纯解码请求中则达到了132%到153%。此外，图表几乎没有显示Helion静态形状版本和非静态形状版本之间的差异。这一事实对于端到端测量至关重要，我们将在下一节进行讨论。预填充性能上的差距可以通过Helion核心算法相对于Triton核心算法更小的启动网格来解释。如前所述，Helion核心算法仅在查询头和批次维度上并行化，而不沿查询本身并行。这与Triton 2D核心算法形成对比，Triton的两个维度是批次维度（与Helion相同）以及查询头和查询令牌的混合作为第二个维度。因此，优化Helion核心算法的启动网格是另一个可行的优化方向。
MI300X微基准测试

图5：MI300X上的微基准测试。延迟按批次内解码请求的比例排序。批次中包含的序列长度是可变的，中位数是最大序列长度的40%，这可能发生在真实的在线推理场景中。横轴表示批次中令牌的总数。
MI300X微基准测试，按最大序列长度排序

图6：MI300X上的微基准测试。设置和结果与上图相同，但此处延迟按批次中的最大序列长度排序，其中预填充、部分预填充和解码组合在一起。

在MI300X平台上，结果略有不同，预填充的差距更大。在这里，Helion核心算法相对于Triton核心算法的性能在预填充请求中介于13%到75%之间，而在纯解码批次中介于58%到107%之间。

然而，即使在此处，Helion核心算法在纯解码请求中也表现出与Triton核心算法持平或超越的性能，因此我们可以认为我们的核心算法实现具备平台和性能可移植性。

vLLM中的端到端表现：模拟真实世界

作为vLLM的忠实支持者，我们自然希望评估我们的Helion分批注意力算法在真实且相关的端到端场景中的表现。因此，我们编写了一个“helion_attn”后端，类似于我们之前的Triton注意力后端，并评估了其性能。我们甚至为此提交了一份草案PR。

为了进行评估，我们使用了vLLM内置的基于流行ShareGPT基准的推理服务基准测试脚本，并禁用了前缀缓存功能。

VLLM_ATTENTION_BACKEND=EXPERIMENTAL_HELION_ATTN vllm serve \
    meta-llama/llama3.1-8b-instruct/ \
    --disable-log-requests --no-enable-prefix-caching
vllm bench serve --model meta-llama/llama3.1-8b-instruct/ \
    --dataset-name sharegpt \
    --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
    --ignore_eos

这个设置通过客户端发送请求给服务器的方式，模拟真实用户体验，从而评估vLLM推理服务器的“端到端”性能。这意味着客户端并不知道推理服务器的内部状态，例如当前运行了多少请求，或者为了获得良好核心算法性能，请求应该具有何种“形状”。在这些实验所用的特定基准测试中，客户端一次性发送1000个请求，vLLM服务器需要尽快处理这些请求，其中包括调度非常大的批次和特别是纯解码批次。

我们将实验性的Helion注意力后端与当前vLLM中的Triton后端进行了比较。两个后端都为混合批次和纯解码批次使用了完整的CUDA/HIP图。一个不同之处是，vLLM中当前的Triton注意力后端不进行“实时”调优，而是通过if-else语句在四种不同配置之间进行选择。这与我们Helion注意力后端的概念验证有所不同，后者在运行时使用自动调优器从6或7种配置中进行选择，具体取决于平台。为了公平起见，我们总是进行两次预热基准测试运行，以允许自动调优器运行，并让Helion和Triton的JIT编译器编译大部分相关的核心算法版本。每张图显示了三个结果：作为基线的Triton注意力后端在vLLM中的性能，具有静态形状的Helion核心算法的性能，以及具有动态形状的Helion核心算法的性能。在这些实验中，我们也对结果进行了Triton结果的归一化。
H100端到端性能测量

图7：在H100上，使用vLLM基准测试Llama3.1 8B和ShareGPT数据集进行的端到端性能测量。

从图表中可以看出，采用静态形状的Helion仅达到了Triton总吞吐量的约26%，而采用动态形状的Helion则实现了总令牌吞吐量的96%，并且在TTFT（Time-to-first-token，即预填充时间）方面也表现出色，与Triton持平，ITL（Inter-token-latency，即解码一个令牌所需的时间）也非常接近。这项实验凸显了推理服务器的一个重要现实：请求的形状多样、数量庞大且无法预知。此外，调度批次的形状还取决于其他多种因素，例如请求到达的顺序。因此，即使进行了相同的预热基准测试，使用静态形状的Helion几乎每次请求都会触发重新编译，因为查询张量的形状很少能完全相同。由于这是一个端到端实验，这些编译时间会反映在测量的延迟和吞吐量中。这种评估核心算法实现性能的方式与微基准测试中考察原始核心算法性能不同，但它反映了vLLM用户将体验到的真实世界性能。因此，由于其庞大的JIT开销抵消了核心算法运行时（微小的）性能提升，使用静态形状的Helion表现不佳。请注意，由于在“实时”自动调优过程中生成的Triton代码出现不可恢复的崩溃，我们不得不为静态形状的端到端实验禁用自动调优，并使用了在微基准测试中确定的最佳解码性能配置。实验的这一限制可能解释了静态形状和动态形状之间TTFT差距的一小部分，但无法解释ITL的差距以及吞吐量的巨大差异。静态形状默认启用，允许Helion通过在生成的Triton代码中使用硬编码的张量形状来优化性能。静态形状通常在Helion中被提及为性能优化，但对于像vLLM这样高度动态的使用场景，它们并非如此。更令人惊讶的是微基准测试的相应结果：即使只看纯粹的核心算法性能，优化静态形状与不优化静态形状之间几乎没有区别。我们怀疑，分批注意力核心算法的输入实际上受形状约束很大，这极大地促成了Helion中静态和动态形状编译之间性能差异极小的事实。例如，矩阵乘法的大小总是需要与vLLM的KV缓存页面大小（或块大小）对齐，并且循环融合等编译时优化无法改变这一点。

这些端到端结果中另一个意料之外的发现是，在这个特定基准测试中，Helion的TTFT实际上与Triton持平，因为这里的预填充批次比我们的微基准测试设置更大且更均匀。
MI300X端到端性能测量

图8：在MI300X上，使用vLLM基准测试Llama3.1 8B和ShareGPT数据集进行的端到端性能测量。

在MI300X平台上，我们的后端同样能正常运行，但使用动态形状时，总令牌吞吐量仅达到59%。这个结果并不令人意外，因为我们的微基准测试已经显示，在MI300X上，Helion核心算法和Triton核心算法在预填充方面的差距更大。

经验与展望：Helion的未来之路

这次实验经历非常愉快，我们认为Helion既是一种极其便捷的语言，也是一个功能强大的自动调优器。总的来说，我们在这个博客文章中描述的实验只花费了不到三周的时间，而所取得的令人印象深刻的结果，着实让我们感到惊喜。当然，为了进一步优化Helion中分批注意力的实现，仍然有许多途径可以探索，例如平衡长预填充、长解码或超大批次等场景。这可能需要Helion中实现不同的启动网格或沿查询进行并行化，并需要进一步研究以确定在不同核心算法版本之间选择的最佳启发式方法，类似于我们在Triton注意力后端中的实现方式。这种缺失的精细化优化也解释了在报告的实验中Triton和Helion实现在性能上的差距。在最好的情况下，我们可以让Helion自动调优器自动完成这种平衡。由于这些权衡都与平台相关，我们认为Helion的自动调优器非常适合可靠且快速地实现这一自动化过程。然而，在这种背景下，我们也需要找到一个好的平衡点，即在重新触发Helion调优和JIT与通过启发式方法选择配置以实现低延迟执行已编译的核心算法之间进行权衡。

我们实验性的vLLM后端另一个可能的优化是，在CUDA/HIP图捕获期间使用静态形状，因为在那里额外的JIT开销并不重要，并且记录的CUDA/HIP图的形状是静态的。因此，在这里可以让编译器更积极地考虑形状进行优化，是安全可行的。然而，之后在运行时我们必须切换回动态形状。

在实验过程中，我们意识到，对于Helion核心算法的开发，我们也需要一套全面、自动化且可靠的微基准测试套件，以便在大量使用场景中理解核心算法的详细性能。这与我们学习开发Triton核心算法的方式相似，因此，我们改编了最初为Triton工作而构建的微基准测试套件。

Helion中最有用的“命令”被证明是tensor.view，它能帮助我们及早了解Helion编译器是否认为张量的形状与我们预期的一致。这大大简化了只打印符号形状的编译器错误的调试过程。最后，我们希望能在Helion中添加一些预训练的启发式方法或决策树，以便在耗时数小时的自动调优和在vLLM等低延迟场景中所有情况只使用一种配置之间找到一个中间地带。

新媒网跨境认为，总而言之，Helion是PyTorch生态系统中一个令人兴奋且非常有用的新成员，我们非常期待它将如何影响vLLM的未来发展，并期待它能为AI推理效率带来更多创新与突破。

新媒网（公号: 新媒网跨境发布），是一个专业的跨境电商、游戏、支付、贸易和广告社区平台，为百万跨境人传递最新的海外淘金精准资讯情报。

本文来源：新媒网 https://nmedialink.com/posts/helion-153-ai-speed-code-half-industry-shock.html