MPK（Mirage Persistent Kernel）源码笔记（3）--- 系统接口

MPK（Mirage Persistent Kernel）源码笔记（3）--- 系统接口


目录

• MPK（Mirage Persistent Kernel）源码笔记（3）--- 系统接口
  
  
  
  • 0x00 概述
    
  • 0x01 流程
    
  • 0x02 初始化
    
  • 0x03 定义计算图
    
  • 0x04 编译
    
  • 0x05 执行
    
  • 0xFF 参考
    
  
  

0x00 概述

因为转译系统需要通过persistent_kernel.py来完成，所以我们先介绍persistent_kernel.py。
persistent_kernel.py是 Persistent Kernel的Python接口，本质是Python到CUDA持久化内核系统的桥梁，允许用户用python定义复杂的计算图，然后在GPU上高效执行。主要功能包括：

• 持久化内核管理。提供了 PersistentKernel 作为接口类来管理和执行持久化CUDA内核。
  
• 内核编译。将Python定义的计算图编译为CUDA代码并生成共享库。集成了nvcc编译器来编译生成CUDA代码。
  
• 内核执行。提供接口来初始化、启动和执行持久化内核。
  

此外，在 HARD_CODE 定义的C函数是底层入口点，具体如下：

• init_func：初始化内核。
  
• launch_func：启动内核执行。会调用到 launch_persistent_kernel。
  
• finalize_func：清理和终止内核。
  

0x01 流程

persistent_kernel.py的工作流程如下：

• 初始化：创建 PersistentKernel  类。
  
• 定义计算图：使用各种layer方法（如embed_layer、attention_layer等）定义计算图。
  
• 编译。调用compile()方法生成和编译CUDA内核。
  
  
  
  • 生成任务图。
    
  • 创建CUDA代码。
    
  • 调用nvcc编译器。
    
  • 创建Python绑定模块
    
  
  
• 执行：调用call()方法启动内核执行。 self.launch_func()
  
• 清理：调用finalize()方法或者自动析构。
  

具体如下图所示。

0x02 初始化

初始化函数会创建 PersistentKernel  类。
因为此处只是系统接口，大部分有意义的工作在C++代码中实现，因此此处略过。

1. class PersistentKernel:
2.     def __init__(
3.         self,
4.         world_size: int,
5.         mpi_rank: int,
6.         num_workers: int,
7.         num_local_schedulers: int,
8.         num_remote_schedulers: int,
9.         max_seq_length: int,
10.         eos_token_id: int64,
11.         meta_tensors: list[torch.Tensor],
12.         profiler_tensor: torch.Tensor,
13.         spec_decode_config: SpecDecodeConfig
14.     ):
15.         self.__finalized__ = False
16.         self._is_compiled = False
17.         self.world_size = world_size
18.         self.mpi_rank = mpi_rank
19.         self.num_workers = num_workers
20.         self.num_local_schedulers = num_local_schedulers
21.         self.num_remote_schedulers = num_remote_schedulers
22.         self.max_seq_length = max_seq_length
23.         self.eos_token_id = eos_token_id
24.         self.kn_graph = KNGraph(CyKNGraph(disable_fingerprint=True))
25.         self.meta_tensors = meta_tensors
26.         self.profiler_tensor = profiler_tensor
27.         self.use_nvshmem = True if world_size > 1 else False
28.         self.spec_decode_config = spec_decode_config
29.         self._spec_decode_handlers = {
30.             "promptlookup": self.prompt_lookup_spec_handler,
31.         }
32.         self._spec_verify_handlers = {
33.             "promptlookup": self.prompt_lookup_verify_handler,
34.         }

复制代码

0x03 定义计算图

persistent_kernel.py 使用各种layer方法（如embed_layer、attention_layer等）定义计算图。简易流程如下：

对应的代码举例如下：

1.     def attach_input(self, torch_tensor: torch.Tensor, name: str = None) -> DTensor:
2.         """
3.         将PyTorch张量附加到计算图，创建对应的DTensor（分布式张量）。
4.         
5.         参数：
6.             torch_tensor: 待附加的PyTorch张量
7.             name: 张量名称（必须指定）
8.         返回：
9.             与输入张量关联的DTensor实例
10.         说明：
11.             仅支持行优先（row-major）内存布局，通过步长校验确保布局正确性
12.         """
13.         # 提取张量维度与步长信息
14.         dims = tuple([d for d in torch_tensor.shape])
15.         strides = tuple([s for s in torch_tensor.stride()])
16.         
17.         # 校验是否为行优先布局（高维步长 = 低维步长 × 低维尺寸）
18.         for d in range(len(dims) - 1):
19.             assert strides[d] == strides[d + 1] * dims[d + 1]
20.         
21.         # 转换PyTorch数据类型为框架内部 dtype
22.         dtype = convert_torch_type_to_dtype(torch_tensor.dtype)
23.         
24.         # 创建输入张量节点
25.         t = self.kn_graph.new_input(dims=dims, strides=strides, dtype=dtype)
26.         
27.         # 断言名称非空（当前实现限制）
28.         assert name is not None
29.         
30.         # 将DTensor与PyTorch张量绑定，并注册到计算图
31.         self.kn_graph.attach_torch_tensor(t, torch_tensor, name)
32.         return t
34.     def new_tensor(
35.         self,
36.         dims: tuple,
37.         strides: tuple = None,
38.         dtype: dtype = bfloat16,
39.         name: str = None,
40.         io_category: str = "cuda_tensor",
41.     ) -> DTensor:
42.         """
43.         创建新的DTensor并根据IO类别附加到计算图。
44.         
45.         参数：
46.             dims: 张量维度元组
47.             strides: 步长元组（默认自动按行优先计算）
48.             dtype: 数据类型（默认bfloat16）
49.             name: 张量名称（必须指定）
50.             io_category: IO类别（"cuda_tensor"或"nvshmem_tensor"）
51.         返回：
52.             新创建的DTensor实例
53.         说明：
54.             支持CUDA本地张量与NVSHMEM分布式张量两种类型
55.         """
56.         # 若指定步长，校验是否为行优先布局
57.         if strides is not None:
58.             for d in range(len(dims) - 1):
59.                 assert strides[d] == strides[d + 1] * dims[d + 1]
60.         
61.         # 创建张量节点
62.         t = self.kn_graph.new_input(dims=dims, strides=strides, dtype=dtype)
63.         
64.         # 断言名称非空（当前实现限制）
65.         assert name is not None
66.         
67.         # 根据IO类别绑定张量到计算图
68.         if io_category == "cuda_tensor":
69.             self.kn_graph.attach_cuda_tensor(t, name)  # 绑定CUDA张量
70.         elif io_category == "nvshmem_tensor":
71.             self.kn_graph.attach_nvshmem_tensor(t, name)  # 绑定NVSHMEM分布式张量
72.         else:
73.             raise RuntimeError(f"Invalid io_category: {io_category}")
74.         return t
76.     def fuse_tensors(
77.         self, inputs: list[DTensor], fused_dim: int, num_groups: int, name: str = None
78.     ) -> DTensor:
79.         """
80.         融合多个张量到单个张量（当前仅支持第0维融合）。
81.         
82.         参数：
83.             inputs: 待融合的DTensor列表
84.             fused_dim: 融合维度（必须为0）
85.             num_groups: 分组数量
86.             name: 融合后张量名称
87.         返回：
88.             融合后的DTensor实例
89.         """
90.         # 当前仅支持第0维融合
91.         assert fused_dim == 0
92.         
93.         # 调用计算图的张量融合接口
94.         t = self.kn_graph.fuse_tensors(inputs, fused_dim, num_groups, name)
95.         return t
97.     def embed_layer(
98.         self,
99.         input: DTensor,  # 输入张量 [batch_size, num_spec_tokens]
100.         weight: DTensor,  # 嵌入权重 [vocab_size, hidden_size]
101.         output: DTensor,  # 输出张量 [batch_size, hidden_size]
102.         grid_dim: tuple,  # CUDA网格维度
103.         block_dim: tuple,  # CUDA块维度
104.         input_source: int = 0,  # 输入源类型（0: 全 tokens, 1: 输入 token）
105.     ):
106.         """
107.         定义嵌入层计算，将输入张量通过嵌入权重映射到隐藏空间。
108.         
109.         参数：
110.             input: 输入张量
111.             weight: 嵌入权重张量
112.             output: 输出张量（用于存储结果）
113.             grid_dim: CUDA kernel的网格维度
114.             block_dim: CUDA kernel的块维度
115.             input_source: 输入源类型标记
116.         说明：
117.             内部创建线程块图（TBGraph），定义输入输出映射关系，并注册为"embedding"任务
118.         """
119.         # 创建线程块图（CyTBGraph为底层实现，64为共享内存大小）
120.         tb_graph = TBGraph(CyTBGraph(grid_dim, block_dim, 1, 64))
121.         
122.         # 定义输入输出张量的维度映射规则
123.         tb_graph.new_input(input, (-1, 1, -1), -1, True)  # 输入张量维度映射
124.         tb_graph.new_input(weight, (1, -1, -1), -1, True)  # 权重张量维度映射
125.         tb_graph.new_input(output, (1, 0, -1), -1, True)  # 输出张量维度映射
126.         
127.         # 将张量与线程块图关联
128.         self.kn_graph.customized([input, weight, output], tb_graph)
129.         
130.         # 注册嵌入层任务，附加输入源参数
131.         self.kn_graph.register_task(tb_graph, "embedding", [input_source])
133.     def rmsnorm_linear_layer(
134.         self,
135.         input: DTensor,  # 输入张量
136.         weight_norm: DTensor,  # 归一化权重
137.         weight_linear: DTensor,  # 线性层权重
138.         output: DTensor,  # 输出张量
139.         grid_dim: tuple,  # CUDA网格维度
140.         block_dim: tuple,  # CUDA块维度
141.     ):
142.         """
143.         定义RMS归一化+线性变换组合层。
144.         
145.         参数：
146.             input: 输入张量（2D）
147.             weight_norm: RMS归一化权重（2D）
148.             weight_linear: 线性层权重（2D）
149.             output: 输出张量（2D）
150.             grid_dim: CUDA kernel的网格维度
151.             block_dim: CUDA kernel的块维度
152.         说明：
153.             先对输入执行RMS归一化，再通过线性层变换，输出结果存储到output
154.         """
155.         # 校验输入张量维度（当前仅支持2D张量）
156.         assert input.num_dims == 2
157.         assert weight_linear.num_dims == 2
158.         assert output.num_dims == 2
159.         
160.         # 创建线程块图
161.         tb_graph = TBGraph(CyTBGraph(grid_dim, block_dim, 1, 64))
162.         
163.         # 定义输入输出维度映射
164.         tb_graph.new_input(input, (-1, -1, -1), 1, True)  # 输入张量
165.         tb_graph.new_input(weight_norm, (-1, -1, -1), 0, True)  # 归一化权重
166.         tb_graph.new_input(weight_linear, (0, -1, -1), 1, True)  # 线性层权重
167.         tb_graph.new_input(output, (1, -1, -1), -1, True)  # 输出张量
168.         
169.         # 关联张量与线程块图
170.         self.kn_graph.customized([input, weight_norm, weight_linear, output], tb_graph)
171.         
172.         # 注册RMS归一化+线性层任务
173.         self.kn_graph.register_task(tb_graph, "rmsnorm_linear")
175.     def attention_layer(
176.         self,
177.         input: DTensor,  # 输入张量 (batch_size, fused_outdim / world_size)
178.         k_cache: DTensor,  # K缓存 (batch_size, seq_len, kv_heads, head_dim)
179.         v_cache: DTensor,  # V缓存 (batch_size, seq_len, kv_heads, head_dim)
180.         q_norm: DTensor,  # Q归一化权重 (可选)
181.         k_norm: DTensor,  # K归一化权重 (可选)
182.         cos_pos_embed: DTensor,  # 余弦位置编码 (可选)
183.         sin_pos_embed: DTensor,  # 正弦位置编码 (可选)
184.         output: DTensor,  # 输出张量 (batch_size, hidden_size / world_size)
185.         grid_dim: tuple,  # CUDA网格维度
186.         block_dim: tuple,  # CUDA块维度
187.     ):
188.         """
189.         定义注意力层计算，支持 rotary 位置编码与 Q/K 归一化。
190.         
191.         参数：
192.             input: 输入张量（2D）
193.             k_cache: 键缓存张量（4D）
194.             v_cache: 值缓存张量（4D）
195.             q_norm: Q归一化权重（可选，1D）
196.             k_norm: K归一化权重（可选，1D）
197.             cos_pos_embed: 余弦位置编码（可选，2D）
198.             sin_pos_embed: 正弦位置编码（可选，2D）
199.             output: 输出张量（2D）
200.             grid_dim: CUDA kernel的网格维度
201.             block_dim: CUDA kernel的块维度
202.         说明：
203.             自动检测是否启用 rotary 编码与 Q/K 归一化，动态调整计算逻辑
204.         """
205.         # 校验输入输出张量维度
206.         assert input.num_dims == 2  # (batch_size, fused_outdim / world_size)
207.         assert output.num_dims == 2  # (batch_size, hidden_size / world_size)
208.         assert k_cache.num_dims == 4  # (batch_size, seq_len, kv_heads, head_dim)
209.         assert v_cache.num_dims == 4  # (batch_size, seq_len, kv_heads, head_dim)
210.         
211.         # 提取注意力头相关参数
212.         head_dim = k_cache.dim(3)  # 头维度
213.         num_kv_heads = k_cache.dim(2)  # KV头数量
214.         num_q_heads = output.dim(1) // head_dim  # Q头数量
215.         
216.         # 检测是否启用 rotary 位置编码
217.         rotary_embed = 0
218.         if cos_pos_embed is not None or sin_pos_embed is not None:
219.             assert cos_pos_embed.num_dims == 2  # (seq_len, head_dim)
220.             assert sin_pos_embed.num_dims == 2  # (seq_len, head_dim)
221.             assert cos_pos_embed.dim(1) == head_dim
222.             assert sin_pos_embed.dim(1) == head_dim
223.             rotary_embed = 1  # 标记启用rotary编码
224.         
225.         # 检测是否启用Q/K归一化
226.         qk_norm = 0
227.         if q_norm is not None or k_norm is not None:
228.             assert q_norm.num_dims == 1  # (head_dim)
229.             assert k_norm.num_dims == 1  # (head_dim)
230.             qk_norm = 1  # 标记启用Q/K归一化
231.             assert q_norm.dim(0) == head_dim
232.             assert k_norm.dim(0) == head_dim
233.         
234.         # 注意力层参数列表
235.         params = [num_q_heads, num_kv_heads, qk_norm, rotary_embed]
236.         
237.         # 创建线程块图
238.         tb_graph = TBGraph(CyTBGraph(grid_dim, block_dim, 1, 64))
239.         
240.         # 定义输入输出维度映射
241.         tb_graph.new_input(input, (0, 1, -1), -1, True)  # 输入张量
242.         tb_graph.new_input(k_cache, (0, 2, -1), 1, True)  # K缓存
243.         tb_graph.new_input(v_cache, (0, 2, -1), 1, True)  # V缓存
244.         tb_graph.new_input(q_norm, (-1, -1, -1), -1, True)  # Q归一化权重
245.         tb_graph.new_input(k_norm, (-1, -1, -1), -1, True)  # K归一化权重
246.         tb_graph.new_input(cos_pos_embed, (-1, -1, -1), -1, True)  # 余弦位置编码
247.         tb_graph.new_input(sin_pos_embed, (-1, -1, -1), -1, True)  # 正弦位置编码
248.         tb_graph.new_input(output, (0, 1, -1), -1, True)  # 输出张量
249.         
250.         # 关联所有张量与线程块图
251.         self.kn_graph.customized(
252.             [
253.                 input,
254.                 k_cache,
255.                 v_cache,
256.                 q_norm,
257.                 k_norm,
258.                 cos_pos_embed,
259.                 sin_pos_embed,
260.                 output,
261.             ],
262.             tb_graph,
263.         )
264.         
265.         # 注册注意力层任务，附加参数
266.         self.kn_graph.register_task(tb_graph, "attention", params)

复制代码

0x04 编译

persistent_kernel.py的compile 函数主要功能是将定义好的内核图（kernel graph）编译成可执行的 CUDA 代码，并生成一个 Python 共享库（.so 文件），以便在 Python 环境中调用执行，具体如下：

• 生成任务图和 CUDA 代码。
  
  
  
  • 调用 self.kn_graph.generate_task_graph 方法，基于当前定义的内核图（KNGraph）生成任务图（task graph）和对应的 CUDA 代码。这一步会根据图中的操作（如矩阵乘法、元素级运算等）生成优化后的 CUDA 实现。
    
  
  
• 准备编译环境。
  
  
  
  • 创建临时目录用于存放生成的代码文件和编译产物。
    
  • 将生成的 CUDA 代码写入 .cu 文件。
    
  • 将任务图的 JSON 表示写入文件，便于调试或后续分析。
    
  
  
• 配置编译参数
  
  
  
  • 获取 CUDA 编译器（nvcc）路径。
    
  • 确定 Python 头文件路径，以便生成的库可以与 Python 交互。
    
  • 获取 Mirage 框架的头文件和依赖库路径。
    
  • 如果使用 NVSHMEM（多 GPU 通信库），则还需要配置 NVSHMEM 和 MPI 的头文件及库路径。
    
  
  
