flashinfer.comm.CakeFp4AllReduceWorkspace

class flashinfer.comm.CakeFp4AllReduceWorkspace(group: ProcessGroup | None = None, *, family: str = 'allreduce', input_format: str = 'bf16', capacity: int = 16384, device: device | None = None)

Per-rank state of the NVFP4-transport all-reduce, created collectively.

Every rank of group constructs one instance with the same family, input_format and capacity. family is "allreduce" (world size 2 or 4; the sum of the partials) or "rmsnorm" (world size 2; the sum fused with the residual add and the RMSNorm, any output format per call); input_format is "bf16" (activations) or "nvfp4" (pre-quantised records). capacity is the largest row count a call may carry (at most 103000 for world size 4). The receive windows are shared through CUDA IPC, armed with their Lamport sentinels and published in the peer window table; the control window (the registry’s control words: one four-word call-state slot per CTA) starts at zero. The CTA count of every launch is the SM count of the device, at most the grid bound of the workspace’s programs (the CTAs their control window provisions for), and must agree on every rank. Construction loads (and on first use compiles) every program this workspace can launch for the device’s exact target. destroy() releases the shared windows and is collective as well.

__init__(group: ProcessGroup | None = None, *, family: str = 'allreduce', input_format: str = 'bf16', capacity: int = 16384, device: device | None = None) → None

Methods

__init__([group, family, input_format, ...])
destroy()	Release the shared receive windows (collective; a second call is a no-op).
module(program)	The built module of program on this workspace's target.
program(rows[, output_format])	The program a call with rows rows (and, for rmsnorm, output_format) runs on this rank.

Attributes

closed