淘宝客如何免费做网站,服务器及网站建设的特点,个人网站注销原因,wordpress版本编辑因为转译系统需要通过persistent_kernel.py来完成#xff0c;所以我们先介绍persistent_kernel.py。persistent_kernel.py是 Persistent Kernel的Python接口#xff0c;本质是Python到CUDA持久化内核系统的桥梁#xff0c;允许用户用python定义复杂的计算图#xff0c;然后…因为转译系统需要通过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()方法或者自动析构。具体如下图所示。3-10x02 初始化初始化函数会创建 PersistentKernel 类。因为此处只是系统接口大部分有意义的工作在C代码中实现因此此处略过。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_fingerprintTrue))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等定义计算图。简易流程如下MPK-3-2对应的代码举例如下def attach_input(self, torch_tensor: torch.Tensor, name: str None) - DTensor:将PyTorch张量附加到计算图创建对应的DTensor分布式张量。参数torch_tensor: 待附加的PyTorch张量name: 张量名称必须指定返回与输入张量关联的DTensor实例说明仅支持行优先row-major内存布局通过步长校验确保布局正确性# 提取张量维度与步长信息dims tuple([d for d in torch_tensor.shape])strides tuple([s for s in torch_tensor.stride()])# 校验是否为行优先布局高维步长 低维步长 × 低维尺寸for d in range(len(dims) - 1):assert strides[d] strides[d 1] * dims[d 1]# 转换PyTorch数据类型为框架内部 dtypedtype convert_torch_type_to_dtype(torch_tensor.dtype)# 创建输入张量节点t self.kn_graph.new_input(dimsdims, stridesstrides, dtypedtype)# 断言名称非空当前实现限制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:创建新的DTensor并根据IO类别附加到计算图。参数dims: 张量维度元组strides: 步长元组默认自动按行优先计算dtype: 数据类型默认bfloat16name: 张量名称必须指定io_category: IO类别cuda_tensor或nvshmem_tensor返回新创建的DTensor实例说明支持CUDA本地张量与NVSHMEM分布式张量两种类型# 若指定步长校验是否为行优先布局if strides is not None:for d in range(len(dims) - 1):assert strides[d] strides[d 1] * dims[d 1]# 创建张量节点t self.kn_graph.new_input(dimsdims, stridesstrides, dtypedtype)# 断言名称非空当前实现限制assert name is not None# 根据IO类别绑定张量到计算图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(fInvalid io_category: {io_category})return tdef fuse_tensors(self, inputs: list[DTensor], fused_dim: int, num_groups: int, name: str None) - DTensor:融合多个张量到单个张量当前仅支持第0维融合。参数inputs: 待融合的DTensor列表fused_dim: 融合维度必须为0num_groups: 分组数量name: 融合后张量名称返回融合后的DTensor实例# 当前仅支持第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网格维度block_dim: tuple, # CUDA块维度input_source: int 0, # 输入源类型0: 全 tokens, 1: 输入 token):定义嵌入层计算将输入张量通过嵌入权重映射到隐藏空间。参数input: 输入张量weight: 嵌入权重张量output: 输出张量用于存储结果grid_dim: CUDA kernel的网格维度block_dim: CUDA kernel的块维度input_source: 输入源类型标记说明内部创建线程块图TBGraph定义输入输出映射关系并注册为embedding任务# 创建线程块图CyTBGraph为底层实现64为共享内存大小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, (1, -1, -1), -1, True) # 权重张量维度映射tb_graph.new_input(output, (1, 0, -1), -1, True) # 输出张量维度映射# 将张量与线程块图关联self.kn_graph.customized([input, weight, output], tb_graph)# 注册嵌入层任务附加输入源参数self.kn_graph.register_task(tb_graph, embedding, [input_source])def rmsnorm_linear_layer(self,input: DTensor, # 输入张量weight_norm: DTensor, # 归一化权重weight_linear: DTensor, # 线性层权重output: DTensor, # 输出张量grid_dim: tuple, # CUDA网格维度block_dim: tuple, # CUDA块维度):定义RMS归一化线性变换组合层。参数input: 输入张量2Dweight_norm: RMS归一化权重2Dweight_linear: 线性层权重2Doutput: 输出张量2Dgrid_dim: CUDA kernel的网格维度block_dim: CUDA kernel的块维度说明先对输入执行RMS归一化再通过线性层变换输出结果存储到output# 校验输入张量维度当前仅支持2D张量assert input.num_dims 2assert weight_linear.num_dims 2assert output.num_dims 2# 创建线程块图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) # 线性层权重tb_graph.new_input(output, (1, -1, -1), -1, True) # 输出张量# 关联张量与线程块图self.kn_graph.customized([input, weight_norm, weight_linear, output], tb_graph)# 注册RMS归一化线性层任务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网格维度block_dim: tuple, # CUDA块维度):定义注意力层计算支持 rotary 位置编码与 Q/K 归一化。参数input: 输入张量2Dk_cache: 键缓存张量4Dv_cache: 值缓存张量4Dq_norm: Q归一化权重可选1Dk_norm: K归一化权重可选1Dcos_pos_embed: 余弦位置编码可选2Dsin_pos_embed: 正弦位置编码可选2Doutput: 输出张量2Dgrid_dim: CUDA kernel的网格维度block_dim: CUDA kernel的块维度说明自动检测是否启用 rotary 编码与 Q/K 归一化动态调整计算逻辑# 校验输入输出张量维度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)# 提取注意力头相关参数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]# 创建线程块图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) # 输出张量# 关联所有张量与线程块图self.kn_graph.customized([input,k_cache,v_cache,q_norm,k_norm,cos_pos_embed,sin_pos_embed,output,],tb_graph,)# 注册注意力层任务附加参数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 对象的成员变量供后续调用。流程图如下MPK-3-3具体代码如下def compile(self,**kwargs,):# 从关键字参数中获取输出目录默认为Noneoutput_dir kwargs.get(output_dir, None)# 获取Mirage相关的核心路径根目录、包含目录、依赖目录MIRAGE_ROOT, INCLUDE_PATH, DEPS_PATH get_key_paths()# 创建临时目录用于存放编译过程中的中间文件tempdir_obj tempfile.TemporaryDirectory()tempdir tempdir_obj.name# 生成任务图根据GPU数量和当前GPU ID划分计算任务results self.kn_graph.generate.generate_task_graph(num_gpusself.world_size, my_gpu_idself.mpi_rank)# 定义CUDA代码和编译产物的临时路径cuda_code_path os.path.join(tempdir, test.cu) # 生成的CUDA源代码路径so_path os.path.join(tempdir, test.cpython-38-x86_64-linux-gnu.so) # 编译后的的共享库路径# 定义任务图JSON文件的临时路径json_file_path os.path.join(tempdir, task_graph.json)# 将任务图数据写入JSON文件with open(json_file_path, w) as f:f.write(results[json_file])# 将生成的CUDA代码与硬编码补充内容合并后写入文件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_okTrue) # 确保输出目录存在已存在则不报错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)) # 复制任务图JSON# 检查nvccCUDA编译器是否存在cc shutil.which(nvcc)if cc is None:# 若未找到nvcc抛出运行时错误提示用户安装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的头文件包含目录用于编译时链接Python库py_include_dir sysconfig.get_paths(schemescheme)[include]# 从环境变量中获取Mirage的安装路径若已设置if MIRAGE_HOME in os.environ:MIRAGE_HOME_PATH os.environ.get(MIRAGE_HOME)# 初始化NVSHMEM和MPI相关的路径变量用于分布式通信NVSHMEM_INC_PATH None # NVSHMEM头文件目录NVSHMEM_LIB_PATH None # NVSHMEM库文件目录MPI_INC_PATH None # MPI头文件目录MPI_LIB_PATH None # MPI库文件目录# 若启用NVSHMEMNVIDIA共享内存库配置其相关路径if self.use_nvshmem:# 配置NVSHMEM头文件路径if NVSHMEM_INC_PATH in os.environ:# 优先使用环境变量中指定的路径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依赖MPIif 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对应A10075对应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_cctarget_cc, # GPU计算能力cccc, # nvcc编译器路径file_namecuda_code_path, # 输入的CUDA源代码py_include_dirpy_include_dir, # Python头文件目录mirage_home_pathMIRAGE_HOME_PATH, # Mirage根目录mirage_inc_pathINCLUDE_PATH, # Mirage头文件目录mirage_deps_pathDEPS_PATH, # Mirage依赖目录nvshmem_inc_pathNVSHMEM_INC_PATH, # NVSHMEM头文件目录nvshmem_lib_pathNVSHMEM_LIB_PATH, # NVSHMEM库目录mpi_inc_pathMPI_INC_PATH, # MPI头文件目录mpi_lib_pathMPI_LIB_PATH, # MPI库目录py_so_pathso_path, # 输出的共享库路径profilingTrue if self.profiler_tensor is not None else False, # 是否启用性能分析use_nvshmemself.use_nvshmem, # 是否使用NVSHMEMnum_workersself.num_workers, # 工作线程数量num_local_schedulersself.num_local_schedulers, # 本地调度器数量num_remote_schedulersself.num_remote_schedulers, # 远程调度器数量)# 执行编译命令生成共享库subprocess.check_call(cc_cmd)# 动态导入编译生成的共享库import importlib.util# 创建模块规格指定模块名称和共享库路径spec importlib.util.spec_from_file_location(__mirage_launcher, so_path)# 从规格创建模块mod importlib.util.module_from_spec(spec)# 执行模块加载spec.loader.exec_module(mod)# 绑定模块中的核心函数初始化、启动、结束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进程编号self.num_workers, # 工作线程数量self.num_local_schedulers, # 本地调度器数量self.num_remote_schedulers, # 远程调度器数量self.max_seq_length, # 最大序列长度self.eos_token_id, # 结束符token ID)# 标记编译完成状态self._is_compiled True总的来说compile 函数的作用是将用户通过 PersistentKernel API 定义的计算图转换为高度优化的 CUDA 代码并将其编译为可在当前系统上运行的 Python 模块从而实现高性能的 GPU 计算。编译成功后用户可以通过调用 init_func, launch_func, finalize_func 来初始化、执行和清理内核。0x05 执行persistent_kernel.py调用call()方法启动内核执行。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, fmirage_{self.mpi_rank}.perfetto-trace)launch_func()函数会调用launch_persistent_kernel()来启动内核。static PyObject *launch_func(PyObject *self, PyObject *args) {launch_persistent_kernel();Py_RETURN_NONE;