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 (pallas_core.TransformedRef) – 右側運算元參考。
- 返回類型：:
無