最近，PyTorch 团队发布了 Helion，这是一种基于 PyTorch 的新型领域专用语言，旨在简化高性能且可移植内核的开发。凭借内置的广泛自动调优功能，Helion 有望将性能可移植性的前沿推向比 Triton 更远的地方。
为了验证这一前景（并学习 Helion），我们着手挑战用 Helion 编写 AI 领域性能要求最严苛的内核之一：Paged Attention（分页注意力机制），即 vLLM 的核心。
在过去的一年中，我们为 vLLM 贡献了一个完全用 Triton 编写的性能和平台可移植的注意力后端，它没有外部依赖，可以在 NVIDIA、AMD 和 Intel GPU 上运行（观看我们的 PyTorch 大会演讲）。因此，我们将其中一个内核 (unified_attention_2d) 在 Helion 中实现，作为 vLLM 中的一个实验性后端 (PR#27293)。
vLLM、Triton 和 Helion 背景简介
vLLM 被广泛用于大模型（LLM）推理，也是 PyTorch 基金会的一部分。vLLM 在生产环境中的采用率越来越高，可以在 NVIDIA、AMD 和 Intel GPU 上运行，也可以在 Google TPU、华为昇腾 NPU、AWS Inferentia 或 IBM Spyre 等定制加速器上运行。vLLM 为几乎所有 LLM 模型提供高效、高性能的推理，这归功于其精心设计的软件架构以及与 torch.compile 的深度集成。
Triton 是一种可以使用 Python 编写的领域专用语言 (DSL)，它提供针对 AMD、Intel 和 NVIDIA GPU 的即时 (JIT) 编译。Triton 内核已证明具有最先进的性能且具备可移植性。例如，我们用 Triton 编写了分页注意力机制，同一份内核代码在 NVIDIA H100 和 AMD MI300 上均达到了最先进的性能（您可以阅读我们的 详细论文 或 相关博文）。为此，我们也以有限的方式利用了 Triton 的自动调优器。然而，Triton 中的自动调优存在严重的局限性，尽管它对性能可移植性有积极影响，但却限制了其在生产环境中的使用。因此，对于我们的 Triton 注意力后端，目前我们使用简单的 if-else 语句作为启发式方法。
此外，Triton 也是 PyTorch Inductor（torch.compile 的编译组件）的输出语言。
Helion 是另一种 DSL，于 10 月底进入测试阶段。Helion 将自己视为“平铺 PyTorch”（tiled PyTorch），其目标主要有两个：第一，将平铺（tiling）引入 PyTorch，以便使用 PyTorch API 编写平铺程序；第二，通过广泛的自动调优增强可移植性。与 Triton 相比，Helion 的自动调优器不仅拥有可用的缓存机制，还拥有更多的自由度。这种更大的自由度源于 Helion 的自动调优器不仅可以调整像 warp 数量或流水线深度这样的底层编译标志，还能改变实现的算法层面。它还具有高级搜索算法，这正是我们之前在 Triton 环境中研究过的方向。
实现细节：如何在平铺 PyTorch 中编写分页注意力机制
启动网格与并行化方法
作为起点，我们希望重新实现我们统一注意力 Triton 内核的简化版“2D”版本。它被称为“2D”是因为该内核具有二维启动网格（详细信息见此处），我们选择此版本是因为我们认为并行平铺 softmax 的实现初期过于复杂。
然而，由于 Helion 处理启动网格的方式与 Triton 不同，我们没有 1:1 地照搬 2D 方法，而是围绕“Q 块”（Q blocks）这一核心概念构建了一个 Helion 内核。这一概念如下图所示。
图 1：Helion 内核中“Q 块”的概念。
在此图中，我们看到了需要计算的一个请求的三个维度。注意力内核需要遍历所有达到 query_length（底轴）的查询 token。在我们的内核中，我们同时获取多个查询 token。此平铺大小 TILE_Q 是可调的。接下来，对于每个 token，有多个查询头和 KV 头（左轴）。我们重新实现了 QGA 优化，以便一次获取一个 KV 头对应的所有查询头。每个 KV 头的查询头数 (QpKV) 是该方向上的平铺大小，称为 TILE_M。最后，我们必须在大小为 TILE_N（对角轴）的可调块中遍历该查询的 KV 缓存，直到当前的上下文长度。在这个内部循环中，实际的注意力计算（包括矩阵乘法，即 hl.dot）正在发生，使用了在线 softmax 实现。在内核中，还有一个额外的循环包裹所有这些，以遍历批处理中的所有请求（图中未显示）。
然而，vLLM 处理的输入张量，其第一维度是序列数量和查询长度的组合（通常称为“扁平化变长”布局）。因此，vLLM 提供了一个额外的张量作为索引，用于识别哪个 token 属于哪个序列。
因此，在尝试了一些实现方案后，我们确定了上述的四循环方法。
# registering tunable block sizes q_block_size = hl.register_block_size(1, q_block_padded_size) num_pages_at_once = hl.register_block_size(1, 32) # outer loop -> becomes the launch grid 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 # loop over the query of one request for tile_q in hl.tile(query_len, block_size=q_block_size): ... # loop over KV cache 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 个 token）。
由此生成的 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）进行循环，第二个平铺中就会出现第二个请求的 token，以及第一个请求的最后三个 token！此外，下一个平铺会混合第二个和第三个请求的 token。但是，我们不能为批处理中的每个序列更改平铺大小，因为每个编译的内核只能有一个平铺大小（请记住，这是编译时常量）。一种解决方案是这里只使用大小为 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 行。点击此处查看。
对我们来说，编写 Helion 内核比编写 Triton 版本要直观得多，尽管由于 Helion 的编程模型我们需要进行算法上的调整。在 Triton 中，大量的时间和精力（以及代码行数）都花在确保所有平铺具有正确的掩码和边界上，这需要手动处理所有维度的张量跨度和偏移量。Helion 自动且正确地完成了这些工作，为开发者节省了大量精力（尤其是调试工作！）。
此外，Helion 通过自动调优处理了其他底层细节，如实际的启动网格实现、平铺大小发现或张量内存分配。
自动调优
Helion 的优势和核心功能之一是自动调优。它不仅可以在各种不同的调优旋钮中进行搜索，还能自行检测这些旋钮所有可能的有效值。用户只需要定义平铺或块大小的上限和下限。这也可以这样完成：
for tile_n in hl.tile(seqlen//page_size, block_size=None):
这与 Triton 形成了鲜明对比，在 Triton 中，用户需要列出所有可能的有效配置（通常通过大量的嵌套 num_stages=s，其中 s 在 [1, 2, 3, 4, 5, 6, 8, 42] 中）。除了 API 不够舒适外，这还存在用户容易错过在某一平台上表现出色的组合的风险。Helion 的方法解决了这个问题，它要求用户仅在符号层面定义形状，然后推导出所有可能的组合。
如果不是最外层的 hl.tile 循环（它将成为启动网格），平铺的边界可以从张量形状推导出来。
Helion 还会将自动调优期间创建的每个内核变体与默认配置进行对比，检查其正确性，并丢弃所有数值误差过大的实验。在我们实验的初期，这个基准（或被 Helion 称为“默认配置”）是自动推导为搜索空间中“中间”位置的配置。然而，这种自动发现为我们带来了一些问题，因为结果对于我们的内核是一个无效的配置，在 vLLM 中无法处理 16 的页面大小。因此，自动调优根本无法工作，我们不得不修补 Helion 并手动定义默认配置。
是的，Helion 仍处于测试阶段，处于活跃开发中，这个问题后来通过允许用户定义一个外部函数作为自动调优基准来修复： autotune_baseline_fn=callable()。这个功能解决了我们的问题，我们随后可以定义我们现有的 Triton 实现作为基准，我们知道它会给出非常好的性能和正确的结果。这极大地简化并增强了我们的自动调优过程。
我们非常欣赏的另一个功能是自动调优的“努力级别”（effort level），Helion 将其作为用户定义的设置：autotune_effort=，可以是 ‘none’, ‘quick’, 或 ‘full’。根据我们开发 vLLM Triton 注意力后端的经验，由于分页内存和由此导致的有限的有效配置数量，自动调优过程是受限的，通常进行数天的自动调优并不划算。因此，一旦此功能可用，我们将努力级别设置为 quick。但通常情况下，“快”和“慢”是相对的，Helion 自动调优器的 quick 模式仍然需要 10 个小时来为 72 个不同的场景（如批大小、序列长度、头大小等）调整我们的内核。这比“full”（或默认）设置所需的 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 个 token (TILE_Q) 和 4 页 (TILE_N，相当于 64 个 token) 的平铺大小是最佳的。它还找出如何寻址所有涉及的张量，或者循环是否应该被展平和重排序。
对所有旋钮的详细讨论超出了本文的范围。
我们调优 Triton 内核的经验告诉我们，针对广泛场景进行调优（在微基准设置中），然后仅选择少数几个配置在 vLLM 中使用，并使用决策树或其他启发式方法在它们之间进行选择，是一种很好的权衡。
然而，当前版本 Helion 的一个缺点是它期望整个内核使用一种配置，并且我们无法像在 Triton 中那样区分对预填充（prefill）批处理或解码（decode）批处理有利的配置。因此，对于本博文的实验，我们选择了 6 种配置，在 NVIDIA 上运行 vLLM 时进行“实时”调优，在 AMD GPU 上则是 7 种。
性能评估
拥有良好的开发体验当然很好，但如果性能比我们使用 Triton 所能做到的还要差，我们可能还是不会使用 Helion。因此，我们在 NVIDIA H100 和 AMD MI300X 上对我们在 Helion 中的新分页注意力机制与我们现有的 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 内核。当然，将动态形状与全面自动调优结合，或将静态形状与快速自动调优结合也是有效的组合，但为了清晰起见，我们没有在此处绘制它们。相反，我们选择了两个“极端”：第一，“每个请求较少”的优化，意味着动态形状和针对所选配置的快速自动调优；第二，拥有最多可能优化的场景，意味着静态形状和针对每种形状的全面调优。
图 3：H100 上的微基准测试。延迟按批处理中解码请求的份额排序。批处理中的序列具有可变长度，中位数为最大序列的 40%，正如在真实在线推理场景中可能出现的那样。批处理中的 token 总数标注在 x 轴上。
图 4：H100 上的微基准测试。与上图相同的设置和结果，但此处延迟是按批处理中的最大序列长度排序，预填充、部分预填充和解码混合在一起。
从 H100 可以看出，我们的 Helion 分页注意力实现已经在解码方面优于 Triton 2D 注意力内核，并且在大批量预填充方面与 Triton 持平。我们的数据表明，Helion 内核在预填充请求方面的性能比 Triton 高出 29% 到 137%，在纯解码请求方面高出 132% 到 153%。此外，图表显示具有静态形状的 Helion 变体和没有该变体的变体之间几乎没有任何差异。这一事实对于端到端测量很重要，将在下一节中讨论。
预填充中的差距可以通过 Helion 内核与 Triton 内核相比更小的启动网格来解释。如上所述，Helion 内核仅在查询头和批处理维度上进行并行化，而不是沿着查询本身。这与 Triton 2d 内核形成对比，后者的两个维度是批处理维度（与 Helion 中一样）以及作为第二个维度的查询头和查询 token 的混合。因此，优化 Helion 内核的启动网格是另一个优化方向。
图 5：MI300X 上的微基准测试。延迟按批处理中解码请求的份额排序。批处理中的序列具有可变长度，中位数为最大序列的 40%，正如在真实在线推理场景中可能出现的那样。批处理中的 token 总数标注在 x 轴上。
图 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 编译器编译大部分相关的内核版本。每个图显示三个结果：当前 vLLM 中 triton_attn 后端的性能（作为基准）、具有静态形状的 Helion 内核的性能以及具有动态形状的 Helion 内核的性能。在这些实验中，我们同样将结果标准化为 Triton 的结果。
图 7：使用 vllm bench 在 H100 上使用 Llama3.1 8B 和 ShareGPT 数据集的端到端性能测量。
如图所示，具有静态形状的 Helion 仅达到 Triton 总吞吐量的约 26%，而具有动态形状的 Helion 达到总 token 吞吐量的 96%，在 TTFT（首个 token 时间，即预填充时间）方面持平，在 ITL（Token 间延迟，即解码一个新 token 的时间）方面也非常接近。
这个实验突显了推理服务器的一个重要现实：请求形状多样、数量众多，且无法预先知晓。此外，调度的批处理形状还取决于其他各种方面，例如请求到达的顺序。因此，即使运行了相同的基准测试作为预热，使用静态形状的 Helion 也会触发几乎每个请求的重新编译，因为查询张量的形状几乎从来都不完全相同。由于这是一个端到端实验，这些编译时间反映在测量的延迟和吞吐量中。评估我们内核实现性能的这种方式不同于微基准测试中查看原始内核性能的方式，但它反映了 vLLM 用户将体验到的真实世界性能。
因此，由于巨大的 JIT 开销超过了内核运行时中的（小）性能增益，具有静态形状的 Helion 表现不佳。请注意，由于“实时”自动调优期间生成的 Triton 代码出现不可恢复的崩溃，我们在静态形状的端到端实验中不得不禁用它，并使用了微基准测试中确定的最佳解码性能配置。实验的这一局限性可能解释了 TTFT 中静态和动态形状之间差距的一小部分，但不能解释 ITL 的差距，也不能解释吞吐量的巨大差异。静态形状在 Helion 中默认启用，并允许 Helion 使用生成的 Triton 代码中的硬编码张量形状来优化性能。静态形状通常在 Helion 中被提及为性能优化，但对于像 vLLM 这样高度动态的用例来说并非如此。
更令人惊讶的是微基准测试的相应结果：即使查看纯内核性能，优化静态形状与否之间也几乎没有区别。我们怀疑分页注意力内核的输入实际上是形状受限的事实，导致 Helion 中静态和动态形状编译之间的性能差异极小。例如，矩阵乘法的大小总是需要与 vLLM 的 KV 缓存页面大小（或块大小）对齐，而像循环融合这样的编译时优化无法改变这一点。
这些端到端结果中的另一个惊喜是，Helion 的 TTFT 在这个特定基准测试中实际上与 Triton 持平，因为这里的预填充批次比我们的微基准测试设置中更大且更均匀。
图 8：使用 vllm bench 在 MI300X 上使用 Llama3.1 8B 和 ShareGPT 数据集的端到端性能测量。
在 MI300X 上，我们的后端同样有效，但使用动态形状仅达到了总 token 吞吐量的 59%。这个结果并不令人惊讶，因为我们的微基准测试已经显示，在 MI300X 上，Helion 内核和 Triton 内核在预填充方面的差距更大。
经验教训与结论
我们很享受这次实验，并认为 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。
鸣谢
这项工作得到了 IBM Research AI 平台团队的支持，特别感谢我们的同事：Thomas Parnell、Jan van Lunteren、Mudhakar Srivatsa 和 Raghu Ganti。此外，我们要感谢 Meta 的 Jason Ansel 和 Helion 团队的反馈与支持，特别是他们修复了我们报告的错误，有时甚至在 24 小时内就完成了修复。