海尔网站建设水平网站标题特效

海尔网站建设水平,网站标题特效,wordpress修改产品价格,建设部网站上怎样查询企业业绩0x00 概述MPK 包含内置 GPU 运行时系统#xff0c;可在单个 GPU 巨型内核内完整执行任务图。这使得系统能在推理过程中无需额外内核启动的情况下#xff0c;实现任务执行与调度的细粒度控制#xff0c;以实现高吞吐量与低延迟。这座超级工厂能全自动运转#xff0c;核心在于…0x00 概述MPK 包含内置 GPU 运行时系统可在单个 GPU 巨型内核内完整执行任务图。这使得系统能在推理过程中无需额外内核启动的情况下实现任务执行与调度的细粒度控制以实现高吞吐量与低延迟。这座超级工厂能全自动运转核心在于MPK设计了一套跑在GPU上的运行时系统。这套系统的精髓在persistent_kernel.py前端接口和persistent_kernel.cuh后端实现里体现得淋漓尽致。由于所有的调度和任务切换都发生在单一内核上下文内任务间的开销极低通常仅需 1-2 微秒从而能够高效地执行多层、多 GPU 的 LLM 工作负载。0x01 SM不同角色为了实现任务执行与调度的细粒度控制MPK 在启动时将 GPU 上所有流式多处理器SM静态分区为两种角色即工作单元Worker和调度器Scheduler。工作 SM 与调度 SM 的数量在内核启动时固定配置且总和等于物理 SM 总数从而彻底避免动态上下文切换开销。下图展示了 MPK 的执行时间线其中每个矩形代表一个在工作单元上运行的任务每个圆圈代表一个事件。当一个任务完成时它会递增其对应触发事件的计数器。当事件计数器达到预设阈值时该事件被视为已激活并被加入Scheduler的事件队列。随后Scheduler会启动所有依赖于该事件的下游任务。这种设计实现了细粒度的软件流水线化并允许计算与通信之间重叠比如矩阵乘法Matmul任务可以与来自不同层的注意力任务并行执行。一旦有部分 matmul 结果可用即可开始 Allreduce 通信。img1.1 Scheduler SM调度决策由 MPK 的分布式Scheduler处理每个Scheduler运行于单个线程束warp上。由于每个流式多处理器SM可以容纳多个线程束因此单 SM 最多可并发运行 4 个Scheduler。每个Scheduler维护激活事件队列并持续执行以下操作事件出队移除依赖已满足的激活事件即所有前置任务均已完成。任务启动调度依赖该激活事件的任务集。这种分布式调度机制在实现跨 SM 可扩展执行的同时最小化协同开销。这些SM不负责计算它们是GPU内部的“调度系统”。它们监控着一系列“事件”Event。一个事件代表一个或多个前置任务已经完成。当Scheduler监测到某个事件被触发例如一个矩阵乘任务完成了它就会查询预先编译好的任务图找到所有依赖这个事件的后续任务然后把这些新任务的IDTaskId投递到工人们的任务队列里。在persistent_kernel.cuh里else分支就是Scheduler的逻辑。它们也是在一个while(true)循环里不断检查sched_queues处理激活的事件并分派新任务。1.2 Worker SM这些SM是纯粹的执行单元负责干具体的计算活比如矩阵乘、向量加法等。每个工人SM都有一个自己的任务队列Task Queue它们的工作就是不断地从队列里取任务、执行、再取下一个。这在persistent_kernel.cuh的源码里体现得很清楚if (blockIdx.x config.num_workers)这个分支里的while(true)循环就是工人SM执行逻辑的直接实现。每个工作单元独占一个流式多处理器SM并维护专属任务队列。其执行遵循以下高效简洁的循环流程获取任务从队列中提取下一待执行任务。执行计算运行任务如矩阵乘法 / 注意力机制 / GPU 间数据传输。事件触发任务完成后通知触发事件。循环执行重复上述过程。该机制既保障了工作单元的持续满载运行又实现了跨层和跨操作的异步任务执行。0x02 推理引擎persistent_kernel.cuh 是运行引擎的入口文件。2.1 初始化init_persistent_kernel是初始化函数负责初始化运行时环境、各种数据结构和分配内存。具体如下参数解析和配置设置接受来自调用方的参数包括元数据张量、性能分析缓冲区、GPU rank、工作线程数、调度器数等。设置全局运行时配置global_runtime_config的各种参数如工作线程数、调度器数、序列长度限制等。初始化NVSHMEM如果启用初始化NVSHMEMNVIDIA SHared Memory环境用于多GPU间的通信。获取当前GPU ID和总GPU数量。调用_init_persistent_kernel函数该函数用于初始化所有任务和事件描述。内存分配和初始化。比如为所有任务任务描述/TaskDesc、所有事件EventDesc分配内存并复制数据。为工作队列worker_queue和调度器队列sched_queue 分配内存并复制数据。为事件计数器EventCounter分配内存并复制数据。为第一个任务分配内存并复制数据。启动内核调用init_kernel进行内核初始化。如果使用NVSHMEM则添加全局屏障确保所有初始化完成。流程图如下mirage-5-1代码如下// 外部C语言接口Mirage持久化内核MPK的初始化函数// 负责配置内核运行时参数、初始化分布式通信如NVSHMEM、分配GPU内存并加载任务/事件数据extern C void init_persistent_kernel(std::vectorvoid * meta_tensors,void *profiler_buffer,int my_rank,int num_workers,int num_local_schedulers,int num_remote_schedulers,int max_seq_length,long long eos_token_id) {// 断言确保元数据张量数量为3对应step、tokens、new_token_nums三个核心张量assert(meta_tensors.size() 3);// 将元数据张量指针转换为对应类型赋值给全局运行时配置global_runtime_config.step static_castint *(meta_tensors[0]); // 步骤标记张量记录当前生成步骤global_runtime_config.tokens static_castlong long *(meta_tensors[1]); // 生成的token序列张量global_runtime_config.new_token_nums static_castint *(meta_tensors[2]);// 新增token数量张量// 配置工作单元、调度器数量参数global_runtime_config.num_workers num_workers; // 工作单元worker总数global_runtime_config.num_local_schedulers num_local_schedulers; // 本地调度器数量global_runtime_config.num_remote_schedulers num_remote_schedulers; // 远程调度器数量分布式场景// 配置序列生成相关参数global_runtime_config.max_seq_length max_seq_length; // 最大序列长度token生成上限global_runtime_config.eos_token_id eos_token_id; // 结束符token的ID// 配置性能分析缓冲区用于存储性能数据global_runtime_config.profiler_buffer profiler_buffer;// 计算总调度器数量本地调度器 远程调度器int num_schedulers num_local_schedulers num_remote_schedulers;// 初始化分布式通信库NVSHMEM仅当启用USE_NVSHMEM宏时执行cudaSetDevice(my_rank); // 设置当前使用的GPU设备对应分布式场景中的GPU编号#ifdef USE_NVSHMEMMPI_Comm mpi_comm MPI_COMM_WORLD; // 初始化MPI通信域NVSHMEM依赖MPI实现分布式通信nvshmemx_init_attr_t attr NVSHMEMX_INIT_ATTR_INITIALIZER; // 初始化NVSHMEM属性结构体attr.mpi_comm mpi_comm; // 将MPI通信域绑定到NVSHMEM属性nvshmemx_init_attr(NVSHMEMX_INIT_WITH_MPI_COMM, attr); // 基于MPI通信域初始化NVSHMEMnvshmem_barrier_all(); // 所有GPU进程同步确保NVSHMEM初始化完成// 获取NVSHMEM相关进程信息int mype nvshmem_my_pe(); // 当前GPU在NVSHMEM中的进程编号int npes nvshmem_n_pes(); // NVSHMEM中的总进程数即GPU总数int mype_node nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); // 当前进程在节点内的编号// 打印进程信息用于调试和日志记录printf(mype(%d) npes(%d) mype_node(%d)\n, mype, npes, mype_node);#else// 未启用NVSHMEM时默认单GPU场景配置int mype 0; // 进程编号默认为0int npes 1; // 总进程数默认为1#endif// 配置队列长度参数每个工作单元、调度器的队列容量global_runtime_config.per_worker_queue_len 1024; // 每个工作单元队列的最大任务数global_runtime_config.per_sched_queue_len 1024; // 每个调度器队列的最大事件数// 配置GPU相关参数global_runtime_config.num_gpus npes; // GPU总数对应NVSHMEM的总进程数global_runtime_config.my_gpu_id mype; // 当前GPU的编号global_runtime_config.num_graphs 1; // 内核图数量默认1个global_runtime_config.split_worker_scheduler true; // 启用工作单元与调度器分离的架构// 声明任务、事件相关数据结构用于存储任务描述、事件描述、初始任务列表std::vectorTaskDesc all_tasks; // 所有任务的描述信息列表std::vectorEventDesc all_events; // 所有事件的描述信息列表用于管理任务依赖std::vectorTaskId first_tasks; // 初始任务ID列表内核启动时首先执行的任务// 调用内部初始化函数填充任务、事件、初始任务数据_init_persistent_kernel(all_tasks, all_events, first_tasks, npes, mype);// 初始化工作单元队列的最后就绪任务ID数组GPU端内存分配// 每个工作单元维护两个队列本地队列 远程队列因此数组长度为2 * 工作单元数global_runtime_config.worker_queue_last_ready_task_id gpu_mallocunsigned long long int((num_workers * 2) *sizeof(unsigned long long int));// 在主机端初始化该数组初始值均为0表示暂无就绪任务std::vectorunsigned long long int host_worker_queue_last_task_id;for (int i 0; i 2 * num_workers; i) {host_worker_queue_last_task_id.push_back(0);}// 将主机端数组数据拷贝到GPU端内存cudaMemcpy(global_runtime_config.worker_queue_last_ready_task_id,host_worker_queue_last_task_id.data(),(num_workers * 2) * sizeof(unsigned long long int),cudaMemcpyHostToDevice);// 初始化调度器队列的最后就绪事件ID和下一个空闲事件ID数组GPU端内存分配// 额外增加1个队列用于全局调度器因此数组长度为 总调度器数 1global_runtime_config.sched_queue_last_ready_event_id gpu_mallocunsigned long long int((num_schedulers 1) *sizeof(unsigned long long int));global_runtime_config.sched_queue_next_free_event_id gpu_mallocunsigned long long int((num_schedulers 1) *sizeof(unsigned long long int));// 在主机端初始化这两个数组初始值均为0std::vectorunsigned long long int host_sched_queue_last_event_id;for (int i 0; i (num_schedulers 1); i) {host_sched_queue_last_event_id.push_back(0);}// 将主机端数组数据拷贝到GPU端内存两个数组初始值相同cudaMemcpy(global_runtime_config.sched_queue_last_ready_event_id,host_sched_queue_last_event_id.data(),(num_schedulers 1) * sizeof(unsigned long long int),cudaMemcpyHostToDevice);cudaMemcpy(global_runtime_config.sched_queue_next_free_event_id,host_sched_queue_last_event_id.data(),(num_schedulers 1) * sizeof(unsigned long long int),cudaMemcpyHostToDevice);// 初始化所有事件的计数器GPU端内存分配global_runtime_config.all_event_counters gpu_mallocEventCounter(all_events.size() * sizeof(EventCounter)); // 事件触发计数器global_runtime_config.all_event_num_triggers gpu_mallocint(all_events.size() * sizeof(int)); // 事件所需触发次数// 在主机端初始化事件所需触发次数数组从事件描述中读取对应值std::vectorint host_all_event_counters;for (size_t i 0; i all_events.size(); i) {host_all_event_counters.push_back(all_events.at(i).num_triggers);}// 将主机端数据拷贝到GPU端内存cudaMemcpy(global_runtime_config.all_event_num_triggers,host_all_event_counters.data(),all_events.size() * sizeof(int),cudaMemcpyHostToDevice);// 将事件触发计数器初始化为0所有事件初始未触发cudaMemset(global_runtime_config.all_event_counters,0,all_events.size() * sizeof(EventCounter));// 初始化所有任务数据将主机端任务描述拷贝到GPU端global_runtime_config.all_tasks gpu_mallocTaskDesc(all_tasks.size() * sizeof(TaskDesc));cudaMemcpy(global_runtime_config.all_tasks,all_tasks.data(),all_tasks.size() * sizeof(TaskDesc),cudaMemcpyHostToDevice);// 初始化所有事件数据将主机端事件描述拷贝到GPU端global_runtime_config.all_events gpu_mallocEventDesc(all_events.size() * sizeof(EventDesc));cudaMemcpy(global_runtime_config.all_events,all_events.data(),all_events.size() * sizeof(EventDesc),cudaMemcpyHostToDevice);// 初始化工作单元队列GPU端内存分配{std::vectorTaskId * host_worker_queues; // 主机端存储每个工作单元队列的GPU内存指针// 为每个工作单元的两个队列分配GPU内存for (int i 0; i (num_workers * 2); i) {TaskId *worker_queue gpu_mallocTaskId(global_runtime_config.per_worker_queue_len * sizeof(TaskId));host_worker_queues.push_back(worker_queue);}// 分配GPU内存存储所有工作单元队列的指针global_runtime_config.worker_queues gpu_mallocTaskId *((num_workers * 2) * sizeof(TaskId *));// 将队列指针从主机端拷贝到GPU端cudaMemcpy(global_runtime_config.worker_queues,host_worker_queues.data(),(num_workers * 2) * sizeof(TaskId *),cudaMemcpyHostToDevice);}// 初始化调度器队列GPU端内存分配{std::vectorEventId * host_sched_queues; // 主机端存储每个调度器队列的GPU内存指针// 为每个调度器队列含全局调度器队列分配GPU内存for (int i 0; i (num_schedulers 1); i) {EventId *sched_queue gpu_mallocEventId(global_runtime_config.per_sched_queue_len * sizeof(EventId));host_sched_queues.push_back(sched_queue);}// 分配GPU内存存储所有调度器队列的指针global_runtime_config.sched_queues gpu_mallocEventId *((num_schedulers 1) * sizeof(EventId *));// 将队列指针从主机端拷贝到GPU端cudaMemcpy(global_runtime_config.sched_queues,host_sched_queues.data(),(num_schedulers 1) * sizeof(EventId *),cudaMemcpyHostToDevice);}// 初始化初始任务数据将主机端初始任务ID拷贝到GPU端{global_runtime_config.first_tasks gpu_mallocTaskId(first_tasks.size() * sizeof(TaskId));cudaMemcpy(global_runtime_config.first_tasks,first_tasks.data(),first_tasks.size() * sizeof(TaskId),cudaMemcpyHostToDevice);}// 启动初始化内核GPU端执行配置网格和线程块维度1个block每个block含128个threadinit_kerneldim3(1, 1, 1), dim3(128, 1, 1)(global_runtime_config);cudaDeviceSynchronize(); // 等待GPU初始化内核执行完成#ifdef USE_NVSHMEMnvshmem_barrier_all(); // 分布式场景下所有GPU进程同步确保初始化全部完成#endif}2.2 启动内核launch_persistent_kernel是启动 CUDA 内核的入口函数具体功能如下获取设备信息。获取当前GPU设备编号查询设备的SM数量用于确定grid大小。根据配置选择内核启动模式模式一分离式内核split_worker_scheduler true)这是默认模式分别启动两个独立的内核。线程块大小128个线程网格大小num_local_schedulers执行scheduler_kernel函数。线程块大小128个线程网格大小num_workers每个线程块对应一个worker执行worker_kernel函数。为所有内核设置最大动态共享内存大小为 MAX_SHARE_MEMORY_SIZE。使用cudaDeviceSynchronize等待内核启动完成。模式二单一持久化内核split_worker_scheduler false网格大小sm_count使用所有的sm线程块大小128个线程调用persistent_kernel函数。如果使用NVSHMEM则使用nvshmemx_collective_launch启动否则使用标准CUDA内核启动。为所有内核设置最大动态共享内存大小为 MAX_SHARE_MEMORY_SIZE。使用cudaDeviceSynchronize等待内核启动完成。流程图如下mirage-5-2代码如下。// 外部C语言接口Mirage持久化内核MPK的启动函数// 根据运行时配置的架构模式工作单元与调度器分离/一体化启动对应的GPU内核extern C void launch_persistent_kernel() {int device;// 获取当前正在使用的GPU设备编号cudaGetDevice(device);int sm_count;// 获取当前GPU设备的流式多处理器SM数量用于一体化内核的网格维度配置cudaDeviceGetAttribute(sm_count, cudaDevAttrMultiProcessorCount, device);// 判断是否启用工作单元与调度器分离的架构模式if (global_runtime_config.split_worker_scheduler) {// 打印日志标识当前启动的是分离模式工作单元内核 调度器内核printf(worker kernel scheduler kernel\n);// 配置并启动工作单元内核与调度器内核// 1. 设置工作单元内核的最大动态共享内存大小使用预定义的最大共享内存常量cudaFuncSetAttribute(worker_kernel,cudaFuncAttributeMaxDynamicSharedMemorySize,MAX_SHARE_MEMORY_SIZE);// 2. 设置调度器内核的最大动态共享内存大小cudaFuncSetAttribute(scheduler_kernel,cudaFuncAttributeMaxDynamicSharedMemorySize,MAX_SHARE_MEMORY_SIZE);// 创建两个独立的CUDA流分别用于工作单元内核和调度器内核的异步执行cudaStream_t worker_stream, scheduler_stream;cudaStreamCreate(worker_stream); // 工作单元内核专属流cudaStreamCreate(scheduler_stream); // 调度器内核专属流// 注分离模式不支持NVSHMEM分布式通信// 原因nvshmemx_collective_launch会串行启动内核阻碍工作单元与调度器内核的交互// 启动工作单元内核worker_kerneldim3(global_runtime_config.num_workers, 1, 1), // 网格维度工作单元数量 × 1 × 1每个工作单元对应一个blockdim3(128, 1, 1), // 线程块维度128个thread × 1 × 1MAX_SHARE_MEMORY_SIZE /* 动态共享内存大小 */,worker_stream /* 绑定到工作单元专属流 */(global_runtime_config); // 传入全局运行时配置作为内核参数// 启动调度器内核scheduler_kerneldim3(global_runtime_config.num_local_schedulers, 1, 1), // 网格维度本地调度器数量 × 1 × 1dim3(32, 1, 1), // 线程块维度32个thread × 1 × 10 /* 调度器内核无需动态共享内存设为0 */,scheduler_stream /* 绑定到调度器专属流 */(global_runtime_config); // 传入全局运行时配置作为内核参数// 等待GPU上所有内核执行完成并检查执行错误cudaError_t err cudaDeviceSynchronize();if (err ! cudaSuccess) {// 若执行出错打印错误信息包含具体错误描述printf(CUDA kernel launch error: %s\n, cudaGetErrorString(err));}// 销毁创建的CUDA流释放资源cudaStreamDestroy(worker_stream);cudaStreamDestroy(scheduler_stream);// 打印日志标识持久化内核启动流程完成printf(Finished Launch Persistent Kernel\n);} else {// 打印日志标识当前启动的是一体化模式单个持久化内核printf(a single persistent kernel\n);// 配置并启动一体化持久化内核// 设置一体化内核的最大动态共享内存大小cudaFuncSetAttribute(persistent_kernel,cudaFuncAttributeMaxDynamicSharedMemorySize,MAX_SHARE_MEMORY_SIZE);#ifdef USE_NVSHMEM// 若启用NVSHMEM分布式通信使用NVSHMEM的集合式启动接口适配多GPU协同void *args[] {global_runtime_config}; // 封装内核参数全局运行时配置nvshmemx_collective_launch((void const *)persistent_kernel, // 待启动的一体化内核函数dim3(sm_count, 1, 1), // 网格维度GPU的SM数量 × 1 × 1每个SM对应一个blockdim3(128, 1, 1), // 线程块维度128个thread × 1 × 1args, // 内核参数数组MAX_SHARE_MEMORY_SIZE /* 动态共享内存大小 */,0 /* 不绑定特定流使用默认流 */);#else// 未启用NVSHMEM时直接启动一体化内核单GPU场景persistent_kerneldim3(sm_count, 1, 1), // 网格维度SM数量 × 1 × 1dim3(128, 1, 1), // 线程块维度128个thread × 1 × 1MAX_SHARE_MEMORY_SIZE /* 动态共享内存大小 */(global_runtime_config); // 传入全局运行时配置作为内核参数#endif// 等待GPU内核执行完成并检查执行错误cudaError_t err cudaDeviceSynchronize();if (err ! cudaSuccess) {// 若执行出错打印错误信息printf(CUDA kernel launch error: %s\n, cudaGetErrorString(err));}// 打印日志标识持久化内核启动流程完成printf(Finished Launch Persistent Kernel\n);}}persistent_kernel、worker_kernel、scheduler_kernel的函数如下__global__ void persistent_kernel(RuntimeConfig config) {persistent_checker(config);if (blockIdx.x config.num_workers) {execute_worker(config);} else {execute_scheduler(config, -(4 * config.num_workers));}}__global__ void worker_kernel(RuntimeConfig config) {worker_checker(config);execute_worker(config);}__global__ void scheduler_kernel(RuntimeConfig config) {scheduler_checker(config);execute_scheduler(config, 0);}2.3 Scheduler 实现2.3.1 功能execute_scheduler是 CUDA 内核中负责任务调度的核心函数主要功能如下初始化和设置调度器每个线程块处理4个调度器int warp_thread_id threadIdx.x % 32 / warp_id 4依据调度器ID确定负责的工作节点范围。区分本地调度器处理本地任务和远程调度器处理跨GPU任务的不同处理逻辑。这允许系统依据实际硬件环境和应用需求灵活调整调度器的数量和分工。事件队列处理循环函数进入一个无限循环持续处理事件队列中的事件。轮询调度队列获取待处理的事件。使用原子操作确保队列同步。支持多个调度队列本地和全局。事件分类处理。依据事件类型进行不同处理。终止事件向所有工作节点发送终止任务。结束调度器运行。任务图结束事件 EVENT_END_OF_TASK_GRAPH。调用prepare_next_batch检查是否继续下一批处理。如果需要终止则调用 terminate_schedulers否则为下一次迭代启动新的任务图。依赖任务事件 EVENT_LAUNCH_DEPENDENT_TASKS增加迭代次数。将任务分割给多个本地调度器只有本地调度器才可以处理这类事件。按照工作节点数量进行任务分组。按照轮询方式将任务分配给工作节点。分配算法如下计算每轮需要的任务数量e.last_task_id - e.first_task_id config.num_workers - 1) / config.num_workers。对每个工作节点进行迭代分配位置索引 e.first_task_id i * config.num_workers j如果索引在有效范围内则分配给对应的工作节点。使用next_worker变量实现轮询机制。大规模任务事件 EVENT_LAUNCH_MASSIVE_TASKS将任务分割给多个本地调度器只有本地调度器才可以处理这类事件。按照顺序将任务分发给对应的不同的工作节点具体是使用get_first_last_ids来为每个调度器分配任务子集。通过将大量任务均匀分配给所有本地调度器来实现负载均衡。每个调度器只处理分配给它的任务范围。普通事件任务 EVENT_LAUNCH_TASKS直接将任务分配给工作节点使用轮询方式确保负载均衡2.3.2 流程流程图如下。mirage-5-32.3.3 要点需要注意的地方如下任务分发机制使用轮询方式将任务分发给工作节点。维护每个工作节点的下一个可用任务位置。使用原子操作确保线程安全的任务队列更新。同步和可见性保证使用 relaxed 和 accuire / release 语义确保内存操作的正确顺序。通过原子操作维护队列状态。支持NVSHMEM的跨GPU通信任务间的依赖关系通过任务描述符中的 dependent_event字段建立任务间的依赖关系。在执行任务前检查依赖是否满足。使用计数机制确保依赖任务按照正确顺序执行。任务索引计算size_t position_index e.first_task_id i * config.num_workers j; 通过这种方式系统可以将连续的任务按顺序分配给不同的工作单元。处理依赖任务时会迭代编号更新。增加iteration_num表示进入新的迭代轮询分配依赖任务通过轮询方式分配给工作队列确保任务在worker之间均匀分布。本地调度器专用依赖任务只能由本地调度器处理远程调度器不处理这种任务。在execute_worker 中系统会检查任务的依赖关系。当一个任务完成时会触发相应事件具备依赖关系的任务会放入广播队列中处理。每个线程块处理4个调度器int warp_thread_id threadIdx.x % 32; --- 128 / 32 4每个warp对应一个调度器实例。每个调度器由足够的线程资源32个这样可以充分利用硬件资源将调度器均匀分布在多个线程块中避免单点瓶颈避免过多调度器并行导致资源竞争简化调度器的同步和通信。在 persistent_checker 中有assert(num_schedulers % 4 0);assert(gridDim.x config.num_workers num_schedulers / 4);这说明如果有 N 个调度器就需要 N/4个线程块来处理每个线程块处理4个调度器。假设12个调度器则需要3个线程块来处理所有调度器。线程块0处理调度器03线程块1处理调度器47。只有warp中第一个线程warp_thread_id 0实际执行调度器策略其它线程处于空闲状态可能未来会扩展或者执行辅助任务。远程调度器使用场景多GPU环境中的跨GPU通信任务。当系统配置多个GPU时需要使用远程调度器处理GPU间的通信任务。尤其是处理ALLREDUCE等操作。NVSHEMEM相关任务。比如处理TASK_NVSHMEM_COPY任务时需要远程调度器。这些任务需要在不同GPU之间复制数据。负责协调不同GPU之间的任务执行顺序处理跨GPU的数据传输确保全局唯一性。通过get_first_last_ids将工作负载划分给不同的调度器。__device__ __forceinline__ voidget_first_last_ids(unsigned long long int num_elements,unsigned long long int num_workers,unsigned long long int my_id,unsigned long long int *my_first_element,unsigned long long int *my_last_element) {unsigned long long int num_elements_per_worker num_elements / num_workers;unsigned long long int reminder num_elements % num_workers;if (my_id reminder) {*my_first_element (num_elements_per_worker 1) * my_id;*my_last_element *my_first_element num_elements_per_worker 1;} else {*my_first_element num_elements_per_worker * my_id reminder;*my_last_element *my_first_element num_elements_per_worker;}}2.3.4 代码// 设备端函数执行调度器逻辑每个线程块仅有一个warp参与需注意同步限制// 功能从调度器队列读取事件解析事件类型并生成对应的任务分配至工作单元队列__device__ void execute_scheduler(RuntimeConfig config, int offset) {// 计算总调度器数量本地调度器 远程调度器int num_schedulers config.num_local_schedulers config.num_remote_schedulers;// 计算当前线程在warp内的ID0-31因每个线程块仅一个warp简化同步逻辑int warp_thread_id threadIdx.x % 32;// 以下逻辑禁止使用__syncthreads()避免跨warp同步导致的效率损失// 仅warp内ID为0的线程执行调度核心逻辑单线程负责调度决策减少资源竞争if (warp_thread_id 0) {// 计算当前调度器的全局ID偏移量块索引区分不同调度器实例int sched_id blockIdx.x offset;// 初始化调度器队列相关参数int num_sched_queues 1; // 调度器需处理的队列数量初始为1size_t iteration_num 0; // 当前迭代次数用于任务ID生成EventId *sched_queues[2]; // 调度器需监听的事件队列指针最多2个int sched_queue_ids[2]; // 队列对应的全局ID// 绑定当前调度器的专属队列sched_queues[0] config.sched_queues[sched_id];sched_queue_ids[0] sched_id;// 用于记录当前调度器管理的工作单元ID范围unsigned long long int my_first_worker, my_last_worker;// 区分本地调度器与远程调度器的逻辑if (sched_id config.num_local_schedulers) {// 本地调度器额外处理全局队列的事件多队列监听sched_queues[num_sched_queues] config.sched_queues[num_schedulers];sched_queue_ids[num_sched_queues] num_schedulers;num_sched_queues; // 队列数量增至2// 计算本地调度器管理的工作单元ID范围从总工作单元中分配get_first_last_ids(config.num_workers,config.num_local_schedulers,sched_id,my_first_worker,my_last_worker);} else {// 远程调度器管理的工作单元ID范围偏移总工作单元数避免与本地冲突get_first_last_ids(config.num_workers,config.num_remote_schedulers,sched_id - config.num_local_schedulers, // 远程调度器的本地索引my_first_worker,my_last_worker);// 远程工作单元ID从总工作单元数开始编号my_first_worker config.num_workers;my_last_worker config.num_workers;}// 调试日志打印调度器ID及管理的工作单元范围仅MPK_ENABLE_VERBOSE启用时#ifdef MPK_ENABLE_VERBOSEprintf([SCHD] sched_id(%d) first_worker(%llu) last_worker(%llu)\n,sched_id,my_first_worker,my_last_worker);#endif// 初始化队列位置跟踪变量当前处理位置与最新事件位置支持2个队列size_t cur_event_pos[2], last_event_pos[2];for (int i 0; i 2; i) {cur_event_pos[i] 0; // 当前已处理的事件索引last_event_pos[i] 0; // 队列中最新的事件索引}// 记录每个工作单元队列的下一个空闲位置避免原子操作开销本地缓存size_t worker_queue_next_free_task_pos[MAX_WORKER_PER_SCHEDULER];for (int i 0; i MAX_WORKER_PER_SCHEDULER; i) {worker_queue_next_free_task_pos[i] 0;}// 特殊初始化ID为0的调度器首个工作单元的队列起始位置设为1if (sched_id 0) {worker_queue_next_free_task_pos[0] 1;}// 任务分配的工作单元迭代变量轮询分配时使用int next_worker my_first_worker;// 当前监听的队列索引初始为0int queue_idx 0;// 调度主循环持续运行直至收到终止事件while (true) {// 等待队列中有新事件循环检查当前队列是否有未处理事件while (cur_event_pos[queue_idx] last_event_pos[queue_idx]) {// 使用acquire语义读取最新的事件位置确保数据可见性last_event_pos[queue_idx] ld_acquire_gpu_u64(config.sched_queue_last_ready_event_id[sched_queue_ids[queue_idx]]);// 若当前队列有新事件退出等待否则切换到下一个队列if (cur_event_pos[queue_idx] last_event_pos[queue_idx]) {break;} else {queue_idx (queue_idx num_sched_queues - 1) ? 0 : queue_idx 1;}// 短暂休眠10纳秒避免空循环占用过多资源__nanosleep(10);}// 断言确保调度器队列未溢出当前位置队列长度需大于最新事件位置assert(cur_event_pos[queue_idx] config.per_sched_queue_len last_event_pos[queue_idx]);// 读取当前待处理的事件ID使用relaxed语义平衡性能与可见性EventId event_id ld_relaxed_gpu_u64(sched_queues[queue_idx][cur_event_pos[queue_idx] % config.per_sched_queue_len]);// 获取事件描述信息EventDesc e config.all_events[event_id];// 检查是否为终止事件退出调度循环if (is_termination_event(event_id, e)) {// 若为本地调度器向其管理的所有工作单元发送终止任务ID0if (sched_id config.num_local_schedulers) {for (int i my_first_worker; i my_last_worker; i) {// 获取工作单元队列的下一个空闲位置size_t last_task_id worker_queue_next_free_task_pos[i - my_first_worker];// 写入终止任务ID0st_relaxed_gpu_u64(config.worker_queues[i][last_task_id %config.per_worker_queue_len],0);// 使用release语义更新工作单元队列的最新就绪任务ID确保任务可见atom_add_release_gpu_u64(config.worker_queue_last_ready_task_id[i],1);}}// 退出调度器return;}// 处理任务图结束事件EVENT_END_OF_TASK_GRAPHif (e.event_type EVENT_END_OF_TASK_GRAPH) {#ifdef MPK_ENABLE_VERBOSEprintf([SCHD] END_OF_TASK_GRAPH\n);#endif// 检查是否需要准备下一批次任务if (!prepare_next_batch(config)) {// 无需继续终止所有调度器terminate_schedulers(config);} else {// 为下一次迭代启动任务图开始任务ID1size_t last_task_id worker_queue_next_free_task_pos[next_worker - my_first_worker];// 计算任务ID迭代次数任务索引并写入工作单元队列st_relaxed_gpu_u64(config.worker_queues[next_worker][last_task_id % config.per_worker_queue_len],compute_task_id(iteration_num 1, 1 /*begin_task_graph*/));// 更新工作单元队列的最新就绪任务IDatom_add_release_gpu_u64(config.worker_queue_last_ready_task_id[next_worker], 1);#ifdef MPK_ENABLE_VERBOSE// 打印调试信息GPU ID、调度器ID、迭代次数、任务索引、工作单元ID等printf([%d][SCHD]EVENT_END_OF_TASK_GRAPH schd_id(%d) iter_num(%llu) task_idx(1) worker_id(%d) worker_last_ready_pos(%llu)\n,config.my_gpu_id,sched_id,iteration_num 1,next_worker,last_task_id 1);#endif// 轮询切换到下一个工作单元next_worker (next_worker my_last_worker - 1) ? my_first_worker: next_worker 1;}}// 处理启动依赖任务事件EVENT_LAUNCH_DEPENDENT_TASKSelse if (e.event_type EVENT_LAUNCH_DEPENDENT_TASKS) {// 迭代次数递增标识新一批任务iteration_num iteration_num 1;// 断言该事件仅由本地调度器处理assert(sched_id config.num_local_schedulers);// 按工作单元数量拆分任务范围轮询分配给管理的工作单元for (size_t i 0;i (e.last_task_id - e.first_task_id config.num_workers - 1) /config.num_workers;i) {for (size_t j my_first_worker; j my_last_worker; j) {// 计算当前任务在全局范围内的索引size_t position_index e.first_task_id i * config.num_workers j;// 仅处理范围内的任务if (position_index e.last_task_id) {// 获取工作单元队列的下一个空闲位置size_t last_task_id worker_queue_next_free_task_pos[next_worker -my_first_worker];// 计算任务ID并写入工作单元队列st_relaxed_gpu_u64(config.worker_queues[next_worker][last_task_id %config.per_worker_queue_len],compute_task_id(iteration_num, position_index));// 更新工作单元队列的最新就绪任务IDatom_add_release_gpu_u64(config.worker_queue_last_ready_task_id[next_worker], 1);// 轮询切换到下一个工作单元next_worker (next_worker my_last_worker - 1)? my_first_worker: next_worker 1;}}}}// 处理其他类型事件如EVENT_LAUNCH_MASSIVE_TASKSelse {// 初始化当前调度器需处理的任务范围TaskId my_first_task e.first_task_id, my_last_task e.last_task_id;// 若为启动大规模任务事件按本地调度器数量拆分任务范围if (e.event_type EVENT_LAUNCH_MASSIVE_TASKS) {assert(sched_id config.num_local_schedulers); // 仅本地调度器处理// 计算当前调度器负责的子任务范围get_first_last_ids(e.last_task_id - e.first_task_id,config.num_local_schedulers,sched_id,my_first_task,my_last_task);// 映射到全局任务ID范围my_first_task e.first_task_id;my_last_task e.first_task_id;}// 遍历任务范围将任务分配给工作单元轮询策略for (size_t i my_first_task; i my_last_task; i) {// 获取工作单元队列的下一个空闲位置size_t last_task_id worker_queue_next_free_task_pos[next_worker - my_first_worker];// 计算任务ID并写入工作单元队列st_relaxed_gpu_u64(config.worker_queues[next_worker][last_task_id % config.per_worker_queue_len],compute_task_id(iteration_num, i));// 更新工作单元队列的最新就绪任务IDatom_add_release_gpu_u64(config.worker_queue_last_ready_task_id[next_worker], 1);// 轮询切换到下一个工作单元next_worker (next_worker my_last_worker - 1) ? my_first_worker: next_worker 1;}}// 移动到下一个事件cur_event_pos[queue_idx] 1;}}}2.4 Worker实现execute_worker 是 Mirage 持久化内核MPK中工作单元的核心设备端函数运行于 GPU 线程块上负责任务的获取、依赖检查、执行及后续事件触发是 MPK 内核中实际承载计算任务的核心组件。其流程逻辑可分为五大核心阶段形成 “任务获取 — 数据准备 — 依赖等待 — 任务执行 — 事件触发” 的完整闭环2.4.1 功能每个block对应一个worker。execute_worker函数负责从任务队列中获取任务并执行主要功能如下初始化和设置。函数首先初始化线程块内的共享内存存储任务 ID 和任务描述并根据当前工作单元 ID 绑定对应的本地任务队列在多 GPU 场景下额外绑定远程队列以接收跨 GPU 任务。同时完成性能分析器的初始化若启用为后续任务执行监控做好准备。此阶段的核心目标是建立工作单元与任务队列的关联明确任务来源。任务获取循环。任务获取。会等待任务队列中有可用任务使用原子操作确保线程安全从共享内存中加载任务描述。仅线程块内第 0 个线程执行任务获取逻辑通过轮询监听本地 / 远程队列使用 acquire 语义读取队列的最新就绪任务位置若队列无新任务则短暂休眠并切换队列若有新任务则读取任务 ID并通过共享内存同步给线程块内其他线程。随后线程块内线程并行将任务描述从全局内存拷贝到共享内存提升访问效率并通过线程块同步确保数据拷贝完成。函数主体是无限循环持续从任务队列中获取任务。其中线程0threadIdx.x 0负责协调其它线程协助数据加载。while(true) {// 获取下一个任务。// 执行任务// 触发事件}任务依赖检查阶段任务依赖检查。检查任务是否有依赖事件如果有依赖则等待依赖事件完成足够的触发次数使用事件计步器进行同步。仅第 0 个线程检查当前任务是否存在依赖事件若存在有效依赖事件计算该任务所需的事件触发次数循环等待直至事件触发次数满足需求使用 acquire 语义读取事件计数器确保数据可见性。此阶段通过精准的依赖管理保证任务执行的顺序正确性避免因数据未就绪导致的计算错误。任务执行根据任务类型执行差异化逻辑终止任务TASK_TERMINATE直接退出工作单元循环结束执行任务图开始任务TASK_BEGIN_TASK_GRAPH无实际计算逻辑仅作为流程标记NVSHMEM 拷贝任务TASK_NVSHMEM_COPY调用分布式通信接口完成跨 GPU 数据传输并触发远程事件信号归约任务TASK_REDUCE调用专用归约内核处理多 GPU 场景下的二维归约计算其他任务通过通用任务执行函数_execute_task(task_desc, config)事件触发与队列更新阶段任务执行完成后仅第 0 个线程触发对应的事件本地事件原子累加事件计数器若触发次数达到当前迭代需求将事件加入调度器队列大规模任务事件加入全局队列其他事件随机分配给本地调度器并通过 CAS 操作更新调度器队列的就绪状态远程事件依赖 NVSHMEM 拷贝任务在数据传输时自动触发信号此处仅打印调试日志。2.4.2 流程流程图如下mirage-5-42.4.3 代码execute_workerexecute_worker的代码如下。// 设备端函数执行工作单元逻辑Worker// 功能从工作单元队列读取任务处理任务依赖执行具体任务逻辑并触发后续事件__device__ void execute_worker(RuntimeConfig config) {// 共享内存变量存储当前待执行的任务ID和任务描述线程块内线程共享__shared__ TaskId cur_task_id;__shared__ TaskDesc task_desc;// 性能分析相关初始化仅MPK_ENABLE_PROFILING宏启用时执行#ifdef MPK_ENABLE_PROFILINGPROFILER_CLOSURE_PARAMS_DECL; // 声明性能分析闭包参数// 初始化性能分析器传入缓冲区指针、GPU ID等参数仅线程块内第0个线程执行初始化PROFILER_INIT(static_castuint64_t *(config.profiler_buffer),0,1,(threadIdx.x % 128 0));#endif// 当前工作单元的ID由线程块索引blockIdx.x标识每个工作单元对应一个线程块int worker_id blockIdx.x;// 共享内存变量存储工作单元队列指针及对应的队列ID支持本地和远程队列__shared__ TaskId *worker_queues[2];__shared__ int worker_queue_ids[2];// 获取当前工作单元的本地队列指针TaskId *local_worker_queue config.worker_queues[worker_id];worker_queues[0] local_worker_queue;worker_queue_ids[0] worker_id;int num_worker_queues 1; // 工作单元队列数量初始为1个本地队列// 多GPU场景下额外绑定远程队列用于接收其他GPU的任务if (config.num_gpus 1) {TaskId *remote_worker_queue config.worker_queues[worker_id config.num_workers];worker_queues[num_worker_queues] remote_worker_queue;worker_queue_ids[num_worker_queues] worker_id config.num_workers;num_worker_queues; // 队列数量增至2本地远程}// 记录每个队列当前处理的任务位置线程私有变量size_t cur_task_pos[2];// 共享内存变量记录每个队列最新的任务位置线程块内同步__shared__ size_t last_task_pos[2];// 初始化当前任务位置为0