网站建设哪家服务态度好wordpress支付宝红包

网站建设哪家服务态度好,wordpress支付宝红包,搭建一个网址,可以免费发布信息的网站有哪些x00 概要 CMU 贾志豪老师团队提出的MPK#xff08;Mirage Persistent Kernel#xff09;是依托 Mirage 编译器生态的创新运行时系统#xff0c;其核心能力在于将多GPU环境下大语言模型#xff08;LLM#xff09;推理任务自动转换为适配GPU架构的高性能巨型内核#xff0…x00 概要CMU 贾志豪老师团队提出的MPKMirage Persistent Kernel是依托 Mirage 编译器生态的创新运行时系统其核心能力在于将多GPU环境下大语言模型LLM推理任务自动转换为适配GPU架构的高性能巨型内核megakernel。MPK的关键优势在于将传统由CPU负责的内核调度和任务依赖管理工作转移到GPU端通过“长期驻留的巨型内核Persistent Kernel”自主完成同时统筹GPU内部计算与跨GPU通信任务。这种设计不仅大幅削减了CPU-GPU交互带来的内核启动开销还通过计算与通信的细粒度重叠将推理延迟优化至接近硬件物理极限显著提升推理效率。0.1 传统LLM推理框架的瓶颈传统LLM推理框架的流程存在固有瓶颈CPU需逐个发起CUDA内核调用如矩阵乘法、激活函数计算待GPU执行完当前内核并反馈后再触发下一个内核。这种“CPU发起-GPU执行-CPU等待”的循环会产生频繁的CPU-GPU通信与内核启动开销尤其在自回归生成场景中单次token生成需多轮内核调用开销会持续累积严重拖累整体推理性能。0.2 MPK的流程重构MPK彻底重构了这一流程它仅需CPU在推理初始化阶段向GPU提交一个“永不主动退出的persistent_kernel”之后所有任务分派如层间计算顺序、依赖管理如等待前一层结果再执行下一层均由GPU内部自主完成。此时CPU的角色从“实时调度的包工头”转变为“启动初始化的门卫”仅负责触发首次内核启动后续不再参与任何具体调度。0.3 MPK的关键优势MPK通过将多GPU的LLM推理任务转换为高性能的巨型内核从根本上改变了GPU的运行模式。它不仅减少了内核启动开销还通过细粒度的软件流水线和计算通信重叠显著提高了推理效率。MPK提供了一种全新的思路将性能优化的重心从“如何调用优化库”转移到了“如何为整个模型生成一个最优的、原生的执行体”在多GPU环境下实现了更高的吞吐量和更低的延迟。0x01 问题重新设计类似 Mirage 的 MegaKernel的优势是将所有计算和通信融合进一个单一的巨型内核也称为持续内核是降低大语言模型推理延迟的最有效方法之一。这种方法通过启动一个GPU内核来执行整个模型从逐层计算到GPU间通信整个过程无需中断。尽管有这些优势将LLM编译成巨型内核仍然极具挑战性。1.1 现有框架问题现有框架难以支持单一的巨型内核。现有的高级ML框架如PyTorch、Triton和TVM并不原生支持端到端巨型内核生成。现代LLM系统由各种不同的专用内核库构建而成这种碎片化使得将整个推理流水线整合进一个单一的、统一的内核变得非常困难。高性能GPU内核的手工编写需要大量的专家知识如何自动生成高性能内核代码是一个痛点问题。传统做法依赖于专家编写好的内核或者手工融合规则但这些方法维护成本高容易漏掉跨内核/层级组合优化的机会。1.2 编程抽象层级从编程抽象层级上来看也缺乏最优系统。1.2.1 GPU架构下图展示了当今 GPU 的层次结构。GPU 上的计算被组织为内核每个内核都是一个函数以单程序多数据SPMD的方式在多个 GPU 核心上同时执行。一个内核包括一个线程块网格每个线程块在一个 GPU 流式多处理器上执行并包括多个线程来对单个数据元素进行计算。每个线程都与一个每线程寄存器文件相关联并且线程块内的所有线程都可以访问共享内存以启用集体操作。最后内核的所有输入和输出都存储在 GPU 设备内存中。下图是GPU hierarchy。gpu_hierarchy下图为GPU 计算架构和编程抽象示意图image1.2.2 编程视角Triton 是一款高级 GPU 编程框架其编程视角主要聚焦于块Block级别。该框架的设计允许开发者以块为单位进行编程而块内部的优化工作则由 Triton 编译器自动完成。这种设计模式使开发者能够将精力集中在高层逻辑的构建上无需深入研究线程Thread级别的细节实现。Triton 的核心优势在于其简洁的编程模型和自动化优化能力这使得它在处理复杂并行任务时具有更高的效率。Cutlass 则属于底层 GPU 编程库其编程视角覆盖了块Block、线程束Warp与线程Thread的完整层级。Cutlass 提供了丰富的 CUDA 模板和底层控制接口开发者可以利用这些工具精细调控每个线程的行为从而实现高度优化的计算内核。这种细粒度的控制能力让 Cutlass 在对性能有极致要求的场景中表现出色但同时也增加了编程的复杂性。正是这种编程视角的层级差异构成了当前高性能 GPU 编程领域的核心挑战缺乏一套能够 “跨内核Kernel、线程块Block、线程Thread三个层级” 联合搜索最优计算方案并自动验证方案正确性的系统。现有框架要么局限于单一层级的优化例如 Triton 仅针对块内部逻辑进行优化而 Cutlass 则需要开发者手动协调全层级的适配要么无法在多层级协同后确保计算结果的准确性。这一问题在大型语言模型LLM推理等复杂张量计算场景中会显著增加开发成本与优化难度。0x02 总体思路MegaKernel 可被视为一种 grid 级网格级的内核抽象。与 CUDA 的 thread 级线程级抽象、Triton 的 block 级块级抽象不同它提供了层次更高的抽象能力允许开发者在 grid 级开展编程工作。这种抽象设计能让开发者更灵活地管理 GPU 上的计算资源进而实现更高效的内核生成与执行。MPK 的工作原理主要包含以下两部分MPK 编译器负责将大语言模型LLM的计算图转换为经过优化的任务图。MPK 运行时系统在单个巨型内核内部执行任务图以此达成高吞吐量与低延迟的目标。2.1 编译过程模型翻译将 PyTorch 框架下的模型翻译为 MPK 的指令集这一步骤本质上相当于用 MPK 的指令重新构建模型的过程。尽管 PyTorch 具备强大的自动微分与优化能力但要将模型完整转换为 MPK 的指令集仍需进行大量手动调整与优化操作。任务图生成编译器会将翻译后的模型进一步转换为细粒度任务图。该任务图属于有向无环图DAG图中每个节点代表一项具体任务节点间的边则代表任务之间的依赖关系。这一步骤要求编译器能够准确识别并优化任务间的依赖关系为后续高效调度奠定基础。2.2 执行过程任务调度将生成的任务图交付调度器执行。调度器负责管理 GPU 中的流式多处理器SM并通过 warp specialization线程束特化技术将 SM 划分为 worker工作单元与 scheduler调度单元。这种设计与数据处理领域的 actor 模型角色模型相似scheduler 负责协调任务的执行顺序worker 则负责具体执行分配到的任务。性能优化在小模型与低 Batch批次场景下MPK 通过多种方式显著降低延迟具体包括消除内核启动开销、打破内核边界限制、实现细粒度的 SM 调度以及对任务特定模式进行融合。0x03 通过代码来打通流程我们以demo_chat.py为例来进行全局打通。在此文件中会将Python模型结构映射为Mirage的计算图表示然后编译为高效的持久化CUDA内核执行。3.1 核心模块说明3.1.1 三层结构化图模型Mirage 实现了多层次计算图表示μGraphs通过 kernel-graph、block-graph 和 thread-graph 这三层结构化图模型精确映射了 GPU 程序从内核到线程的执行逻辑与存储层级。这种三层结构与 CUDA 程序的执行层级及 GPU 的存储体系紧密对应每层都清晰定义了“算子类型 - 张量存储 - 核心功能”的关联。三层图功能如下Kernel Graph 是最高计算图定义整个执行流程。通过自定义操作管理多个block graphBlock Graph 是嵌套在自定义操作中定义线程块执行序列Thread Graph是最低层定义线程级别执行细节3.1.2 PersistentKernelPersistentKernel 作为计算图的容器和执行器提供了从计算图构建、优化到执行的过程。persistent_kernel.py是 PersistentKernel的Python接口本质是Python到CUDA持久化内核系统的桥梁允许用户用python定义复杂的计算图然后在GPU上高效执行。3.1.3 层级关系计算图与 PersistentKernel 的关系如下包含关系PersistentKernel 内部包含并管理一个 Kernel graph构建关系通过 PersistentKernel 的各种layer方法构建计算图。转换关系PersistentKernel 将计算图转换为可执行的任务图执行关系PersistentKernel 是计算图的执行引擎。3.1.4 数据流关系数据流关系可以近似如下图所示应用层PersistentKernel.py创建并管理kernel graph││▼输入张量││▼计算图节点各种layer方法添加││▼任务层kernel graph包括所有操作和计算流即定义张量数据流││▼并行层block graph嵌套在自定义操作中定义线程块执行序列即定义内存访问模式││▼执行层task graphkernel graph生成的可执行任务图taskDesc是可执行任务EventDesc管理事件同步和依赖││▼运行时环境PersistentKernel 执行引擎││▼硬件层Thread graph在实际GPU线程中执行具体操作3.2 main()代码demo_chat.py的main()如下。def main():world_size, rank setup_distributed_environment()model, tokenizer load_model_and_tokenizer(rank)tokens torch.full((1, MAX_SEQ_LEN), 0, dtypetorch.long, devicecuda)step_tensor torch.tensor([0], dtypetorch.int32, devicecuda)mpk Noneif args.use_mirage:# 构建计算图mpk build_mirage_graph(model, world_size, rank, args, tokens, step_tensor)positions torch.arange(MAX_SEQ_LEN).unsqueeze(0).to(model.device)position_embeddings model.model.rotary_emb(positions)messages [{role: system, content: SYSTEM_PROMPT}]while True:prompt_container [None]if rank 0:try:prompt input( User: )prompt_container[0] promptexcept EOFError:prompt_container[0] exitif world_size 1:dist.broadcast_object_list(prompt_container, src0)prompt prompt_container[0]messages.append({role: user, content: prompt})text tokenizer.apply_chat_template(messages, tokenizeFalse, add_generation_promptTrue)model_inputs tokenizer([text], return_tensorspt).to(model.device)new_prompt_len model_inputs.input_ids.shape[-1]tokens[0, :new_prompt_len] model_inputs.input_ids[0]if new_prompt_len tokens.shape[1]:tokens[0, new_prompt_len:] 0prompt_len new_prompt_lenif args.use_mirage:end_pos, run_time, generated_len run_mirage_generation(model, mpk, tokens, prompt_len, step_tensor, position_embeddings)else:end_pos, run_time, generated_len run_pytorch_generation(model, tokens, prompt_len, step_tensor, position_embeddings)if rank 0:assistant_response_ids tokens[0, prompt_len:end_pos]assistant_response tokenizer.decode(assistant_response_ids, skip_special_tokensTrue)if world_size 1:dist.destroy_process_group()print(Exiting demo.)总体过程如下模型定义阶段使用PyTorch/HuggingFace定义模型结构。加载预训练权重初始化输入张量和相关参数。任务图构建阶段通过KNOperator定义计算操作构建完整的计算图结构设置任务配置参数任务图优化阶段分析任务间的依赖关系生成事件描述以管理依赖对任务进行合理分组以优化执行任务图转换阶段生成TaskDesc描述每个计算任务生成EventDesc描述任务间同步事件生成CUDA可执行代码输出JSON配置文件用于运行时加载。运行时初始化阶段配置GPU资源worker调度器等分配GPU内存给任务队列和事件队列初始化工作队列和调度队列。设置事件计数器和相关同步机制持久化内核运行阶段worker执行具体计算任务调度器负责任务调度和事件管理通过事件机制协调任务间的依赖关系支持多GPU环境下的分布式执行。3.3 关键步骤3.3.1 计算图构建过程此处对应模型翻译过程即将 PyTorch 框架下的模型翻译为 MPK 的指令集这一步骤本质上相当于用 MPK 的指令重新构建模型的过程。尽管 PyTorch 具备强大的自动微分与优化能力但要将模型完整转换为 MPK 的指令集仍需进行大量手动调整与优化操作。模型转换为计算图的工作是在build_mirage_graph函数中其主要步骤如下初始化持久化内核首先构建PersistentKernel实例。mpk mi.PersistentKernel(world_sizeworld_size,mpi_rankrank,num_workers96,num_local_schedulers48,num_remote_schedulers0,max_seq_length4096,eos_token_idmodel.config.eos_token_id,meta_tensors[step_tensor, tokens_tensor],profiler_tensorprofiler_tensor,)定义张量将模型权重和中间张量添加到计算图中。# 输入张量x mpk.attach_input(torch_tensorinput_tokens, nameinput_token)# 位置编码positions torch.arange(MAX_SEQ_LEN).unsqueeze(0).to(model.device)position_embeddings model.model.rotary_emb(positions)x mpk.attach_input(torch_tensorinput_tokens, nameinput_token)cos_pos_embed mpk.attach_input(torch_tensorposition_embeddings[0][0, :MAX_CONTEXT_LEN, :],namecos_position_embedding,)sin_pos_embed mpk.attach_input(torch_tensorposition_embeddings[1][0, :MAX_CONTEXT_LEN, :],namesin_position_embedding,)# 计算图的中间结果张量embed_out mpk.new_tensor(dims(batch_size, hidden_size), dtypemi.bfloat16, nameembed_out)attn_in mpk.new_tensor(dims(batch_size, fused_outdim_1 // world_size), dtypemi.bfloat16, nameattn_in)attn_out mpk.new_tensor(dims(batch_size, num_local_q_heads * head_dim), dtypemi.bfloat16, nameattn_out)is_nvshmem nvshmem_tensor if world_size 1 else cuda_tensorattn_proj_out mpk.new_tensor(dims(batch_size, hidden_size), dtypemi.bfloat16, nameattn_proj_out, io_categoryis_nvshmem)allreduce_buf mpk.new_tensor(dims(world_size, batch_size, hidden_size), dtypemi.bfloat16, nameall_reduce_buf, io_categoryis_nvshmem)attn_allreduce_out mpk.new_tensor(dims(batch_size, hidden_size), dtypemi.bfloat16, nameattn_allreduce_out, io_categoryis_nvshmem)mlp_mid mpk.new_tensor(dims(batch_size, fused_outdim_2 // world_size), dtypemi.bfloat16, namemlp_mid)mlp_out mpk.new_tensor(dims(batch_size, hidden_size), dtypemi.bfloat16, namemlp_out, io_categoryis_nvshmem)mlp_final mpk.new_tensor(dims(batch_size, hidden_size), dtypemi.bfloat16, namemlp_final, io_categoryis_nvshmem)argmax_in mpk.new_tensor(dims(batch_size, vocab_size), dtypemi.bfloat16, nameargmax_in)argmax_part_value mpk.new_tensor(dims(batch_size, 96), dtypemi.bfloat16, nameargmax_part_value)argmax_part_index mpk.new_tensor(dims(batch_size, 96), dtypemi.int64, nameargmax_part_index)argmax_out mpk.new_tensor(dims(batch_size, 1), dtypemi.int64, nameargmax_out)构建计算层通过调用各种layer方法将模型层添加到计算图。此处会把HuggingFace模型权重映射到Mirage张量。也可以融合张量以提高计算效率。# --- Define the Model Graph ---w_embed mpk.attach_input(torch_tensormodel.model.embed_tokens.weight, nameembed_tokens)mpk.embed_layer(inputx, weightw_embed, outputembed_out, grid_dim(1, 1, 1), block_dim(128, 1, 1))x embed_outfor i, layer in enumerate(model.model.layers):# Attention blockw_norm_attn mpk.attach_input(torch_tensorlayer.input_layernorm.weight, nameflayer_{i}_input_layernorm)w_q mpk.attach_input(torch_tensorlayer.self_attn.q_proj.weight, nameflayer_{i}_q_proj)w_k mpk.attach_input(torch_tensorlayer.self_attn.k_proj.weight, nameflayer_{i}_k_proj)w_v mpk.attach_input(torch_tensorlayer.self_attn.v_proj.weight, nameflayer_{i}_v_proj)w_qkv mpk.fuse_tensors(inputs[w_q, w_k, w_v], fused_dim0, num_groupsnum_local_kv_heads, nameflayer_{i}_qkv_proj)mpk.rmsnorm_linear_layer(inputx, weight_normw_norm_attn, weight_linearw_qkv, outputattn_in, grid_dim(96, 1, 1), block_dim(128, 1, 1))w_q_norm mpk.attach_input(torch_tensorlayer.self_attn.q_norm.weight, nameflayer_{i}_q_norm)w_k_norm mpk.attach_input(torch_tensorlayer.self_attn.k_norm.weight, nameflayer_{i}_k_norm)k_cache mpk.attach_input(torch_tensormodel.model.kv_cache[0][i], nameflayer_{i}_k_cache)v_cache mpk.attach_input(torch_tensormodel.model.kv_cache[1][i], nameflayer_{i}_v_cache)mpk.attention_layer(inputattn_in, q_normw_q_norm, k_normw_k_norm, k_cachek_cache, v_cachev_cache, cos_pos_embedcos_pos_embed, sin_pos_embedsin_pos_embed, outputattn_out, grid_dim(batch_size, num_local_kv_heads, 1), block_dim(128, 1, 1))w_o_proj mpk.attach_input(torch_tensorlayer.self_attn.o_proj.weight, nameflayer_{i}_o_proj)mpk.linear_with_residual_layer(inputattn_out, weightw_o_proj, residualx, outputattn_proj_out, grid_dim(hidden_size // 64, 1, 1), block_dim(128, 1, 1))x attn_proj_outif world_size 1:mpk.allreduce_layer(inputattn_proj_out, bufferallreduce_buf, outputattn_allreduce_out, grid_dim(hidden_size // 64, 1, 1), block_dim(128, 1, 1))x attn_allreduce_out# MLP blockresidual_mlp xw_norm_mlp mpk.attach_input(torch_tensorlayer.post_attention_layernorm.weight, nameflayer_{i}_post_attn_layernorm)w_gate_proj mpk.attach_input(torch_tensorlayer.mlp.gate_proj.weight, nameflayer_{i}_gate_proj)w_up_proj mpk.attach_input(torch_tensorlayer.mlp.up_proj.weight, nameflayer_{i}_up_proj)w_gatedup mpk.fuse_tensors(inputs[w_gate_proj, w_up_proj], fused_dim0, num_groups1, nameflayer_{i}_gatedup_proj)mpk.rmsnorm_linear_layer(inputx, weight_normw_norm_mlp, weight_linearw_gatedup, outputmlp_mid, grid_dim(96, 1, 1), block_dim(128, 1, 1))w_down_proj mpk.attach_input(torch_tensorlayer.mlp.down_proj.weight, nameflayer_{i}_down_proj)mpk.silu_mul_linear_with_residual_layer(inputmlp_mid, weightw_down_proj, residualresidual_mlp, outputmlp_out, grid_dim(hidden_size // 64, 1, 1), block_dim(128, 1, 1))x mlp_outif world_size 1:mpk.allreduce_layer(inputmlp_out, bufferallreduce_buf, outputmlp_final, grid_dim(hidden_size // 64, 1, 1), block_dim(128, 1, 1))x mlp_final# Final layerw_final_norm mpk.attach_input(torch_tensormodel.model.norm.weight, namemodel_norm_weight)w_lm_head mpk.attach_input(torch_tensorlm_head_weight, namelm_head)mpk.rmsnorm_linear_layer(inputx, weight_normw_final_norm, weight_linearw_lm_head, outputargmax_in, grid_dim(96, 1, 1), block_dim(128, 1, 1))# Argmaxmpk.argmax_partial_layer(inputargmax_in, output(argmax_part_value, argmax_part_index), grid_dim(96, 1, 1), block_dim(128, 1, 1))mpk.argmax_reduce_layer(input(argmax_part_value, argmax_part_index), outputargmax_out, grid_dim(1, 1, 1), block_dim(128, 1, 1))3.3.2 任务图生成此处对应任务图生成编译器会将翻译后的模型进一步转换为细粒度任务图。该任务图属于有向无环图DAG图中每个节点代表一项具体任务节点间的边则代表任务之间的依赖关系。这一步骤要求编译器能够准确识别并优化任务间的依赖关系为后续高效调度奠定基础。调用compile()方法生成最终的执行图。compile()函数内会执行生成任务图。创建CUDA代码。调用nvcc编译器。创建Python绑定模块。mpk.compile()print(Mirage graph compiled.)return mpk