CS 336

Basics

• 资源量的计算：两个方面：memory & compute

total_flops公式：$6 \times token_num \times param_num$ (对每一个输入的token ，前向要跑过所有的参数，每一个参数都要参与矩阵乘，每个元素都需要经过一次加法和乘法，所以前向需要2TP，反向需要对输出和权重各做一次相同的操作，所以一共需要6TP)

• Scaling Law的建立：用小模型拟合出scale分布曲线（scaling recipe），要求建立比较细致的recipe，包括BS改变的影响等；这是因为一次大规模资源量消耗太大，不可能不断通过大型实验来发现最佳超参；

Tokenization

Encode <-> Decode

• 可以通过增大词表(vocab size)的方法，来提高压缩比（string -> indices）；（因为每个token可以表示的信息可能会变得更多），这会导致：

  A. 序列长度更短（对attention友好）

  B. sparsity增大，因为有很多embedding可能都不会被学到，这不是好事

BPE_tokenizer

思路：在原始数据上训练tokenizer，得到一个贴合数据的vocabulary；最终让常见的序列可以用一个token来表示，不常见的序列用很多个token来表示；

方法：一开始将每个byte当作一个token，之后不断merge常见的相邻的tokens。

def merge(indices: list[int], pair: tuple[int, int], new_index: int) -> list[int]:  
    """Return `indices`, but with all instances of `pair` replaced with `new_index`."""
    new_indices = []  
    i = 0  
    while i < len(indices):
        if i + 1 < len(indices) and indices[i] == pair[0] and indices[i + 1] == pair[1]:
            new_indices.append(new_index)
            i += 2
        else:
            new_indices.append(indices[i])
            i += 1
    return new_indices

  
@dataclass(frozen=True)
class BPETokenizerParams:
    """All you need to specify a BPETokenizer."""
    vocab: dict[int, bytes]     # index -> bytes
    merges: dict[tuple[int, int], int]  # index1,index2 -> new_index
  
  
def train_bpe(string: str, num_merges: int) -> BPETokenizerParams:  
    Start with the list of bytes of string.
    indices = list(map(int, string.encode("utf-8")))  
    merges: dict[tuple[int, int], int] = {}  # index1, index2 => merged index
    vocab: dict[int, bytes] = {x: bytes([x]) for x in range(256)}  # index -> bytes
    for i in range(num_merges):
        # Count the number of occurrences of each pair of tokens
        counts = count_adjacent_pairs(indices)  
        # Find the most common pair
        pair = max(counts, key=counts.get)  
        # Merge that pair
        new_index = 256 + i  
        merges[pair] = new_index  
        vocab[new_index] = vocab[pair[0]] + vocab[pair[1]]  
        indices = merge(indices, pair, new_index)  
    compression_ratio = get_compression_ratio(string, indices)  
    return BPETokenizerParams(vocab=vocab, merges=merges)

def count_adjacent_pairs(indices: list[int]) -> dict[tuple[int, int], int]:
    """Return a dictionary mapping each adjacent pair of tokens in `indices` to the number of times it occurs."""
    counts = defaultdict(int)
    for index1, index2 in zip(indices, indices[1:]):
        counts[(index1, index2)] += 1
    return counts

  
class BPETokenizer(Tokenizer):
    """BPE tokenizer given a set of merges and a vocabulary."""
    def __init__(self, params: BPETokenizerParams):
        self.params = params
    def encode(self, string: str) -> list[int]:
        indices = list(map(int, string.encode("utf-8")))  
        # Note: this is a very slow implementation
        for pair, new_index in self.params.merges.items():  
            indices = merge(indices, pair, new_index)  
        return indices
    def decode(self, indices: list[int]) -> str:
        bytes_list = list(map(self.params.vocab.get, indices))  
        string = b"".join(bytes_list).decode("utf-8")  
        return string

Resource Accounting

Einops

• motivation: 因为使用普通的pytorch矩阵行列变换操作很容易出错，einops方便我们对dimension进行操作；

x = torch.ones(3, 4)
y = torch.ones(4, 3)

# --- sum ---
z = x @ y
# 等价于
z = einsum(x, y, "seq1 hidden, hidden seq2 -> seq1 seq2")

z = x @ y.transpose(-2, -1)
# 等价于 (或者手动将...写成'batch')
z = einsum(x, y, "... seq1 hidden, ... seq2 hidden -> ... seq1 seq2")

# --- reduce ---
y = x.sum(dim=-1)
# 等价于
y = reduce(x, "...hidden -> ...", "sum")

# --- rearrange ---
# (3, 8) -> (3, 2, 4)，或者反过来也可以
x = rearrange(x, "...(heads hidden1) -> ... heads hidden1", heads=2)

Flops of matmul operation

x: (B, D), w: (D, K) 每个位置包含一次加法和乘法，故flops为$2*BDK$

• B是点的数量
• (DK)是参数的数量

那么对于一次前向的矩阵乘法，flops就是$2*tokens*params$，也是前边整体flops计算公式的由来。

MFU

MFU: Model FlOPS Utilization = actual flop_per_second / promised flop_per_second

promised flop per second: 可以在设备指标中找到

一般大于等于0.5的MFU就算是很好了.

arithmetic intensity

两个组件：计算单元& memory，所以计算耗时取决于两个因素：

1. Accelerator speed (FLOP/s)
2. memory bandwith (bytes/s)

衡量程序是compute boundh还是memory bound有两种方法：

首先可以比较communication time和compute time：

• communication time: bytes / h100_bytes_per_second
• compute time: flops / h100_flop_per_second

我们假设通信和计算可以overlap，那么：

• memory bound: communication time > compute time
• compute bound: compute time > communication time

另一种等价的衡量方式：

• Accelarator intensity: h100_flop_per_second / h100_bytes_per_second
• Arithmetic intensity: flops / bytes

那么：

• memory bound: Accelerator intensity > arithmetic intensity
• compute bound: Arithmetic intensity > accelerator intensity

example: matmul

bytes: 2nn+2nn+2nn

Flops: nn(2n-1)

只要matrix足够大，就是一个compute bound的操作。

为什么推理过程是memory bound的？因为推理大多做的是matrix-vector multiplication.

    n = 1024
    x = torch.ones(n, dtype=torch.bfloat16, device=cuda_if_available())
    w = torch.ones(n, n, dtype=torch.bfloat16, device=cuda_if_available())
    y = x @ w
    bytes = (2 * n) + (2 * n * n) + (2 * n)  # Read x, read w, write y
    flops = n * (2 * n - 1)  # n dot-products
    arithmetic_intensity = flops / bytes  # ~1 

H100_accelerator_intensity » arithmetic intensity

roofline plots

在前期，如果搬运每个byte所对应的计算操作很少，那么显然大部分资源都被浪费在memory上，而非计算单元内。

Architecture

Layernorm

Pre-vs-post norm

将$x_i$->$x_{i+1}$的通路称为residual stream，目前主流的方法是右侧的pre-norm，即在mha和FFN之前做layer norm. 因为不在residual stream上，所以也被称为non-residual norm.

优势：

• 即使不经过warmup，也可以有相比post-norm更好的稳定性，更少的gradient spike现象。

现在还有一种方法是在计算之后也加上layernorm：

也被叫做double norm.

Layernorm vs. RMSNorm

观察这两者的公式区别：

RMSNorm相比layernorm，有着更少的操作（不需要计算mean）和更少的参数（没后bias term），事实上这两个diff对性能的提升至关重要：

可以看到虽然normalization和element-wise的bias term计算的flop很少，但是占用的时间却很长。

在现代的LLM transformer架构中，bias term也经常是没有的.

Activations

一个architecture设计的经验之谈：gating往往很有帮助；其实就是一个矩阵乘法。

