Triton & 九齿 2024 冬季作业

[voltjia]

本作业的目标是让学员实现一个基于 Triton 的注意力计算内核，并确保它与 PyTorch 中的 scaled_dot_product_attention 函数的输出一致。具体来说，作业要求实现一个带有 is_causal 参数的注意力计算内核，该参数控制是否使用因果注意力。你需要通过 Triton 实现这一功能，并验证其正确性。

作业内容

1. 理解注意力机制

在注意力机制中，给定查询（Query）、键（Key）和值（Value），计算过程通常为：

1. 计算查询与键的点积得到注意力得分（Attention Scores）；
2. 根据得分计算权重，并应用于值（Value）以得到最终的注意力输出。

其中，因果注意力（Causal Attention）通过掩蔽未来时间步（确保当前位置只与之前的位置进行交互）来实现。

PyTorch 中的 scaled_dot_product_attention 实现了该机制。你将基于 Triton 实现相同的计算。

感兴趣的同学可以阅读 Attention is All You Need。

2. 实现步骤

1. 安装所需的 Python 库

pip install triton torch pytest

2. 理解函数签名

给定 attention(query, key, value, is_causal=False, scale=None) 函数的签名，你需要实现核心的注意力计算。在这里：

• query, key, value 是注意力机制中的三个张量，分别表示查询、键和值；
• is_causal 用于控制是否应用因果掩蔽；
• scale 控制是否进行缩放处理，通常在计算点积时除以键的维度的平方根。

3. 使用 Triton 进行实现

你可以参考 FlashAttention 等算法进行实现。

4. 验证结果

在 test_attention_kernel 中，我们提供了一个与 PyTorch 结果对比的测试函数 compare_results。你需要确保你的 Triton 实现与 PyTorch 的 scaled_dot_product_attention 输出一致。请使用以下命令运行测试：

pytest <path-to-the-script>

3. 提交要求

• 需要提交一个使用 Triton 实现的注意力计算内核；
• 将会使用以下测试代码验证计算内核的实现，确保输出与 PyTorch 的结果一致；
• 请勿修改测试代码及 attention 函数签名；
• 仅导入必要的包，添加必要的 Triton 计算内核函数和修改 attention 函数体以调用 Triton 内核；
• 提交代码时请确保注释清晰。

测试代码

import pytest
import torch
from torch.nn.functional import scaled_dot_product_attention


def attention(query, key, value, is_causal=False, scale=None):
    """Mock compute kernel. Replace with your actual implementation."""
    raise NotImplementedError("Replace this with your actual implementation.")


def compare_results(triton_result, torch_result, rtol=1e-2, atol=1e-2):
    """Compares the Triton kernel result with the PyTorch result."""
    if not torch.allclose(triton_result, torch_result, rtol=rtol, atol=atol):
        diff = (triton_result - torch_result).abs().max()
        raise AssertionError(f"Results do not match. Max difference: {diff}.")


@pytest.mark.parametrize(
    "batch_size, num_heads, seq_len, head_dim",
    [(2, 4, 16, 32), (4, 8, 64, 64), (1, 16, 128, 128), (2, 4, 1024, 64)],
)
@pytest.mark.parametrize("scale", [1.0, None])
@pytest.mark.parametrize("is_causal", [False, True])
def test_attention_kernel(batch_size, num_heads, seq_len, head_dim, scale, is_causal):
    torch.manual_seed(0)

    shape = (batch_size, num_heads, seq_len, head_dim)
    dtype = torch.float16
    device = "cuda"

    q = torch.randn(shape, dtype=dtype, device=device)
    k = torch.randn(shape, dtype=dtype, device=device)
    v = torch.randn(shape, dtype=dtype, device=device)

    torch_result = scaled_dot_product_attention(
        q, k, v, is_causal=is_causal, scale=scale
    )
    triton_result = attention(q, k, v, is_causal=is_causal, scale=scale)

    compare_results(triton_result, torch_result)

希望这个作业帮助大家更好地理解 Triton 与深度学习中的并行计算。

[Lfan-ke]

感觉九齿的符号定义类似于Sympy，或者说早年的静态图神经网络框架的网络定义，就是符号运算后还是一个符号，编译后有了真实数据后会带入编译后的图进行运算（因为现在才继续看九齿的课，所以没能及时在在线直播的评论区吐槽，就转Github讨论区了，若不合时宜可直接删除此评论，就是Sympy的例子更加直观，可体验，作为示范可以加深理解）

[voltjia]

对，其实九齿最早就是受到了 SymPy 和 Graphene 的启发，然后基于 Triton 开发的，所以感觉跟它们之间有相似性是非常对劲的。而且其实之后是有计划在九齿当中引入 SymPy，从而可以更加全面的对表达式进行化简的。另外就是，马上九齿也会加入一个可视化功能，也就是咱们训练营最后一节课那些图示的生成器，这样咱们在做元变换的时候，就可以更加直观地搞清楚发生了什么了。