jax.experimental.pallas.mosaic_gpu.emit_pipeline#

jax.experimental.pallas.mosaic_gpu.emit_pipeline(body, *, grid, in_specs=(), out_specs=(), max_concurrent_steps=1, delay_release=0)[原始碼]#

建立一個函式以在 Pallas 核心中發射手動管線。

參數:

  • body (Callable[..., None]) – 管線主體。

  • grid (pallas_core.StaticGrid) – 用於管線的網格。

  • in_specs (Sequence[pallas_core.BlockSpec]) – 輸入的區塊規格。

  • out_specs (Sequence[pallas_core.BlockSpec]) – 輸出的區塊規格。

  • max_concurrent_steps (int) – 同時處於活動狀態的最大循序階段數。預設為 1。

  • delay_release (int) – 在重複使用輸入/輸出參考之前等待的步數。預設為 0，且必須嚴格小於 max_concurrent_steps。一般來說，如果不在主體中等待 WGMMA，則會希望將其設定為 1。