比如在activation的设计中，从relu到**GLU的演化就是多了一个gate function:

各种GLU的不同就在于新增加的参数矩阵V的选择的不同，对于目前最通用的SwiGLU，选择的是一个$sigmoid(x)$:

这里的一个小细节是为了让总参数量和之前一样（因为增加了一个新的权重矩阵v），对dim需要进行缩减。

Serial vs. Parallel layers

GPTJ ,PaLM, GPT-NeoX等模型提出了将原本序列化运算的transformer结构改造成parallel的：

但目前不常用。

位置编码

参考资料：

• https://kazemnejad.com/blog/transformer_architecture_positional_encoding/
• https://zhuanlan.zhihu.com/p/721032991
• https://spaces.ac.cn/archives/8231
• https://zhuanlan.zhihu.com/p/642884818
• https://mp.weixin.qq.com/s/-1xVXjoM0imXMC7DKqo-Gw
• https://kexue.fm/archives/8265/comment-page-2
• https://mp.weixin.qq.com/s?__biz=MzA3MTgwODE1Ng==&mid=2247484826&idx=1&sn=8935f0bcb2e09f438cbf3ae63825d671&chksm=9f26a069a851297f568ba7cd111082e603108716928b8444a253457233f24d09d3a18447d6b9&cur_album_id=3199751010206973953&scene=189#wechat_redirect

• 为什么要有位置编码？

因为attention结构本身无法捕捉token顺序。

位置编码有以下几个要求：

1. 能够表示一个token在序列中的绝对位置；
2. 能够用绝对位置表示token间的相对位置；
3. 具有外推性，即可以表示模型在训练过程中没有见过的长度；

Sinusoidal位置编码

正余弦位置编码的思路来自于位置本身的二进制表示，提供了一种有界又连续的编码方法：

• 为什么$\omega_k = \frac{1}{10000^{2k / d}}$?

为了满足设想：相关距离越远的embedding，相关性应该越小；也即远程衰减性：

• 绝对位置编码如何表达相对位置信息？

ROPE

对于二维向量：

即在向量上乘上一个旋转矩阵，同样地对于多偶数维向量，可以将其两两分组(注意这里的$\theta$对每个d的值是不同的)，我们接下来会证明为什么这个形式是可以表达相对位置信息的；

上式中的旋转矩阵十分稀疏，为了节省算力，可以以下面的方式等效实现：

上述公式里$\theta$的取值可以复用先前正余弦位置编码的方法，同样带来一定的远程衰减性。

• ROPE是如何用绝对位置编码表示相对位置信息的？

我们考察二维下的ROPE，注意到其相当于在embedding上乘了一个旋转矩阵，设旋转矩阵为$R$，那我们尝试证明：

$$<R_aX,R_bY>=<X,R_{b-a}Y>$$

注意到旋转矩阵的性质：

1. $$R_a^T=R_{-a}$$
2. $$R_aR_b=R_{a+b}$$

则：

$$<R_aX,R_bY>=(R_aX)^TR_bY=X^TR_a^TR_bY=X^TR_{b-a}Y=<X,R_{b-a}Y>$$

那么对于高维向量，由于内积具有线性性质，即$<a,b>=a_0b_0+a_1b_1+a_2b_2+a_3b_3+…=<a^0,b^0>+<a^1b^1>+…$，其中$a^0=[a_0,a_1]$，以此类推；所以将高维向量做两两分组并分别应用旋转矩阵后，上述在二维空间推导出的性质仍然成立。

• ROPE为何具有外推性？

本质上是因为旋转矩阵的存在，让位置编码具备了周期性和远程衰减性，这两个性质允许我们做类似线性插值（将推理时没有见过的旋转角度恢复到训练时见过的角度范围内），以及后续的优化高频信息的NTK插值等方法，通过缩小旋转弧度$m\theta_i$达到长度扩展的目的，具体参见参考文章的最后一篇内容。

