文章摘要
本文围绕GLM-5.2-FP8大模型decode阶段进入Sparse FlashMLA前的计算流程展开优化研究，实验基于NVIDIA B300平台，先后测试了Kernel Fusion、基于Megakernel思路的MegaRTP方案，最终以CUDA Graph+Multistream+PDL组合方案将原80us的总耗时降至36.992us，同时分析了不同方案的性能差异与成因。

针对GLM-5.2-FP8模型在decode阶段的Transformer layer优化，本文聚焦该层进入Sparse FlashMLA前的计算流程，研究范围涵盖从Add、RMSNorm、Quant到Paged Indexer Score的计算阶段，不包含后续的TopK和Sparse FlashMLA步骤。实验基于NVIDIA B300平台，一次处理两个请求，每个请求输入8个token，对应参数为M=16、KV=8192。

最初的实现包含约30个串行短kernel，总耗时约80us。通过Kernel Fusion合并相邻计算后，kernel数量降至11个，串行执行时间缩短至61.631us。但这一优化未解决GPU利用率问题：11个kernel在同一条CUDA stream上串行执行，前一个kernel未完成时，后续kernel无法启动，同一时刻仅当前kernel的CTA驻留在SM上，导致SM利用率不足。

为进一步优化，研究团队参考Megakernel思路，将各kernel拆分为tile粒度的指令，放入一个常驻kernel中执行，实现了名为MegaRTP的方案。该方案将kernel拆分为tile instruction，由常驻worker CTA按指令队列执行，扣除固定框架开销后，估算耗时降至39.601us，完整运行耗时为44.272us，其中固定执行框架开销为4.671us。

在分析收益来源后，团队采用CUDA Graph + Multistream + Programmatic Dependent Launch（PDL）重建计算图，将这套组合简称为CMP，最终实现耗时36.992us。其中，Multistream并行执行无依赖分支，将执行时间从61.631us降至44.432us；PDL让有依赖的consumer提前进入GPU，完成prologue和权重预取，进一步将耗时降至36.992us。

PDL的收益并非仅消除kernel启动气泡：单条依赖边的收益取决于consumer在GDC前可完成的有效工作。若仅让consumer提前进入GPU而无额外prologue或权重预取，收益通常不足1us；若能提前执行prologue和搬运权重，收益可达1-3us；若consumer提前驻留导致资源竞争过强，收益甚至可能回退。

Megakernel也存在类似问题：它比PDL更激进地提前调度后续任务，可能导致多个未满足数据依赖的任务提前驻留，与当前kernel争抢SM、L2和显存带宽。

Megakernel的收益可分为三类：一是无依赖任务的并行，可通过Multistream实现；二是有依赖任务的提前调度、prologue和权重预取，可通过PDL实现；三是Megakernel额外提供的tile级依赖和跨指令流水，这是其特有的收益来源。

本文的工作负载中，tile级流水的收益有限：为保持GEMM效率，blockM会在一定范围内随token数增大，导致M方向的tile数不随token数等比增长；此外，GEMM后常紧跟RMSNorm，RMSNorm需等待对应行的所有GEMM N tile写回，而非单个tile完成即可开始，进一步减少了可利用的细粒度依赖。当M=16时，关键路径上基本没有可持续推进的tile级流水，即使存在细粒度依赖，也需跨越多个wave且位于关键路径上，才可能转化为端到端收益。因此，除非工作负载如MegaMoE般能通过大量跨wave的细粒度依赖持续驱动不同硬件，否则Megakernel相比Multistream + PDL很难获得明显额外收益。

Megakernel未能超越CUDA Graph + Multistream + PDL的原因主要有三点：

1. MegaRTP当前未启用2-CTA GEMM：这类GEMM需要两个CTA成对协同执行，但MegaRTP的worker CTA从各自队列独立取任务，调度协议无法保证配对CTA同时取到对应的GEMM指令，因此只能放弃这部分高性能配置。

2. Op间依赖通过gmem counter的原子操作保证数据可见性：consumer轮询counter时，load命中L2的延迟约为300-400个cycle，未命中进入HBM后，单次访问延迟随带宽压力从约800个cycle增长到4000个cycle左右。在高带宽压力下，这种全局counter轮询的代价可能超过CUDA Graph + PDL中约300ns的kernel气泡开销。

3. Megakernel会更激进地提前调度未满足依赖的任务并预取权重，导致多个任务共同争抢带宽。此外，很多Megakernel工作集中在batch=1，并非其无法支持batch=1，而是当工作负载缺乏可重叠的细粒度依赖时，Megakernel的收益主要来自缩短kernel handoff和提前执行prologue、权重预取，这些收益的绝对时间固定（0.5-3us），在batch=1时kernel本身耗时短，固定收益占比大，更易被观测到。

Kernel Fusion后，GPU仍有超过一半的SM未被充分利用。最初的实现约有30个CUDA kernel，将相邻计算合并为11个kernel后，串行时间降至61.631us，但这11个kernel在同一条stream上串行执行，grid提交的CTA数量有限，导致SM利用率低：例如12-CTA的B1 producer最多仅能让12个SM运行目标CTA，其余SM需等待，无法提前执行后续独立分支。

