flashinfer.comm.CakeFp4AllReduceFusedNormSpec

class flashinfer.comm.CakeFp4AllReduceFusedNormSpec(weight: Tensor, eps: float = 1e-06, weight_bias: float = 1.0, output: CakeFp4AllReduceFusedNormOutput = CakeFp4AllReduceFusedNormOutput.BF16, global_scale: Tensor | None = None)

Request of the fused epilogue: residual add, then RMSNorm scaled by weight_bias + weight.

weight is a contiguous bfloat16 [6144] tensor on the workspace device, eps the RMSNorm epsilon and weight_bias the constant added to the weight (Gemma: 1.0); both are runtime scalars. output is the precision of the normalised result. global_scale is required for the NVFP4 output: a float32 tensor of at least one element on the device holding the quantiser’s input global scale (448 * 6 / amax(y) in the convention of fp4_quantize); the kernel reads it, the host never does, so the call stays graph capturable.

__init__(weight: Tensor, eps: float = 1e-06, weight_bias: float = 1.0, output: CakeFp4AllReduceFusedNormOutput = CakeFp4AllReduceFusedNormOutput.BF16, global_scale: Tensor | None = None) → None

Methods

__init__(weight[, eps, weight_bias, output, ...])

Attributes

eps
global_scale
output
weight_bias
weight