Hyperparameters

• Feedforward-model dimension ratio

对transformer中的FFN层的一般形式：

$$FFN(x)=max(0,xW_1+b_1)W_2+b_2$$

一般都有$d_{ff}=4d_{model}$或者$d_ff=2.66d_{model}$.

• Head_dim * num_heads to model-dim ratio

基本会保持model dim是head dim * num_heads的整数倍，大部分是1:

• Aspect ratios

表征模型的宽度和深度的比值，主要考量在于如果模型过深，可能需要通过PP来做并行切分，对性能有影响，对效果的影响则并非主要因素。大部分模型的d_model / n_layer都在100左右：

• vocab sizes

• Dropout and other regularization

在预训练时，因为有很多数据，同时SGD只在语料库上跑一遍，所以想要overfit不太容易，所以有weight decay和dropout的必要吗？大部分现代LLM仍然会做dropout & weight decay，但其目的并非为了防止overfitting，而是在动态优化上（比如和lr decay结合给模型带来的收敛加速）上有优势：

stability issue

Softmax

• Output softmax stability - z-loss

通过增加一个z-loss的正则化项，因为我们在尝试最小化loss，所以这样可以让Z(X)贴近1，从而达到稳定Z(x)的目的.

• Attention softmax stability - QK norm

Logit soft-capping

Attention heads

除了以下几个例外，大部分模型对attention heads都不会有改动：

• GQA/MQA (Reduce attention head cost)

以上图为例，计算操作的结果源于n<d，所以在projection和attention两个矩阵乘法中，前者占了上风，如果此时的场景换成长文本，即n » d，那么结果应当为$O(bn^2d)$。

上述场景发生在训练以及推理的prefill阶段中，但是在decode阶段，假设此时也有N个query token逐次进来：

• arithmetic operations: 仍然是proj占据上风，n次的(b*1*d) @ (d*d)，所以仍然是O(bnd^2);
• total memory access: n次的(b*n*d)还有n次的对proj矩阵(d*d)的访问;

这里忽略了softmax的访存，因为比kv读取少一个n的量级，同时因为推理阶段没有backward，所以可以不把softmaxx结果写回HBM.

在decode阶段的计算强度很低，最好需要大batch+短序列，或者模型dim很大，对于小模型不太友好，

MQA正是为了解决上述痛点。

这里多出的第一项是对Q的读取，先前MHA没有列出来，是因为当时有$bn^2d$的存在。

但是MQA的问题在于因为head太少，确实会丢失expressiveness (key-query ratio)，所以变成了GQA：

在训练时repeat

• Sparse / sliding window attention

现在比较流行的做法是在full attention和sparse attention之间交替。

Long-range info via NoPE, short-range info via RoPE + SWA.

Attention Alternatives and MOE

Linear attention

Linear attention的核心思路是思考：如何将(QK^T)V变成Q(K^TV)，其好处在于将O(N^2d_k+N^2d_v)的计算复杂度降低到O(2N*d_v*d_k).

核心问题在于softmax不是一个满足结合律的操作，即做上述交换之后，效果不等价，所以当前linear attention会使用elu/silu等其他函数，具体细节可以参考网上的其他文章。

linear attention的优劣明显：

• 优势：降低计算复杂度，适合推理，原因：可以表示成类似RNN的形式：

因为推理时的token是逐个喂进来的，所以这种串行形式适合推理。

• 劣势：不适合训练，因为无法并行：由于casual mask的存在，导致对每个Q token，不能使用一致的K^TV矩阵，必须按照上边展示的那种kv逐次递增的方法来做。

但是这种方案会有效果问题，所以在实际使用中，例如Minimax M1，使用了hybrid attention的方案，即interleave full attention和linear attention，根据研究表明，两者比值并非线性关系，但有一些证据表明在较低的linear/ratio比值下，模型效果较好。

Lightening attention

