jax.experimental.pallas.mosaic_gpu 模块

Pallas 针对 H100 的实验性 GPU 后端。

这些 API 非常不稳定，每周都可能发生变化。使用风险自负。

类

Barrier(num_arrivals[, num_barriers])
GPUBlockSpec([block_shape, index_map, ...])
GPUCompilerParams(*[, approx_math, ...])	Mosaic GPU 编译器参数。
GPUMemorySpace(value)	一个枚举。
Layout(value)	一个枚举。
SwizzleTransform(swizzle)
TilingTransform(tiling)	表示内存引用的平铺变换。
TransposeTransform(permutation)	转置一个平铺的 memref。
WGMMAAccumulatorRef(shape, dtype, _init)

函数

barrier_arrive(barrier)	到达给定的屏障。
barrier_wait(barrier)	等待给定的屏障。
commit_smem()	提交对 SMEM 的所有写入，使其对加载、TMA 和 WGMMA 可见。
copy_gmem_to_smem(src, dst, barrier)	将 GMEM 引用异步复制到 SMEM 引用。
copy_smem_to_gmem(src, dst[, predicate])	将 SMEM 引用异步复制到 GMEM 引用。
emit_pipeline(body, *, grid[, in_specs, ...])	创建一个函数以在 Pallas 内核中发出手动流水线。
layout_cast(x, new_layout)	转换给定数组的布局。
set_max_registers(n, *, action)	设置一个 warp 拥有的最大寄存器数。
wait_smem_to_gmem(n[, wait_read_only])	等待直到正在进行的 SMEM->GMEM 复制不超过 n 个。
wgmma(acc, a, b)	在给定引用上执行异步 warp 组矩阵乘法累加。
wgmma_wait(n)	等待直到正在进行的 WGMMA 操作不超过 n 个。

别名

ACC	WGMMAAccumulatorRef的别名
GMEM	jax.experimental.pallas.mosaic_gpu.GPUMemorySpace.GMEM 的别名。
SMEM	jax.experimental.pallas.mosaic_gpu.GPUMemorySpace.SMEM 的别名。