jax.experimental.pallas.mosaic_gpu.wgmma#
- jax.experimental.pallas.mosaic_gpu.wgmma(acc, a, b)[source]#
在给定引用上执行异步 warp group matmul-accumulate 操作。
从概念上讲，这等效于执行
acc[...] += a[...] @ b[...]，但计算是异步执行的。- 参数:
acc (gpu_core.WGMMAAbstractAccumulatorRef) – 累加器引用。需要通过调用
jax.experimental.pallas.run_scoped()并使用jax.experimental.pallas.mosaic_gpu.WGMMAAccumulatorRef()分配。a – 左侧操作数引用。
b – 右侧操作数引用。
- 返回类型:
None