在linear attention的基础上，产生了lightening attention。

其在linear attention的基础之上，融合了flash attention的想法，即将完整的token序列切成多个block，分段计算attention。

对于每一个需要依赖先前pos 0-m (属于block1)的pos m+t (属于block2)，从0-m段的attention计算时，缓存中间结果K^TV，并使用linear attention递推到第m位，这样做的好处在于对于长序列场景，前边的所有位置都降低到线性时间复杂度；这在方案中被称为inter block；而对于block2内部的[m+1, m+t)则仍然采用parallel形式的QK^T做计算，充分利用tensor core加速，这在方案中被称为intra block。

此外，采用了类似FA的cache策略，即做inter_ret + intra_ret的cumsum时，在SRAM中进行等。

Future Reading: https://www.zhihu.com/question/9740764576

Mamba-2

可以理解成在linear attention上加了一个gating：

作用是动态遗忘或保留历史信息，从而更有表达力。

Nemotron 3使用了该方案.

Gated delta net

在mamba-2的基础上衍生而来，通过一个投影矩阵$k_tk^T_t$消除历史信息的影响：

至于为什么该矩阵是一个project out的作用，可以参考投影矩阵对应的介绍资料，这里不再赘述。

Qwen 3.5 / Qwen Next使用了该方案.

Sparse adaptation

典型例子：DSA：

虽然计算复杂度仍然是平方级别的，但因为有indexer的存在，导致复杂度的常数项小了很多，整体复杂度降低。

参考资料：https://zhuanlan.zhihu.com/p/1959636888123049941

MOE

老式的MOE做法是：

先计算出所有expert的routing logits，过一层softmax，把softmax的输出结果作为topk的score，再将topk的结果作为gating function，但这种方法会导致最后topk的概率和不为1，所以在之后的模型结构中，大部分改为在topk之后，只在selected experts上做softmax。

• Shared experts

对于shared expert能否提高效果，说法不一，但是将expert切的更细，即fine-grained expert。

Train MOEs

虽然sparsity带来了训练阶段的高效性，但因为gating+topk操作不可微分，这给通过正常的梯度下降更新带来了困难；所以训练MOE模型需要一些trick。

• 使用强化学习更新门控策略: 可以做，但是太复杂，不常用；
• 增加随机扰动项(stochastic perturbations)：

• 启发式平衡loss(Heuristic balancing losses)

在原版的total loss基础上，增加一个辅助loss（auxiliary loss），目的：如果一个expert获得了过多的token，那么会压制接下来token选择该expert的概率。由两部分构成：

$f_i$表示被路由到$E_i$的比例，$P_i$表示被路由到$E_i$的平均概率，那么：

$$loss=\alpha N \sum_{i=1}^N f_i P_i$$

• 需要有f，因为被路由的概率只是一个软性的指标，不代表最后dispatch的结果；
• 需要有p，因为f不可微分，需要p作为可微入口，将梯度回传router；

注意loss对$p_i(x)$的梯度为：$\frac{\alpha N}{T^2} \sum 1_{argmax\ p(x)=o}$，这意味着对某个expert更频繁的使用会导致梯度上升，从而对$p_i(x)$本身带来更强的压制（梯度下降更新）。

一个典型例子是switch transformer.

除了上述的per-expert balancing之外，DeepSeek V1-2还引入了per-device balancing，用来平衡不同device之间的负载：

在DeepSeek-V3中，又引入了per-expert biases, 也被称为auxiliary loss free balancing（其实并不能完全做到loss free）：

在打分结果上加上一个bias，对于接受token数量超过平均值的expert，降低bias，从而达到负载均衡的效果。

MLA (Multi-Head Latent Attention)

在DS V3中，使用了MLA来做KV状态的压缩。

• 先前的问题是什么？

虽然KV cache这种用空间换时间（存储换计算）的方法，将计算复杂度从$O(N^2)$降低到了$O(N)$，但存在以下问题：

