Blankke's Blog
Search K

Appearance

问题1：算子1、3融合 ​

针对问题1之前的错误，由于理解有误，之前书写的实际上是纵向融合。题目中提到三个算子串行执行时，算子1和算子2各自都只能用到GPU的一小部分算力，但它们仍然各自占用一次kernel launch和一次完整的显存读写。 然而纵向融合只能解决各自占用一次launch的问题，并没有解决放在显存读写中的问题。

老师把问题给我改成了算子1和算子3的融合，并说这两个更便于并行。

输入 X: [M, H]
 算子1 (降维): Y = X @ A          其中 A: [H, r], Y: [M, r] 
 算子2 (升维): Z = Y @ B          其中 B: [r, H'], Z: [M, H'] 
 算子3 (主干): W = X @ C         其中 C: [H, H'], W: [M, H'] 
 最终输出: O = W + Z 即 O: [M, H']
 O=X @ C+X @ A @ B

方案1 ​

空分复用： 我在这里的理解是并不是算子融合，而是两个kernel分别launch。 根据我们的参数表格，给出一个空分复用可以算出来的SM级别控制方案呢？

示例

老师给出的例子是： 空分共享 GPU，就是两个 kernel 划分使用 GPU 的不同 SM。

假设一个 GPU 上有 80 个 SM，是否可以支持发射 160 个 block：block 1 - 40 完成计算部分 1，block 41 - 80 完成计算部分 2，block 81 - 120 完成计算部分 1，block 121 - 160 完成计算部分 2；然后每个 block 内进行 persistent thread block 编程。通过这样的形式，是否可以支持两个 block 完成相应 SM 划分调度？

Concise answer: 思路是对的，但不完全可靠。你假设的是调度器按 pid 顺序做 round-robin 填充（pid 0→SM0, pid 1→SM1, ..., pid 79→SM79, pid 80→SM0, pid 81→SM1, ...）。如果这个假设成立，那你的方案确实能让每个 SM 上同时驻留一个 Part1 block 和一个 Part2 block。问题在于：NVIDIA 不保证这个调度顺序，但在实践中它大致成立。

方案图解：

text
Launch 160 blocks, GPU 有 80 个 SM，每个 SM occupancy = 2 blocks

假设 round-robin 调度:

第一轮 (pid 0-79):
SM0  ← pid 0  (Part1)
SM1  ← pid 1  (Part1)
...
SM39 ← pid 39 (Part1)
SM40 ← pid 40 (Part2)
SM41 ← pid 41 (Part2)
...
SM79 ← pid 79 (Part2)

第二轮 (pid 80-159):
SM0  ← pid 80  (Part1)
...
SM39 ← pid 119 (Part1)
SM40 ← pid 120 (Part2)
...
SM79 ← pid 159 (Part2)
符号数值含义
M64batch size，等价于一次输入的 token / row 数
K4096输入 hidden dimension，也就是 hidden_size
N28672gateup_proj 的 fused 输出维度
r8低秩分解的 rank，不来自 config，是题目给定的降维秩

此方案可能是：

python
GPU: 80 SM

算子1：
Y = X @ A
X: [64, 4096]
A: [4096, 8]
Y: [64, 8]

采用 BLOCK_M = 16, BLOCK_N = 16，则： num_tiles_1 = ceil(64 / 16) * ceil(8 / 16) = 4


算子3：
W = X @ C
X: [64, 4096]
C: [4096, 28672]
W: [64, 28672]

采用 BLOCK_M = 16, BLOCK_N = 128，则： num_tiles_3 = ceil(64 / 16) * ceil(28672 / 128) = 896

因此，算子1最多只能同时提供 4 个 block/program，单独运行时最多只能占用约 4 个 SM，剩余 SM 空闲。为了利用这些空闲 SM，可以将算子1和算子3并发 launch，并进行 SM 级空间划分：

SM0-SM3    → kernel1，计算 Y = X @ A
SM4-SM79   → kernel3，计算 W = X @ C

这里使用Stream级调度方法，具体代码如下

python
### 方案1
空分复用：
我在这里的理解是并不是算子融合，而是两个kernel分别launch。
根据我们的参数表格，给出一个空分复用可以算出来的SM级别控制方案呢？
::: details 示例 {data-callout="example" data-callout-fold="closed"}
老师给出的例子是：
空分共享 GPU，就是两个 kernel 划分使用 GPU 的不同 SM。

假设一个 GPU 上有 80 个 SM，是否可以支持发射 160 个 block：block 1 - 40 完成计算部分 1，block 41 - 80 完成计算部分 2，block 81 - 120 完成计算部分 1，block 121 - 160 完成计算部分 2；然后每个 block 内进行 persistent thread block 编程。通过这样的形式，是否可以支持两个 block 完成相应 SM 划分调度？

**Concise answer:** 思路是对的，但**不完全可靠**。你假设的是调度器按 pid 顺序做 round-robin 填充（pid 0→SM0, pid 1→SM1, ..., pid 79→SM79, pid 80→SM0, pid 81→SM1, ...）。如果这个假设成立，那你的方案确实能让每个 SM 上同时驻留一个 Part1 block 和一个 Part2 block。问题在于：**NVIDIA 不保证这个调度顺序，但在实践中它大致成立**。

方案图解：

```text
Launch 160 blocks, GPU 有 80 个 SM，每个 SM occupancy = 2 blocks

假设 round-robin 调度:

第一轮 (pid 0-79):
SM0  ← pid 0  (Part1)
SM1  ← pid 1  (Part1)
...
SM39 ← pid 39 (Part1)
SM40 ← pid 40 (Part2)
SM41 ← pid 41 (Part2)
...
SM79 ← pid 79 (Part2)

第二轮 (pid 80-159):
SM0  ← pid 80  (Part1)
...
SM39 ← pid 119 (Part1)
SM40 ← pid 120 (Part2)
...
SM79 ← pid 159 (Part2)

:::

符号数值含义
M64batch size，等价于一次输入的 token / row 数
K4096输入 hidden dimension，也就是 hidden_size
N28672gateup_proj 的 fused 输出维度
r8低秩分解的 rank，不来自 config，是题目给定的降维秩

此方案可能是：

python
GPU: 80 SM

算子1：
Y = X @ A
X: [64, 4096]
A: [4096, 8]
Y: [64, 8]

采用 BLOCK_M = 16, BLOCK_N = 16，则： num_tiles_1 = ceil(64 / 16) * ceil(8 / 16) = 4


算子3：
W = X @ C
X: [64, 4096]
C: [4096, 28672]
W: [64, 28672]

采用 BLOCK_M = 16, BLOCK_N = 128，则： num_tiles_3 = ceil(64 / 16) * ceil(28672 / 128) = 896

因此，算子1最多只能同时提供 4 个 block/program，单独运行时最多只能占用约 4 个 SM，剩余 SM 空闲。为了利用这些空闲 SM，可以将算子1和算子3并发 launch，并进行 SM 级空间划分：

SM0-SM3    → kernel1，计算 Y = X @ A
SM4-SM79   → kernel3，计算 W = X @ C

这里使用Stream级调度方法，具体代码如下

python
current = torch.cuda.current_stream()
stream_pair = _get_stream_pair(x.device)
down_stream = stream_pair.down_stream
main_stream = stream_pair.main_stream
down_stream.wait_stream(current)
main_stream.wait_stream(current) 
with torch.cuda.stream(down_stream):
    launch_triton_matmul(x, a, y)
with torch.cuda.stream(main_stream):
    launch_triton_matmul(x, c, w)
current.wait_stream(down_stream)
current.wait_stream(main_stream)
return y, w


### 方案2
横向融合算子
因为 launch kernel 的时候，grid 数量就是我们自己指定的。例如普通 matmul：
```python
grid = (num_tiles,)matmul_kernel[grid](...)

那么 Triton 会启动：pid = 0, 1, 2, ..., num_tiles - 1如果是横向融合两个 matmul：算子1需要 num_tiles_1 个 program算子3需要 num_tiles_3 个 program那么总 program 数就是：total_tiles = num_tiles_1 + num_tiles_3,launch：

python
grid = (total_tiles,)horizontal_fused_kernel[grid](...)

然后 kernel 内：

python
pid = tl.program_id(0)
if pid < num_tiles_1:
    # 算子1
else:   
     pid3 = pid - num_tiles_1   
     # 算子3

这里 pid3 = pid - num_tiles_1 很重要。因为算子3内部也需要从 0 开始编号自己的 tile。 这种方案似乎是一起launch，那么SM上就会同时排放算子1和算子3的工作，也就是可能同一个SM上做的是不同的计算工作，所以不算严格的空分复用。

  • 可以通过Profile让他先测出在当前配置下一个SM上最多能驻留多少个block,多少个warp
  • 第1个wave就是前Sm数量乘x每个sm最大驻留的block数
  • 只有这些block是能够确定映射到哪些SM上的
  • 如果你的任务负载大于了这些block数，那最好就启动这么多数量的block作为persistent驻留
  • 每个persistent block自己去取任务，而不是交给GPU去调度了。这是后面的方案2+

方案3 ​

注意到算子1和3的左乘矩阵相同，那么似乎可以真正意义地把右边的两个矩阵拼接成一个矩阵。因为是成立的。 那么使用逻辑拼接直接在kernel外把AC两个矩阵拼起来，输出也拼起来。这样就可以使用普通的GEMM算子进行乘法了。 不过问题在于会不会block划分本来是对齐的，现在这样就没有对齐了，造成额外的开销。

这时根据我们的数据，主干算子的tile数目为896,降维算子只有4个，那么算下来尾块只有那4个。

疑问

这里我提出疑问。就像我的观察，可以把算子 1、3 的矩阵拼起来，然后计算的就是一步乘法，以及第二步就是可以把 和 的结果加起来，整个步骤合成两步一样。

为什么不能让 直接变成 呢，因为 LoRA 是微调的时候调很小的参数，这里是微调完后相当于冻结 C，把整体加回去，那么只需要做一个可以预计算的矩阵加法，然后再乘，如此显而易见的融合为什么没有被采用呢。

我上网搜索了这个思路，说这个就是 LoRA-merge，是真实存在的一个方法。

训练时原始参数 W 保持不动，更新矩阵 A/B。合并时 W-merged = W + BA 代替原有的 W。好像这就是一个简单快捷的方法。不过想到这里，我又想起了 Punica 提到的多 LoRA adapter 的场景，Y := X@W + (x1@A1@B1, x2@A2@B2, ..., xn@An@Bn) 的场景下就无法直接 merge，这个方法也就用不了了。所以我提出的方案三也用不了。

方案2+ ​

前面三个方案考虑过后认为方案2其实最符合空分复用的思想，这里我重新学习了一下，学到了persistent block的思想(参见笔记)，所以可以把方案二稍微升级一下。

主体思想 ​

是启动 NumBlocks ≈ NumSMs，例如 108 个 persistent blocks。

每个 block 根据自己的 block_id 被分到两个 worker pool：

Pool 0: main workers，负责 W = X @ C 的 896 个 tile
Pool 1: lora-down workers，负责 Y = X @ A 的 4 个 tile

每个 worker 在 kernel 内部循环取 tile：

python

 while true:
     tile_id = atomic_add(global_counter, 1)
     if tile_id >= total_tiles: break
     compute tile

更好的设计是“两阶段 worker”：

if block_id < num_down_workers:
先处理 down tiles: Y = X @ A
down tiles 做完后，转去处理 main tiles: W = X @ C
else:
直接处理 main tiles: W = X @ C

细节问题 ​

这里有一些实现细节上会出来的问题

Grouping如何保持 ​

不能将 persistent kernel 实现为 naive atomic work queue，否则原版 GEMM 中的 grouped ordering 以及由此带来的 L2 locality 可能被破坏。persistent 版应保留原有的 grouped_order(linear_tile_id) 映射函数。区别只是：原版中 linear_tile_id 直接来自 tl.program_id(0)，而 persistent 版中 linear_tile_id 由 worker 在循环中根据 chunk_id 和 CHUNK_SIZE 生成。每个 chunk 内部包含连续的 grouped tile id，因此可以在 persistent 调度下尽量保留原 GEMM 的 tile traversal locality。

同时，tile 内部的 BLOCK_M/BLOCK_N/BLOCK_K/GROUP_SIZE_M/num_warps/num_stages 仍可复用原 matmul autotune 配置。

两阶段worker如何融合 ​

这里指的是算子A计算完成后，我试想的是使空闲的block接着帮忙计算算子C的更大部分 这里要特别精确地区分两种“帮忙”。第一种是动态帮忙：谁先做完，谁去抢 main tile。这通常需要 atomic counter：

tile_id = tl.atomic_add(main_counter, 1)

但是会效率变低，因为原子操作就是变成了串行执行。

第二种就是静态预留，为降维的算子单独预留一部分tile.且预留条带不是重新做一套 group。它是在已经排好的 grouped tile 序列上，给不同 worker 分配不同的子序列。 但是可能负载不均衡，尤其是这种情况：

  • worker 0~3: 先做 down，再做自己的 main stripes
  • worker 4~107: 直接做自己的 main stripes

如果 down workers 做 down 花了比较久，那么它们的 main stripes 会延迟开始。最后可能出现：

  • worker 4~107 已经做完了
  • worker 0~3 还在补 main

这就会形成 tail，拖慢整个 kernel。

main 有 896 个 tile，假设 A100 有 108 个 SM。若启动num_workers = 108每个 worker 平均处理896 / 108 ≈ 8.3 个 main tile

如果 effective_num_down_workers = 4，那么 4 个 down worker 先做 down，再各自处理大约 8 个 main tile。

上次更新于: