MPK（Mirage Persistent Kernel）源碼筆記（3）--- 系統接口

MPK（Mirage Persistent Kernel）源碼筆記（3）--- 系統接口

目錄

• MPK（Mirage Persistent Kernel）源碼筆記（3）--- 系統接口
  • 0x00 概述
  • 0x01 流程
  • 0x02 初始化
  • 0x03 定義計算圖
  • 0x04 編譯
  • 0x05 執行
  • 0xFF 參考

0x00 概述

因為轉譯系統需要通過(guò)persistent_kernel.py來(lái)完成，所以我們先介紹persistent_kernel.py。

persistent_kernel.py是 Persistent Kernel的Python接口，本質(zhì)是Python到CUDA持久化內核系統的橋梁，允許用戶(hù)用python定義復雜的計算圖，然后在GPU上高效執行。主要功能包括：

• 持久化內核管理。提供了 PersistentKernel 作為接口類(lèi)來(lái)管理和執行持久化CUDA內核。
• 內核編譯。將Python定義的計算圖編譯為CUDA代碼并生成共享庫。集成了nvcc編譯器來(lái)編譯生成CUDA代碼。
• 內核執行。提供接口來(lái)初始化、啟動(dòng)和執行持久化內核。

此外，在 HARD_CODE 定義的C函數是底層入口點(diǎn)，具體如下：

• init_func：初始化內核。
• launch_func：?jiǎn)?dòng)內核執行。會(huì )調用到 launch_persistent_kernel。
• finalize_func：清理和終止內核。

0x01 流程

persistent_kernel.py的工作流程如下：

• 初始化：創(chuàng )建 PersistentKernel 類(lèi)。
• 定義計算圖：使用各種layer方法（如embed_layer、attention_layer等）定義計算圖。
• 編譯。調用compile()方法生成和編譯CUDA內核。
  • 生成任務(wù)圖。
  • 創(chuàng )建CUDA代碼。
  • 調用nvcc編譯器。
  • 創(chuàng )建Python綁定模塊
• 執行：調用call()方法啟動(dòng)內核執行。 self.launch_func()
• 清理：調用finalize()方法或者自動(dòng)析構。

具體如下圖所示。

0x02 初始化

初始化函數會(huì )創(chuàng )建 PersistentKernel 類(lèi)。

因為此處只是系統接口，大部分有意義的工作在C++代碼中實(shí)現，因此此處略過(guò)。

class PersistentKernel:def __init__(self,world_size: int,mpi_rank: int,num_workers: int,num_local_schedulers: int,num_remote_schedulers: int,max_seq_length: int,eos_token_id: int64,meta_tensors: list[torch.Tensor],profiler_tensor: torch.Tensor,spec_decode_config: SpecDecodeConfig):self.__finalized__ = Falseself._is_compiled = Falseself.world_size = world_sizeself.mpi_rank = mpi_rankself.num_workers = num_workersself.num_local_schedulers = num_local_schedulersself.num_remote_schedulers = num_remote_schedulersself.max_seq_length = max_seq_lengthself.eos_token_id = eos_token_idself.kn_graph = KNGraph(CyKNGraph(disable_fingerprint=True))self.meta_tensors = meta_tensorsself.profiler_tensor = profiler_tensorself.use_nvshmem = True if world_size > 1 else Falseself.spec_decode_config = spec_decode_configself._spec_decode_handlers = {"promptlookup": self.prompt_lookup_spec_handler,}self._spec_verify_handlers = {"promptlookup": self.prompt_lookup_verify_handler,}

0x03 定義計算圖

persistent_kernel.py 使用各種layer方法（如embed_layer、attention_layer等）定義計算圖。簡(jiǎn)易流程如下：

對應的代碼舉例如下：

    def attach_input(self, torch_tensor: torch.Tensor, name: str = None) -> DTensor:"""將PyTorch張量附加到計算圖，創(chuàng  )建對應的DTensor（分布式張量）。參數：torch_tensor: 待附加的PyTorch張量name: 張量名稱(chēng)（必須指定）返回：與輸入張量關(guān)聯(lián)的DTensor實(shí)例說(shuō)明：僅支持行優(yōu)先（row-major）內存布局，通過(guò)步長(cháng)校驗確保布局正確性"""# 提取張量維度與步長(cháng)信息dims = tuple([d for d in torch_tensor.shape])strides = tuple([s for s in torch_tensor.stride()])# 校驗是否為行優(yōu)先布局（高維步長(cháng) = 低維步長(cháng) × 低維尺寸）for d in range(len(dims) - 1):assert strides[d] == strides[d + 1] * dims[d + 1]# 轉換PyTorch數據類(lèi)型為框架內部 dtypedtype = convert_torch_type_to_dtype(torch_tensor.dtype)# 創(chuàng  )建輸入張量節點(diǎn)t = self.kn_graph.new_input(dims=dims, strides=strides, dtype=dtype)# 斷言名稱(chēng)非空（當前實(shí)現限制）assert name is not None# 將DTensor與PyTorch張量綁定，并注冊到計算圖self.kn_graph.attach_torch_tensor(t, torch_tensor, name)return tdef new_tensor(self,dims: tuple,strides: tuple = None,dtype: dtype = bfloat16,name: str = None,io_category: str = "cuda_tensor",) -> DTensor:"""創(chuàng  )建新的DTensor并根據IO類(lèi)別附加到計算圖。參數：dims: 張量維度元組strides: 步長(cháng)元組（默認自動(dòng)按行優(yōu)先計算）dtype: 數據類(lèi)型（默認bfloat16）name: 張量名稱(chēng)（必須指定）io_category: IO類(lèi)別（"cuda_tensor"或"nvshmem_tensor"）返回：新創(chuàng  )建的DTensor實(shí)例說(shuō)明：支持CUDA本地張量與NVSHMEM分布式張量?jì)煞N類(lèi)型"""# 若指定步長(cháng)，校驗是否為行優(yōu)先布局if strides is not None:for d in range(len(dims) - 1):assert strides[d] == strides[d + 1] * dims[d + 1]# 創(chuàng  )建張量節點(diǎn)t = self.kn_graph.new_input(dims=dims, strides=strides, dtype=dtype)# 斷言名稱(chēng)非空（當前實(shí)現限制）assert name is not None# 根據IO類(lèi)別綁定張量到計算圖if io_category == "cuda_tensor":self.kn_graph.attach_cuda_tensor(t, name)  # 綁定CUDA張量elif io_category == "nvshmem_tensor":self.kn_graph.attach_nvshmem_tensor(t, name)  # 綁定NVSHMEM分布式張量else:raise RuntimeError(f"Invalid io_category: {io_category}")return tdef fuse_tensors(self, inputs: list[DTensor], fused_dim: int, num_groups: int, name: str = None) -> DTensor:"""融合多個(gè)張量到單個(gè)張量（當前僅支持第0維融合）。參數：inputs: 待融合的DTensor列表fused_dim: 融合維度（必須為0）num_groups: 分組數量name: 融合后張量名稱(chēng)返回：融合后的DTensor實(shí)例"""# 當前僅支持第0維融合assert fused_dim == 0# 調用計算圖的張量融合接口t = self.kn_graph.fuse_tensors(inputs, fused_dim, num_groups, name)return tdef embed_layer(self,input: DTensor,  # 輸入張量 [batch_size, num_spec_tokens]weight: DTensor,  # 嵌入權重 [vocab_size, hidden_size]output: DTensor,  # 輸出張量 [batch_size, hidden_size]grid_dim: tuple,  # CUDA網(wǎng)格維度block_dim: tuple,  # CUDA塊維度input_source: int = 0,  # 輸入源類(lèi)型（0: 全 tokens, 1: 輸入 token）):"""定義嵌入層計算，將輸入張量通過(guò)嵌入權重映射到隱藏空間。參數：input: 輸入張量weight: 嵌入權重張量output: 輸出張量（用于存儲結果）grid_dim: CUDA kernel的網(wǎng)格維度block_dim: CUDA kernel的塊維度input_source: 輸入源類(lèi)型標記說(shuō)明：內部創(chuàng  )建線(xiàn)程塊圖（TBGraph），定義輸入輸出映射關(guān)系，并注冊為"embedding"任務(wù)"""# 創(chuàng  )建線(xiàn)程塊圖（CyTBGraph為底層實(shí)現，64為共享內存大?。﹖b_graph = TBGraph(CyTBGraph(grid_dim, block_dim, 1, 64))# 定義輸入輸出張量的維度映射規則tb_graph.new_input(input, (-1, 1, -1), -1, True)  # 輸入張量維度映射tb_graph.new_input(weight, (1, -1, -1), -1, True)  # 權重張量維度映射tb_graph.new_input(output, (1, 0, -1), -1, True)  # 輸出張量維度映射# 將張量與線(xiàn)程塊圖關(guān)聯(lián)self.kn_graph.customized([input, weight, output], tb_graph)# 注冊嵌入層任務(wù)，附加輸入源參數self.kn_graph.register_task(tb_graph, "embedding", [input_source])def rmsnorm_linear_layer(self,input: DTensor,  # 輸入張量weight_norm: DTensor,  # 歸一化權重weight_linear: DTensor,  # 線(xiàn)性層權重output: DTensor,  # 輸出張量grid_dim: tuple,  # CUDA網(wǎng)格維度block_dim: tuple,  # CUDA塊維度):"""定義RMS歸一化+線(xiàn)性變換組合層。參數：input: 輸入張量（2D）weight_norm: RMS歸一化權重（2D）weight_linear: 線(xiàn)性層權重（2D）output: 輸出張量（2D）grid_dim: CUDA kernel的網(wǎng)格維度block_dim: CUDA kernel的塊維度說(shuō)明：先對輸入執行RMS歸一化，再通過(guò)線(xiàn)性層變換，輸出結果存儲到output"""# 校驗輸入張量維度（當前僅支持2D張量）assert input.num_dims == 2assert weight_linear.num_dims == 2assert output.num_dims == 2# 創(chuàng  )建線(xiàn)程塊圖tb_graph = TBGraph(CyTBGraph(grid_dim, block_dim, 1, 64))# 定義輸入輸出維度映射tb_graph.new_input(input, (-1, -1, -1), 1, True)  # 輸入張量tb_graph.new_input(weight_norm, (-1, -1, -1), 0, True)  # 歸一化權重tb_graph.new_input(weight_linear, (0, -1, -1), 1, True)  # 線(xiàn)性層權重tb_graph.new_input(output, (1, -1, -1), -1, True)  # 輸出張量# 關(guān)聯(lián)張量與線(xiàn)程塊圖self.kn_graph.customized([input, weight_norm, weight_linear, output], tb_graph)# 注冊RMS歸一化+線(xiàn)性層任務(wù)self.kn_graph.register_task(tb_graph, "rmsnorm_linear")def attention_layer(self,input: DTensor,  # 輸入張量 (batch_size, fused_outdim / world_size)k_cache: DTensor,  # K緩存 (batch_size, seq_len, kv_heads, head_dim)v_cache: DTensor,  # V緩存 (batch_size, seq_len, kv_heads, head_dim)q_norm: DTensor,  # Q歸一化權重 (可選)k_norm: DTensor,  # K歸一化權重 (可選)cos_pos_embed: DTensor,  # 余弦位置編碼 (可選)sin_pos_embed: DTensor,  # 正弦位置編碼 (可選)output: DTensor,  # 輸出張量 (batch_size, hidden_size / world_size)grid_dim: tuple,  # CUDA網(wǎng)格維度block_dim: tuple,  # CUDA塊維度):"""定義注意力層計算，支持 rotary 位置編碼與 Q/K 歸一化。參數：input: 輸入張量（2D）k_cache: 鍵緩存張量（4D）v_cache: 值緩存張量（4D）q_norm: Q歸一化權重（可選，1D）k_norm: K歸一化權重（可選，1D）cos_pos_embed: 余弦位置編碼（可選，2D）sin_pos_embed: 正弦位置編碼（可選，2D）output: 輸出張量（2D）grid_dim: CUDA kernel的網(wǎng)格維度block_dim: CUDA kernel的塊維度說(shuō)明：自動(dòng)檢測是否啟用 rotary 編碼與 Q/K 歸一化，動(dòng)態(tài)調整計算邏輯"""# 校驗輸入輸出張量維度assert input.num_dims == 2  # (batch_size, fused_outdim / world_size)assert output.num_dims == 2  # (batch_size, hidden_size / world_size)assert k_cache.num_dims == 4  # (batch_size, seq_len, kv_heads, head_dim)assert v_cache.num_dims == 4  # (batch_size, seq_len, kv_heads, head_dim)# 提取注意力頭相關(guān)參數head_dim = k_cache.dim(3)  # 頭維度num_kv_heads = k_cache.dim(2)  # KV頭數量num_q_heads = output.dim(1) // head_dim  # Q頭數量# 檢測是否啟用 rotary 位置編碼rotary_embed = 0if cos_pos_embed is not None or sin_pos_embed is not None:assert cos_pos_embed.num_dims == 2  # (seq_len, head_dim)assert sin_pos_embed.num_dims == 2  # (seq_len, head_dim)assert cos_pos_embed.dim(1) == head_dimassert sin_pos_embed.dim(1) == head_dimrotary_embed = 1  # 標記啟用rotary編碼# 檢測是否啟用Q/K歸一化qk_norm = 0if q_norm is not None or k_norm is not None:assert q_norm.num_dims == 1  # (head_dim)assert k_norm.num_dims == 1  # (head_dim)qk_norm = 1  # 標記啟用Q/K歸一化assert q_norm.dim(0) == head_dimassert k_norm.dim(0) == head_dim# 注意力層參數列表params = [num_q_heads, num_kv_heads, qk_norm, rotary_embed]# 創(chuàng  )建線(xiàn)程塊圖tb_graph = TBGraph(CyTBGraph(grid_dim, block_dim, 1, 64))# 定義輸入輸出維度映射tb_graph.new_input(input, (0, 1, -1), -1, True)  # 輸入張量tb_graph.new_input(k_cache, (0, 2, -1), 1, True)  # K緩存tb_graph.new_input(v_cache, (0, 2, -1), 1, True)  # V緩存tb_graph.new_input(q_norm, (-1, -1, -1), -1, True)  # Q歸一化權重tb_graph.new_input(k_norm, (-1, -1, -1), -1, True)  # K歸一化權重tb_graph.new_input(cos_pos_embed, (-1, -1, -1), -1, True)  # 余弦位置編碼tb_graph.new_input(sin_pos_embed, (-1, -1, -1), -1, True)  # 正弦位置編碼tb_graph.new_input(output, (0, 1, -1), -1, True)  # 輸出張量# 關(guān)聯(lián)所有張量與線(xiàn)程塊圖self.kn_graph.customized([input,k_cache,v_cache,q_norm,k_norm,cos_pos_embed,sin_pos_embed,output,],tb_graph,)# 注冊注意力層任務(wù)，附加參數self.kn_graph.register_task(tb_graph, "attention", params)

0x04 編譯

persistent_kernel.py的compile 函數主要功能是將定義好的內核圖（kernel graph）編譯成可執行的 CUDA 代碼，并生成一個(gè) Python 共享庫（.so 文件），以便在 Python 環(huán)境中調用執行，具體如下：

• 生成任務(wù)圖和 CUDA 代碼。

  • 調用 self.kn_graph.generate_task_graph 方法，基于當前定義的內核圖（KNGraph）生成任務(wù)圖（task graph）和對應的 CUDA 代碼。這一步會(huì )根據圖中的操作（如矩陣乘法、元素級運算等）生成優(yōu)化后的 CUDA 實(shí)現。

• 準備編譯環(huán)境。

  • 創(chuàng )建臨時(shí)目錄用于存放生成的代碼文件和編譯產(chǎn)物。
  • 將生成的 CUDA 代碼寫(xiě)入 .cu 文件。
  • 將任務(wù)圖的 JSON 表示寫(xiě)入文件，便于調試或后續分析。

• 配置編譯參數

  • 獲取 CUDA 編譯器（nvcc）路徑。
  • 確定 Python 頭文件路徑，以便生成的庫可以與 Python 交互。
  • 獲取 Mirage 框架的頭文件和依賴(lài)庫路徑。
  • 如果使用 NVSHMEM（多 GPU 通信庫），則還需要配置 NVSHMEM 和 MPI 的頭文件及庫路徑。

• 執行編譯

  • 構建完整的 nvcc 編譯命令，包括源文件、包含路徑、編譯選項、目標架構等。
  • 調用 subprocess.check_call 執行編譯命令，生成一個(gè) Python 可導入的共享庫（.so 文件）。

• 加載編譯結果

  • 使用 importlib.util.spec_from_file_location 和 importlib.util.module_from_spec動(dòng)態(tài)加載編譯生成的 .so 文件作為 Python 模塊。
  • 從加載的模塊中獲取初始化、執行和終結函數（init_func, launch_func, finalize_func），并保存為 PersistentKernel 對象的成員變量，供后續調用。

流程圖如下：

具體代碼如下：

def compile(self,**kwargs,
):      # 從關(guān)鍵字參數中獲取輸出目錄，默認為Noneoutput_dir = kwargs.get("output_dir", None)# 獲取Mirage相關(guān)的核心路徑（根目錄、包含目錄、依賴(lài)目錄）MIRAGE_ROOT, INCLUDE_PATH, DEPS_PATH = get_key_paths()# 創(chuàng  )建臨時(shí)目錄用于存放編譯過(guò)程中的中間文件tempdir_obj = tempfile.TemporaryDirectory()tempdir = tempdir_obj.name# 生成任務(wù)圖：根據GPU數量和當前GPU ID劃分計算任務(wù)results = self.kn_graph.generate.generate_task_graph(num_gpus=self.world_size, my_gpu_id=self.mpi_rank)# 定義CUDA代碼和編譯產(chǎn)物的臨時(shí)路徑cuda_code_path = os.path.join(tempdir, "test.cu")  # 生成的CUDA源代碼路徑so_path = os.path.join(tempdir, "test.cpython-38-x86_64-linux-gnu.so")  # 編譯后的的共享庫路徑# 定義任務(wù)圖JSON文件的臨時(shí)路徑j(luò )son_file_path = os.path.join(tempdir, "task_graph.json")# 將任務(wù)圖數據寫(xiě)入JSON文件with open(json_file_path, "w") as f:f.write(results["json_file"])# 將生成的CUDA代碼與硬編碼補充內容合并后寫(xiě)入文件with open(cuda_code_path, "w") as f:f.write(results["cuda_code"] + HARD_CODE)# 若指定了輸出目錄，將生成的CUDA代碼和JSON文件復制到該目錄if output_dir is not None:os.makedirs(output_dir, exist_ok=True)  # 確保輸出目錄存在（已存在則不報錯）shutil.copy(cuda_code_path, os.path.join(output_dir, "test.cu"))  # 復制CUDA代碼shutil.copy(json_file_path, os.path.join(output_dir, "task_graph.json"))  # 復制任務(wù)圖JSON# 檢查nvcc（CUDA編譯器）是否存在cc = shutil.which("nvcc")if cc is None:# 若未找到nvcc，拋出運行時(shí)錯誤提示用戶(hù)安裝CUDAraise RuntimeError("nvcc not found. Please make sure you have installed CUDA.")# 確定Python的默認安裝路徑方案（適配不同Python版本的API差異）# Python 3.10及以上版本使用get_default_scheme方法if hasattr(sysconfig, "get_default_scheme"):scheme = sysconfig.get_default_scheme()else:# 舊版本Python使用內部方法_get_default_schemescheme = sysconfig._get_default_scheme()# 修正Debian系統中的路徑方案，確保與系統Python兼容if scheme == "posix_local":scheme = "posix_prefix"# 獲取Python的頭文件包含目錄（用于編譯時(shí)鏈接Python庫）py_include_dir = sysconfig.get_paths(scheme=scheme)["include"]# 從環(huán)境變量中獲取Mirage的安裝路徑（若已設置）if "MIRAGE_HOME" in os.environ:MIRAGE_HOME_PATH = os.environ.get("MIRAGE_HOME")# 初始化NVSHMEM和MPI相關(guān)的路徑變量（用于分布式通信）NVSHMEM_INC_PATH = None  # NVSHMEM頭文件目錄NVSHMEM_LIB_PATH = None  # NVSHMEM庫文件目錄MPI_INC_PATH = None  # MPI頭文件目錄MPI_LIB_PATH = None  # MPI庫文件目錄# 若啟用NVSHMEM（NVIDIA共享內存庫），配置其相關(guān)路徑if self.use_nvshmem:# 配置NVSHMEM頭文件路徑if "NVSHMEM_INC_PATH" in os.environ:# 優(yōu)先使用環(huán)境變量中指定的路徑NVSHMEM_INC_PATH = os.environ.get("NVSHMEM_INC_PATH")header_file_path = os.path.join(NVSHMEM_INC_PATH, "nvshmem.h")else:# 未指定則使用默認路徑NVSHMEM_INC_PATH = "/usr/include/nvshmem_12/"header_file_path = os.path.join(NVSHMEM_INC_PATH, "nvshmem.h")# 配置NVSHMEM庫文件路徑if "NVSHMEM_LIB_PATH" in os.environ:NVSHMEM_LIB_PATH = os.environ.get("NVSHMEM_LIB_PATH")lib_file_path = os.path.join(NVSHMEM_LIB_PATH, "libnvshmem.a")else:NVSHMEM_LIB_PATH = "/usr/lib/x86_64-linux-gnu/"lib_file_path = os.path.join(NVSHMEM_LIB_PATH, "libnvshmem.a")# 配置MPI頭文件路徑（NVSHMEM依賴(lài)MPI）if "MPI_INC_PATH" in os.environ:MPI_INC_PATH = os.environ.get("MPI_INC_PATH")header_file_path = os.path.join(MPI_INC_PATH, "mpi.h")else:MPI_INC_PATH = "/usr/include/"header_file_path = os.path.join(MPI_INC_PATH, "mpi.h")# 配置MPI庫文件路徑if "MPI_LIB_PATH" in os.environ:MPI_LIB_PATH = os.environ.get("MPI_LIB_PATH")lib_file_path = os.path.join(MPI_LIB_PATH, "libmpi.so")else:NVSHMEM_LIB_PATH = "/usr/lib/"lib_file_path = os.path.join(MPI_LIB_PATH, "libmpi.so")# 獲取當前GPU的計算能力（如86對應A100，75對應T4等）target_cc = (torch.cuda.get_device_properties(0).major * 10+ torch.cuda.get_device_properties(0).minor)# 生成CUDA編譯命令cc_cmd = get_compile_command(target_cc=target_cc,  # GPU計算能力cc=cc,  # nvcc編譯器路徑file_name=cuda_code_path,  # 輸入的CUDA源代碼py_include_dir=py_include_dir,  # Python頭文件目錄mirage_home_path=MIRAGE_HOME_PATH,  # Mirage根目錄mirage_inc_path=INCLUDE_PATH,  # Mirage頭文件目錄mirage_deps_path=DEPS_PATH,  # Mirage依賴(lài)目錄nvshmem_inc_path=NVSHMEM_INC_PATH,  # NVSHMEM頭文件目錄nvshmem_lib_path=NVSHMEM_LIB_PATH,  # NVSHMEM庫目錄mpi_inc_path=MPI_INC_PATH,  # MPI頭文件目錄mpi_lib_path=MPI_LIB_PATH,  # MPI庫目錄py_so_path=so_path,  # 輸出的共享庫路徑profiling=True if self.profiler_tensor is not None else False,  # 是否啟用性能分析use_nvshmem=self.use_nvshmem,  # 是否使用NVSHMEMnum_workers=self.num_workers,  # 工作線(xiàn)程數量num_local_schedulers=self.num_local_schedulers,  # 本地調度器數量num_remote_schedulers=self.num_remote_schedulers,  # 遠程調度器數量)# 執行編譯命令，生成共享庫subprocess.check_call(cc_cmd)# 動(dòng)態(tài)導入編譯生成的共享庫import importlib.util# 創(chuàng  )建模塊規格：指定模塊名稱(chēng)和共享庫路徑spec = importlib.util.spec_from_file_location("__mirage_launcher", so_path)# 從規格創(chuàng  )建模塊mod = importlib.util.module_from_spec(spec)# 執行模塊加載spec.loader.exec_module(mod)# 綁定模塊中的核心函數（初始化、啟動(dòng)、結束）self.init_func = getattr(mod, "init_func")self.launch_func = getattr(mod, "launch_func")self.finalize_func = getattr(mod, "finalize_func")# 打印編譯完成提示print("Finished megakernel compilation...")# 收集元數據張量的內存地址指針meta_tensors_ptr = [tensor.data_ptr() for tensor in self.meta_tensors]# 獲取性能分析緩沖區的內存地址（若啟用性能分析）profiler_buffer_ptr = (self.profiler_tensor.data_ptr() if self.profiler_tensor is not None else 0)# 調用初始化函數，傳入必要的參數self.init_func(meta_tensors_ptr,  # 元數據張量指針列表profiler_buffer_ptr,  # 性能分析緩沖區指針self.mpi_rank,  # 當前MPI進(jìn)程編號self.num_workers,  # 工作線(xiàn)程數量self.num_local_schedulers,  # 本地調度器數量self.num_remote_schedulers,  # 遠程調度器數量self.max_seq_length,  # 最大序列長(cháng)度self.eos_token_id,  # 結束符token ID)# 標記編譯完成狀態(tài)self._is_compiled = True

