Kernel 运行时¶
LogicalGraph 中的 Operator 只携带元信息。执行时,框架把每个 Operator 转化为若干 Kernel(最小执行单元),分配到对应设备的 KernelStream(异步执行流)上运行,由 Blob(持有 torch::Tensor)承接数据读写。完整链路:
构图: 推导元信息 + 确定算子要在哪些设备上执行
│
│ 为每个 (设备, 流类型) 创建一个 Kernel,准备输入/输出 Blob
▼
Kernel ──► 推入 KernelStream 的队列,立即返回(异步)
│
▼
流线程: 循环取出 Kernel → 执行 → 立即释放(降低峰值显存)
1. Blob — 张量数据的物理容器¶
Blob 持有 torch::Tensor,位于 Operand(元信息)与 Kernel(计算)之间,把逻辑数据节点关联到物理内存。分布式张量在每个设备上各有一个 Blob,按全局设备 ID 索引;Kernel 只访问本设备的 Blob,不跨设备读写。
2. Kernel — 最小执行单元¶
Kernel 持有父 Operator、本设备的输入/输出 Blob 以及所属流。Run() 依次完成:从输入 Blob 取出 torch::Tensor → 调用 Compute()(默认委托给 LibTorch 算子)→ 结果写回输出 Blob。部分 Kernel 有特化实现,部分还能在 Kernel 中调用 Python 代码(见 Python Kernel)。
3. KernelStream — 异步执行流¶
KernelStream 封装一个设备上的一条执行流(CPU 线程或 CUDA Stream),由专用线程按序执行 Kernel 队列。流分为两类:
| 类型 | 用途 |
|---|---|
kCompute |
常规算子计算(matmul、relu、conv2d 等) |
kCommunicate |
集合通信(all-reduce、all-gather 等) |
同一设备可同时拥有两类流,利用 CUDA Stream 并发实现计算与通信重叠。
执行是异步的:Kernel 被推入流的队列后立即返回,流线程循环取 Kernel 执行、执行完立即释放。需要同步时(如取值),等待流中已有 Kernel 全部完成。
4. 一个 Operator → 多个 Kernel¶
一个 Operator 需要在其输入/输出 Operand 涉及的每个设备上执行,因此会实例化为多个 Kernel:构图时框架根据 DeviceMesh 确定算子要执行的设备集合;执行时为每个设备创建一个 Kernel,各 Kernel 只负责本设备的计算。
Operator(DeviceMesh = {GPU:0, GPU:1})
│
├──► Kernel(GPU:0) ──► KernelStream(GPU:0) ← 只访问 GPU 0 的数据
└──► Kernel(GPU:1) ──► KernelStream(GPU:1) ← 只访问 GPU 1 的数据
5. 源文件索引¶
| 文件 | 说明 |
|---|---|
dtorch/core/blob.h .cc |
Blob — 张量数据物理容器 |
dtorch/core/kernel/kernel.h .cc |
Kernel — 最小执行单元 |
dtorch/core/kernel_stream/kernel_stream.h .cc |
KernelStream — 异步执行流 |
dtorch/core/kernel_stream/cpu_kernel_stream.h / cuda_kernel_stream.h .cc |
CPU / CUDA 流的实现 |
dtorch/core/kernel_stream/kernel_stream_manager.h .cc |
流管理器 |
dtorch/core/operators/operator_assign_info.h |
算子到流的分配信息 |
dtorch/core/runner/node_runner_base.h / naive_runner.cc |
Runner — 创建并分派 Kernel |