← FLAGLANG / 研发专栏

用 GEMM + bias 读懂值式接口

减少重复声明,同时保持数值顺序、域身份和就绪关系。

从一段实际源码开始

下面摘自开发树的 gemm_values.py。完整示例先声明 TensorArg、静态域、Schedule 与外部 buffer;片段只展示计算主体,需要放在被捕获的函数体中,不能独立运行。

with F.ValueRegion(program, scope=(128,), transfer_scope=(32,)) as region:
    a = region.read(lhs)
    b = region.read(rhs)
    acc = region.contract(a, b, domain=output_domain,
                          iteration_domain=iteration_domain,
                          accumulator_dtype="f32")
    c = acc.astype("f16")
    bias_value = region.read(bias)
    result = F.map(lambda x, b: x + b, c, bias_value)
    region.write(result, output)

read 返回可消费的值;contract 的输入域来自值,但输出域和归约域仍由作者声明。map 再根据输入推导共同域和 dtype。两个域即使大小相同,也不会因此自动成为同一个名义域。

一次 cast 决定了哪个数学程序

示例先用 f32 累加 GEMM,再把结果转换为 f16,最后与 f16 bias 相加并写出。它的参考计算是:

C = f16(f16(f32(A) @ f32(B)) + bias)

若先在 f32 中加 bias 再统一转 f16,舍入位置就变了。两者在实数代数中相似,在有限精度中却可能给出不同结果。值式接口缩短了写法,但不应把这种语义变化当作无条件优化。

更短的代码并没有消除执行关系

ValueRegion 收集构造记录,退出区域时分阶段展开:先声明逻辑 buffer 和 storage,再产生传输、填充、收缩、转换、逐点计算与写回,最后加入完成等待、共享发布和依赖边。结果仍进入原有 Schedule 与编译流程。

这种展开减少作者重复描述可推导信息的负担。它没有承诺最少存储、最少 barrier 或全局最优流水线,也没有把私有构造记录提升为新的用户输入格式。

异步读:完成关系要显式消费

pending = region.read_async(lhs)
a = region.wait(pending)

pending 本身不是一个已就绪的值。源码检查拒绝直接消费它、重复等待、跨区域使用以及带未等待读取退出区域。调用方必须使用 wait 返回的新句柄。

这里描述的是设备计算依赖,不是 Python 主机线程是否阻塞。即便依赖合法,也还需要独立的调度、生成代码和测量证据,才能声称访存与计算发生了有效重叠。

示例的支持边界

VALUES 文档记录的正向 GEMM+bias 范围是 M=N=64、K=32/64;M=128+bias 明确拒绝。普通 GEMM 接受的形状不能直接推广到其后处理组合。

这组 alpha 接口面向静态整域与直线区域;广播、动态边界、跨 CTA 推导和原地更新不在该文档保证范围。源码中存在某个标量运算符入口,也不说明每个后端都能降级相应表达。

本篇核对了值式实现与错误路径测试源码,没有重跑 H20。冻结预览包和当前开发树的新接口仍需按各自版本使用。

源码依据

基于本地开发快照撰写。下列为源码中的相对路径与仓库提交号,不是公开下载链接;本文没有重新执行 GPU 测试,也不代表冻结发布包的全部能力。

  • flaglang/VALUES.mdd589a008c7
  • flaglang/examples/preview/gemm_values.pyd589a008c7
  • flaglang/python/flaglang/values.pyd589a008c7
  • flaglang/test/test_values.pyd589a008c7
LET’S BUILD / CLASS OF 2027

一起把系统做深。

AI 编译器 · GPU Kernel · 大模型推理系统

邮件
联系 ↗