總的來(lái)說(shuō)，compile 函數的作用是將用戶(hù)通過(guò) PersistentKernel API 定義的計算圖轉換為高度優(yōu)化的 CUDA 代碼，并將其編譯為可在當前系統上運行的 Python 模塊，從而實(shí)現高性能的 GPU 計算。編譯成功后，用戶(hù)可以通過(guò)調用 init_func, launch_func, finalize_func 來(lái)初始化、執行和清理內核。

0x05 執行

persistent_kernel.py調用call()方法啟動(dòng)內核執行。

    def __call__(self, **kwargs):self.launch_func()if self.profiler_tensor is not None:from .profiler_persistent import export_to_perfetto_traceexport_to_perfetto_trace(self.profiler_tensor, f"mirage_{self.mpi_rank}.perfetto-trace")

launch_func()函數會(huì )調用launch_persistent_kernel()來(lái)啟動(dòng)內核。

static PyObject *launch_func(PyObject *self, PyObject *args) {launch_persistent_kernel();Py_RETURN_NONE;
}

0xFF 參考

如何評價(jià)CMU將LLM轉化為巨型內核的Mirage Persistent Kernel(MPK)工作？

Mirage: A Multi-Level Superoptimizer for Tensor Programs 簡(jiǎn)記 塵伊光

OSDI2025論文筆記：Mirage: A Multi-Level Superoptimizer for Tensor Programs 畫(huà)餅充饑

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