1. kv cache的显存大小成为decoding瓶颈；
2. 计算量的下降，让decoding过程成为memory bound；
3. 为了提高计算强度，BS的增大又受到了1的制约；

所以思考，能否存在一种折中方案？即沿用先前空间换时间的优化思路，但是不要那么激进？

一个常用的改变计算强度的优化方法就是利用矩阵结合律，和linear attention类似：

这种方法也被叫做矩阵吸收。经过改造后，原有的KV cache也被替代为：缓存前置的prefill的$N\times d$输入，即$X$.

但定量分析后会发现，减少的KV cache比例远低于增加的计算量。所以需要一些技巧来进一步压缩存储：即降维X：

这样节省的KV cache比例进一步增大，可以平衡带来的计算开销。$X^TW_{DKV}$就是论文中的$C$，也即需要缓存的压缩表示。

矩阵吸收的另一个潜在好处是，假使前提是要缓存$C$，矩阵吸收增加的计算复杂度（相比缓存正常的KV cache）相比于计算$W_k^{’}$和$W_v^{’}$来恢复原本的$K$ $V$矩阵增加的计算复杂度更低。

• 为什么training和prefill阶段不需要做矩阵吸收？

• ROPE是如何与MLA共存的？

首先明确，为什么ROPE不能直接应用在MLA上？因为ROPE矩阵是一个和位置相关的矩阵，不是固定的，导致每次新的token都要重新计算，从而降低推理效率。

解决方法：

具体参考：

• https://zhuanlan.zhihu.com/p/1911795330434986569
• https://zhuanlan.zhihu.com/p/16730036197

MOE stability

other train methods

• Upcycling

使用预先训练好的dense模型，load到MOE模型上：

GPUs TPUs

TPU

TPU与GPU在很多设计上相似，核心区别在于其处理矩阵乘的单元是一个大单元，而GPU是很多个小的tensor core单元来加速matmul的；同时两者对tensor core的定义不一样。

Making GPUs go fast

这里主要讨论如何优化memory pass。

• Control divergence (not a memory issur)

在同一时刻，同一warp中的所有线程处于同一代码段，如果不需要执行对应分支，则等待：

• Low precision computation

可以使用16bit(BF16/FP16)的操作：matrix ops、大部分pointwise操作（relu/add/sub/mul）；

需要更高精度(FP32/FP16)的操作：reduction (sum/softmax/norm)，因为较小的值累加很容易出现rounding errors；

需要使用更大range(FP32/BF16)的操作：返回结果比输入大很多的pointwise ops (exp, log, pow)，比如loss function；

FP8 training:

因为transpose会改变数据排布，需要重新计算scaling，所以MXFP8在内部quantize时，会一次性得到两个矩阵，其中一个用于transpose。

FP4省略。

• Operator fusion

• recomputation

• Memory coalescing and DRAM

• tiling

wave quantization

波量化（Wave Quantization）：当计算任务超出GPU SM数量时，需要将计算任务分成多个waves进行执行，而这些wave被线性执行需要等待，导致性能下降。

Flash Attention

理解从3-pass softmax -> online softmax -> 1-pass flash attention的数学推导：https://zhuanlan.zhihu.com/p/668888063

Kernel, Triton, XLA

kernel basic concepts

一些基本概念。

occupancy

• Each thread can use between 0 and 255 registers.
• The more registers threads use, the fewer threads can be scheduled on an SM (low occupancy).
• Low occupancy isn’t necessarily bad if each thread is doing more work.
• Example: thread coarsening (each thread processes multiple elements).
• Example: thread block has 64 threads, each using 160 registers, SM has 65536 registers