• 执行编译
  
  
  
  • 构建完整的 nvcc 编译命令，包括源文件、包含路径、编译选项、目标架构等。
    
  • 调用 subprocess.check_call 执行编译命令，生成一个 Python 可导入的共享库（.so 文件）。
    
  
  
• 加载编译结果
  
  
  
  • 使用 importlib.util.spec_from_file_location 和 importlib.util.module_from_spec动态加载编译生成的 .so 文件作为 Python 模块。
    
  • 从加载的模块中获取初始化、执行和终结函数（init_func, launch_func, finalize_func），并保存为 PersistentKernel 对象的成员变量，供后续调用。
    
  
  

流程图如下：

具体代码如下：

1. def compile(
2.     self,
3.     **kwargs,
4. ):      
5.     # 从关键字参数中获取输出目录，默认为None
6.     output_dir = kwargs.get("output_dir", None)
7.     # 获取Mirage相关的核心路径（根目录、包含目录、依赖目录）
8.     MIRAGE_ROOT, INCLUDE_PATH, DEPS_PATH = get_key_paths()
9.     # 创建临时目录用于存放编译过程中的中间文件
10.     tempdir_obj = tempfile.TemporaryDirectory()
11.     tempdir = tempdir_obj.name
12.     # 生成任务图：根据GPU数量和当前GPU ID划分计算任务
13.     results = self.kn_graph.generate.generate_task_graph(num_gpus=self.world_size, my_gpu_id=self.mpi_rank)
15.     # 定义CUDA代码和编译产物的临时路径
16.     cuda_code_path = os.path.join(tempdir, "test.cu")  # 生成的CUDA源代码路径
17.     so_path = os.path.join(tempdir, "test.cpython-38-x86_64-linux-gnu.so")  # 编译后的的共享库路径
18.     # 定义任务图JSON文件的临时路径
19.     json_file_path = os.path.join(tempdir, "task_graph.json")
20.     # 将任务图数据写入JSON文件
21.     with open(json_file_path, "w") as f:
22.         f.write(results["json_file"])
23.     # 将生成的CUDA代码与硬编码补充内容合并后写入文件
24.     with open(cuda_code_path, "w") as f:
25.         f.write(results["cuda_code"] + HARD_CODE)
26.         
27.     # 若指定了输出目录，将生成的CUDA代码和JSON文件复制到该目录
28.     if output_dir is not None:
29.         os.makedirs(output_dir, exist_ok=True)  # 确保输出目录存在（已存在则不报错）
30.         shutil.copy(cuda_code_path, os.path.join(output_dir, "test.cu"))  # 复制CUDA代码
31.         shutil.copy(json_file_path, os.path.join(output_dir, "task_graph.json"))  # 复制任务图JSON
33.     # 检查nvcc（CUDA编译器）是否存在
34.     cc = shutil.which("nvcc")
35.     if cc is None:
36.         # 若未找到nvcc，抛出运行时错误提示用户安装CUDA
37.         raise RuntimeError(
38.             "nvcc not found. Please make sure you have installed CUDA."
39.         )
40.     # 确定Python的默认安装路径方案（适配不同Python版本的API差异）
41.     # Python 3.10及以上版本使用get_default_scheme方法
42.     if hasattr(sysconfig, "get_default_scheme"):
43.         scheme = sysconfig.get_default_scheme()
44.     else:
45.         # 旧版本Python使用内部方法_get_default_scheme
46.         scheme = sysconfig._get_default_scheme()
47.     # 修正Debian系统中的路径方案，确保与系统Python兼容
48.     if scheme == "posix_local":
49.         scheme = "posix_prefix"
50.     # 获取Python的头文件包含目录（用于编译时链接Python库）
51.     py_include_dir = sysconfig.get_paths(scheme=scheme)["include"]
53.     # 从环境变量中获取Mirage的安装路径（若已设置）
54.     if "MIRAGE_HOME" in os.environ:
55.         MIRAGE_HOME_PATH = os.environ.get("MIRAGE_HOME")
57.     # 初始化NVSHMEM和MPI相关的路径变量（用于分布式通信）
58.     NVSHMEM_INC_PATH = None  # NVSHMEM头文件目录
59.     NVSHMEM_LIB_PATH = None  # NVSHMEM库文件目录
60.     MPI_INC_PATH = None  # MPI头文件目录
61.     MPI_LIB_PATH = None  # MPI库文件目录
62.     # 若启用NVSHMEM（NVIDIA共享内存库），配置其相关路径
63.     if self.use_nvshmem:
64.         # 配置NVSHMEM头文件路径
65.         if "NVSHMEM_INC_PATH" in os.environ:
66.             # 优先使用环境变量中指定的路径
67.             NVSHMEM_INC_PATH = os.environ.get("NVSHMEM_INC_PATH")
68.             header_file_path = os.path.join(NVSHMEM_INC_PATH, "nvshmem.h")
69.         else:
70.             # 未指定则使用默认路径
71.             NVSHMEM_INC_PATH = "/usr/include/nvshmem_12/"
72.             header_file_path = os.path.join(NVSHMEM_INC_PATH, "nvshmem.h")
73.         # 配置NVSHMEM库文件路径
74.         if "NVSHMEM_LIB_PATH" in os.environ:
75.             NVSHMEM_LIB_PATH = os.environ.get("NVSHMEM_LIB_PATH")
76.             lib_file_path = os.path.join(NVSHMEM_LIB_PATH, "libnvshmem.a")
77.         else:
78.             NVSHMEM_LIB_PATH = "/usr/lib/x86_64-linux-gnu/"
79.             lib_file_path = os.path.join(NVSHMEM_LIB_PATH, "libnvshmem.a")
80.         # 配置MPI头文件路径（NVSHMEM依赖MPI）
81.         if "MPI_INC_PATH" in os.environ:
82.             MPI_INC_PATH = os.environ.get("MPI_INC_PATH")
83.             header_file_path = os.path.join(MPI_INC_PATH, "mpi.h")
84.         else:
85.             MPI_INC_PATH = "/usr/include/"
86.             header_file_path = os.path.join(MPI_INC_PATH, "mpi.h")
87.         # 配置MPI库文件路径
88.         if "MPI_LIB_PATH" in os.environ:
89.             MPI_LIB_PATH = os.environ.get("MPI_LIB_PATH")
90.             lib_file_path = os.path.join(MPI_LIB_PATH, "libmpi.so")
91.         else:
92.             NVSHMEM_LIB_PATH = "/usr/lib/"
93.             lib_file_path = os.path.join(MPI_LIB_PATH, "libmpi.so")
94.     # 获取当前GPU的计算能力（如86对应A100，75对应T4等）
95.     target_cc = (
96.         torch.cuda.get_device_properties(0).major * 10
97.         + torch.cuda.get_device_properties(0).minor
98.     )
100.     # 生成CUDA编译命令
101.     cc_cmd = get_compile_command(
102.         target_cc=target_cc,  # GPU计算能力
103.         cc=cc,  # nvcc编译器路径
104.         file_name=cuda_code_path,  # 输入的CUDA源代码
105.         py_include_dir=py_include_dir,  # Python头文件目录
106.         mirage_home_path=MIRAGE_HOME_PATH,  # Mirage根目录
107.         mirage_inc_path=INCLUDE_PATH,  # Mirage头文件目录
108.         mirage_deps_path=DEPS_PATH,  # Mirage依赖目录
109.         nvshmem_inc_path=NVSHMEM_INC_PATH,  # NVSHMEM头文件目录
110.         nvshmem_lib_path=NVSHMEM_LIB_PATH,  # NVSHMEM库目录
111.         mpi_inc_path=MPI_INC_PATH,  # MPI头文件目录
112.         mpi_lib_path=MPI_LIB_PATH,  # MPI库目录
113.         py_so_path=so_path,  # 输出的共享库路径
114.         profiling=True if self.profiler_tensor is not None else False,  # 是否启用性能分析
115.         use_nvshmem=self.use_nvshmem,  # 是否使用NVSHMEM
116.         num_workers=self.num_workers,  # 工作线程数量
117.         num_local_schedulers=self.num_local_schedulers,  # 本地调度器数量
118.         num_remote_schedulers=self.num_remote_schedulers,  # 远程调度器数量
119.     )
120.     # 执行编译命令，生成共享库
121.     subprocess.check_call(cc_cmd)
123.     # 动态导入编译生成的共享库
124.     import importlib.util
125.     # 创建模块规格：指定模块名称和共享库路径
126.     spec = importlib.util.spec_from_file_location("__mirage_launcher", so_path)
127.     # 从规格创建模块
128.     mod = importlib.util.module_from_spec(spec)
129.     # 执行模块加载
130.     spec.loader.exec_module(mod)
131.     # 绑定模块中的核心函数（初始化、启动、结束）
132.     self.init_func = getattr(mod, "init_func")
133.     self.launch_func = getattr(mod, "launch_func")
134.     self.finalize_func = getattr(mod, "finalize_func")
135.     # 打印编译完成提示
136.     print("Finished megakernel compilation...")
138.     # 收集元数据张量的内存地址指针
139.     meta_tensors_ptr = [tensor.data_ptr() for tensor in self.meta_tensors]
140.     # 获取性能分析缓冲区的内存地址（若启用性能分析）
141.     profiler_buffer_ptr = (
142.         self.profiler_tensor.data_ptr() if self.profiler_tensor is not None else 0
143.     )
144.     # 调用初始化函数，传入必要的参数
145.     self.init_func(
146.         meta_tensors_ptr,  # 元数据张量指针列表
147.         profiler_buffer_ptr,  # 性能分析缓冲区指针
148.         self.mpi_rank,  # 当前MPI进程编号
149.         self.num_workers,  # 工作线程数量
150.         self.num_local_schedulers,  # 本地调度器数量
151.         self.num_remote_schedulers,  # 远程调度器数量
152.         self.max_seq_length,  # 最大序列长度
153.         self.eos_token_id,  # 结束符token ID
154.     )
156.     # 标记编译完成状态
157.     self._is_compiled = True