通过Kineto记录的执行区间和grid CTA数，按各kernel持续时间加权，整段执行平均最多仅43.7%的SM在工作。Kernel间的bubble时间仅占1.08%，因此主要优化空间并非压缩kernel启动开销，而是调整执行顺序：11个kernel并非串行依赖链，QKV-A、Head-Gate和Indexer-K三条路径互不依赖，QKV-A的结果又分别进入主attention路径和Indexer-Q路径，将这些并行任务放在同一条stream上，导致可同时执行的工作被串行排列，因此需将调度粒度下沉到tile级。

Megakernel通过tile粒度调度打破串行执行边界：将传统kernel拆分为tile级任务，每个tile对应CTA级的工作单元，RMSNorm的tile对应一行，GEMM的tile对应(blockM, blockN)输出块。在MegaRTP中，Host将每个tile编码为一条指令，常驻kernel的worker CTA按CPU预先生成的指令队列执行，无需在runtime重新推导计算图。

Megakernel实现了三类并行：

1. 无依赖路径并行：以QKV-A、Head-Gate和Indexer-K为例，三者都依赖Add+RMSNorm+Quant，但彼此无依赖，Add完成对应tile并更新event counter后，等待结果的指令可交错推进，共同使用空闲SM，无需逐个完成整个kernel。

2. 有依赖任务提前准备：任务B依赖任务A的输出时，B在读取A的输出前，可提前执行与A无关的工作，如读取任务描述、构造调度状态、初始化barrier、准备TMA descriptor、加载自身权重和scale，仅让负责读取输入的warp等待A的结果，从而将B的prologue与A的计算重叠。

3. tile级细粒度依赖：传统kernel以整个grid完成为依赖边界，Megakernel将依赖细化到tile，A完成一个tile后更新event counter，B的指令等待该counter达到阈值后再读取并计算该tile，无需等待整个grid结束。这种依赖仅在下游能消费上游已完成的tile时才有收益，当M较小时，上游kernel在一个wave内完成，单个tile耗时接近整个kernel，下游无空闲CTA提前执行，无法缩短关键路径。

相关工作包括Stanford的Megakernels、Mirage、Event Tensor和TileRT等：

Stanford的Look Ma, No Bubbles!面向单GPU、batch size=1的Llama decode，将一轮forward的执行计划放入GPU内的轻量指令执行器，Host构造算子DAG并拆分为可由worker CTA执行的指令，分配到静态队列，Device侧启动常驻kernel，每个CTA分工为controller warp、loader warp、consumer warp、storer warp，通过instruction-arrived、stage semaphore和page semaphore协作，实现控制信息、权重搬运、Tensor Core计算和写回的并行。

Mirage Persistent Kernel（MPK）将算子拆分为SM级task，根据读写范围建立task间的event，而非使用粗粒度barrier，生成tGraph时通过Event Fusion合并相同producer或consumer的event，tGraph Normalization规整task的触发和依赖关系，tGraph Linearization重排task，优化metadata大小，实现细粒度依赖的收益。

TileRT的CUDA kernel实现闭源，其多卡分工方式为：Attention采用两种GPU角色，rank 0执行Indexer Q/K、扫描KI cache、计算score和Top-2048，发布索引给其余七张卡；rank 1-7以TP7分片持有主MLA权重，完成projection和cache准备后，根据索引gather稀疏KV，完成后续计算。MoE则采用expert内部的TP8，每个rank执行相同的router和Top-8，持有部分中间维，完成归约和residual。

MegaRTP的设计思路：

1. 模型图可扩展：同一个runtime可执行不同模型的完整图或任意子图，M和KV长度改变时，仅重新生成Host侧计划，不修改常驻kernel的执行循环。

2. 大M下的高性能GEMM：GLM5的batch或MTP增大后，需高吞吐的Blackwell GEMM，因此接入可按M特化的DeepGEMM。

Host侧将模型图拆分为tile，为每个tile生成一条指令，分配到148个worker CTA的队列；Device侧启动常驻kernel，每个worker CTA从队列读取指令，根据opcode调用对应的算子实现，输入未就绪时等待event，执行完成后更新event，通知后续指令。

Host侧先捕获OpGraph，记录算子和tensor的def-use关系，再展开为TaskGraph，根据producer和consumer的配对策略生成event。配对策略包括四种：whole（保守的fallback）、tile（region对应）、tile_cover（多块覆盖一块）、tile_reduce（多块归约到一块），每种策略对应不同的event生成方式。

KernelPlan将732条tile task按round-robin分配到148个worker queue，其中140个worker有5条有效指令，8个worker有4条，用NoOp补齐。Device侧的张量形状固定为[148,5,32]，无需维护ready queue，每条指令为固定大小的任务描述，包含opcode、wait_count、signal_count、operand_row_id等信息，通过OperandPtrTable优化地址存储，避免重复存储CUDA指针，仅在输入变化时更新该表。

Device侧启动148 CTA × 384 threads的常驻kernel，每个CTA长期循环执行指令，warp specialization分工为controller warp、loader warp、consumer warp、storer warp等，通过instruction_arrived、stage semaphore和page semaphore协作，实现指令间的重叠执行。每个CTA配置两个slot，重叠当前指令的计算和下一条指令的准备，通过共享内存分页，将SMEM切成13个16KiB的physical page，实现tile级的内存复用，避免重复加载权重。

以上内容不代表本平台立场，仅供读者参考