# What we want to run
num_threads_per_block = 128
num_registers_per_thread = 160
# What hardware offers
max_registers = 65536  # Registers allowed per SM
max_warps = 64         # Concurrent warps allowed per SM
# What we can run at once
assert num_registers_per_thread <= 255
num_registers_per_block = num_threads_per_block * num_registers_per_thread  
num_blocks = max_registers // num_registers_per_block  # Limited by registers 
num_warps = num_blocks * num_threads_per_block / 32  
occupancy = num_warps / max_warps

Bank Conflicts

同一个warp中的每个线程对share mem访问的是同一个bank中的地址（不是完全一样的地址，否则会触发broadcast）：

Memory coalescing

针对HBM：

    
When the 32 threads in a warp access HBM, memory accesses combined into transactions of 128 bytes (cache lines).
    M00 M01 M02 M03 M04 M05 M06 M07 M08 M09 M10 M11 M12 M13 M14 M15 M16 M17 M18 M19 M20 M21 M22 M23 M24 M25 M26 M27 M28 M29 M30 M31
    M32 M33 M34 M35 M36 M37 M38 M39 M40 M41 M42 M43 M44 M45 M46 M47 M48 M49 M50 M51 M52 M53 M54 M55 M56 M57 M58 M59 M60 M61 M62 M63
    
Best case: full coalescing, all threads access the same cache line (32 threads x 4 bytes = 128 bytes).

Block occupancy

    
Thread blocks scheduled onto SMs in waves.
    
B200 has 148 SMs, if we launch 160 thread blocks, first wave has 148 blocks, second wave has 12 blocks.
    
Wave quantization problem: last wave has fewer thread blocks, leaving some SMs idle (low block occupancy).
    
Solution: make number of thread blocks divide # SMs.

Profiling

torch自带的profiler，示例：

def profile(run: Callable, num_warmups: int = 1):
    # Warmup
    for _ in range(num_warmups):
        run()
    torch.cuda.synchronize()
    # Run the code with the profiler
    with torch.profiler.profile(activities=[ProfilerActivity.CUDA],
            experimental_config=torch._C._profiler._ExperimentalConfig(verbose=True)) as prof:
        run()
        torch.cuda.synchronize()
    # Print out table
    table = prof.key_averages().table(sort_by="cuda_time_total",
                                      max_name_column_width=100,
                                      row_limit=10)
    # Append to profiles.txt
    with open("var/profiles.txt", "a") as f:
        f.write(f"Profile at {time.ctime()}:\n")
        f.write(table)
        f.write("\n\n")
    return table

Triton

指定thread block做什么。

会将数据加载到shared memory中，再写回global memory。

Element-wise example:

@triton.jit
def triton_gelu_kernel(x_ptr, y_ptr, num_elements, BLOCK_SIZE: tl.constexpr):
    # Input starts at `x_ptr`
    # Output starts at `y_ptr`
    # | T T T T T T T T | T T T T T T T T | T T T T T T T T | T T T T T T T T |
    # |    Block 0      |    Block 1      |     Block 2      |    Block 3     |
    pid = tl.program_id(axis=0)      # Identifies the block
    start = pid * BLOCK_SIZE         # Starting index of this block
    # Indices where this thread block should operate
    offsets = start + tl.arange(0, BLOCK_SIZE)
    # Don't read/write past the end of the tensor
    mask = offsets < num_elements
    # Read
    x = tl.load(x_ptr + offsets, mask=mask)
    # Approx gelu is 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
    # Compute (tl.tanh doesn't exist, use tanh(a) = (exp(2a) - 1) / (exp(2a) + 1)
    a = 0.79788456 * (x + 0.044715 * x * x * x)
    exp = tl.exp(2 * a)
    tanh = (exp - 1) / (exp + 1)
    y = 0.5 * x * (1 + tanh)
    # Store
    tl.store(y_ptr + offsets, y, mask=mask)

triton经过compile后生成PTX.

Row-wise Example:

@triton.jit
def triton_softmax_kernel(x_ptr, y_ptr, x_row_stride, y_row_stride, num_cols, BLOCK_SIZE: tl.constexpr):
    assert num_cols <= BLOCK_SIZE
    # Process each row independently
    row_idx = tl.program_id(0)
    col_offsets = tl.arange(0, BLOCK_SIZE)
    # Read from global memory
    x_start_ptr = x_ptr + row_idx * x_row_stride
    x_ptrs = x_start_ptr + col_offsets
    x_row = tl.load(x_ptrs, mask=col_offsets < num_cols, other=float("-inf"))
    # Compute
    x_row = x_row - tl.max(x_row, axis=0)
    numerator = tl.exp(x_row)
    denominator = tl.sum(numerator, axis=0)
    y_row = numerator / denominator
    # Write back to global memory
    y_start_ptr = y_ptr + row_idx * y_row_stride
    y_ptrs = y_start_ptr + col_offsets
    tl.store(y_ptrs, y_row, mask=col_offsets < num_cols)

如果需要切分tile：

@triton.jit
def row_sum_kernel(x_ptr, out_ptr, N, BLOCK_SIZE: tl.constexpr):
    row = tl.program_id(0)  # Which row are we processing?
    # Accumulator for each thread
    # One row: T1 T2 T3 T4 | T1 T2 T3 T4 | T1 T2 T3 T4 (N = 12, BLOCK_SIZE = 4)
    acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
    # Loop over tiles
    for start in range(0, N, BLOCK_SIZE):
        cols = start + tl.arange(0, BLOCK_SIZE)
        mask = cols < N
        x = tl.load(x_ptr + row * N + cols, mask=mask, other=0.0)
        acc += x
    # Final reduction from BLOCK_SIZE (all threads) to a scalar
    result = tl.sum(acc, axis=0)
    tl.store(out_ptr + row, result)

Matmul example:

@triton.jit
def matmul_relu_kernel(
    a_ptr, b_ptr, c_ptr,    # Compute c = a 
    M, N, K,                # a is M x K, b is K x N, c is M x N
    stride_am, stride_ak,   # How to navigate a
    stride_bk, stride_bn,   # How to navigate b
    stride_cm, stride_cn,   # How to navigate c
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    # We are working on the (m, n)-th tile
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    # Indices
    indices_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)  # Row indices of a [BLOCK_M]
    indices_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)  # Column indices of b [BLOCK_N]
    indices_k = tl.arange(0, BLOCK_K)                    # Row indices of a = column indices of b [BLOCK_K]
    # Initial matrix of pointers of a and b
    a_ptrs = a_ptr + indices_m[:, None] * stride_am + indices_k[None, :] * stride_ak  # [BLOCK_M, BLOCK_K]
    b_ptrs = b_ptr + indices_k[:, None] * stride_bk + indices_n[None, :] * stride_bn  # [BLOCK_K, BLOCK_N]
    acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
    # Move along row tiles of a, column tiles of b
    for k in range(0, K, BLOCK_K):
        a = tl.load(a_ptrs, mask=(indices_m[:, None] < M) & (indices_k[None, :] + k < K), other=0.0)
        b = tl.load(b_ptrs, mask=(indices_k[:, None] + k < K) & (indices_n[None, :] < N), other=0.0)
        acc += tl.dot(a, b)
        a_ptrs += BLOCK_K * stride_ak  # Advance to the next row tile of a
        b_ptrs += BLOCK_K * stride_bk  # Advance to the next column tile of b
    # Apply activation function (e.g., ReLU)
    acc = tl.maximum(acc, 0.0)
    # Write output tile
    c_ptrs = c_ptr + indices_m[:, None] * stride_cm + indices_n[None, :] * stride_cn
    tl.store(c_ptrs, acc, mask=(indices_m[:, None] < M) & (indices_n[None, :] < N))

这行代码将1d的indice_m和indice_k通过broadcast，变成2d的shape：

a_ptrs = a_ptr + indices_m[:, None] * stride_am + indices_k[None, :] * stride_ak