复制代码

总的来说，compile 函数的作用是将用户通过 PersistentKernel API 定义的计算图转换为高度优化的 CUDA 代码，并将其编译为可在当前系统上运行的 Python 模块，从而实现高性能的 GPU 计算。编译成功后，用户可以通过调用 init_func, launch_func, finalize_func 来初始化、执行和清理内核。
0x05 执行

persistent_kernel.py调用call()方法启动内核执行。

1.     def __call__(self, **kwargs):
2.         self.launch_func()
3.         if self.profiler_tensor is not None:
4.             from .profiler_persistent import export_to_perfetto_trace
6.             export_to_perfetto_trace(
7.                 self.profiler_tensor, f"mirage_{self.mpi_rank}.perfetto-trace"
8.             )

复制代码

launch_func()函数会调用launch_persistent_kernel()来启动内核。

1. static PyObject *launch_func(PyObject *self, PyObject *args) {
2.   launch_persistent_kernel();
4.   Py_RETURN_NONE;
5. }

复制代码

0xFF 参考

如何评价CMU将LLM转化为巨型内核的Mirage Persistent Kernel(MPK)工作？
Mirage: A Multi-Level Superoptimizer for Tensor Programs 简记  尘伊光
OSDI2025论文笔记：Mirage: A Multi-Level Superoptimizer for Tensor Programs  画饼充饥
Mirage: A Compiler for High-Performance Tensor Programs on GPUs
https://mirage-project.readthedocs.io/en/latest/mugraph.html
https://mirage-project.readthedocs.io/en/latest/transpiler.html
https://zhihaojia.medium.com/compiling-llms-into-a-megakernel-a-path-to-low-latency-inference-cf7840913c17
舍弃CUDA编程！CMU等用代码将LLM编译成巨型内核，推理延迟降6.7倍  机器之心Pro

来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！

懂技术并乐意极积无私分享的人越来越少。珍惜