Skip to main content

Crate pagesift_gpu

Crate pagesift_gpu 

Source
Expand description

NVIDIA GPU 推理后端:给 pagesift-infer 的重算子(矩阵乘、卷积、激活)提供同语义的 GPU 版本。

  • 不依赖 CUDA 工具链:运行时 dlopen("libcuda.so.1") 调驱动 API(见 driver.rs);
  • 内核是仓库里 kernels/ 用 nightly 的 nvptx 目标编出的 PTX,include_str! 嵌进来,由驱动 JIT。 有两份:sm_70 原版给算力 ≥ 7.0 的卡,.target 改成 sm_61 的一份给 Pascal(6.1、6.2);
  • 没有驱动、没有 N 卡、算力不够时 Gpu::open 返回错误,调用方退回 CPU。

结果与 pagesift_infer::ops 的 CPU 实现只差浮点舍入(累加顺序不同、GPU 用 FMA)。

let gpu = pagesift_gpu::Gpu::open()?;
let a = pagesift_infer::Tensor::new(vec![2, 3], vec![1.0; 6]);
let b = pagesift_infer::Tensor::new(vec![3, 4], vec![1.0; 12]);
let c = gpu.matmul(&a, &b)?; // [2, 4],行优先
assert_eq!(c, vec![3.0; 8]);

Structs§

BatchExec
一次批量前向的执行器。一批一个,用完检查 BatchExec::finish。
ConvShape
一次卷积的全部尺寸,与 pagesift_infer::ops::conv2d 的参数一一对应: 输入 [c, h, w],权重 [o, c/groups, kh, kw]。
Costs
实测的开销,单位都是秒。
Ctx
工人线程上的长期状态:显存池和权重缓存。
DevTensor
显存里的一个批量张量。
DeviceBuffer
显存里的一段 f32。借用创建它的 Gpu,保证在上下文释放前先释放。
Gpu
一块 N 卡上的上下文、载好的内核模块和函数表。所有方法取 &self,可以跨线程共享。
GpuBackend
透明分流的后端。
GpuInfo
一块 N 卡的情况(给 doctor 和模型选择用)。
Timer
用一对 CUDA 事件量设备上一段工作的耗时(只算 GPU 执行,不含主机侧开销)。
Worker
工人句柄,丢掉时线程退出并释放上下文。

Enums§

WaitError
Worker::call_timeout 没等到结果的原因。

Constants§

OOM_TAG
显存不够的错误都带这个前缀,调用方据此减小批次,而不是放弃 GPU。

Functions§

device_count
本机可见的 N 卡数量;没有驱动或驱动初始化失败时为 0。用来区分“没有 GPU”(该退回 CPU)和 “有 GPU 但 Gpu::open 失败”(PTX、驱动出了问题,应当报出来)。
disabled_by_env
PAGESIFT_GPU=off 时返回原因字符串:这时 Gpu::open、Gpu::open_index、probe 都报错, 调用方照常退回 CPU。
is_oom
错误是否是显存不足(减小批次还有救)。
probe
查第 index 块 N 卡的情况,不载内核(不 JIT,比 Gpu::open_index 轻)。 没有驱动、没有这块卡、算力不够或 PAGESIFT_GPU=off 时,错误就是可以给用户看的原因。

Type Aliases§

Owner
拥有权重的对象(模型)。权重缓存按地址查,缓存项持有它,地址在缓存存续期间就一直有效。
SharedWorker
把工人包成可以共享的句柄。