flashinfer.norm.rmsnorm_sigmoid_gate

flashinfer.norm.rmsnorm_sigmoid_gate(input: Tensor, gate: Tensor, eps: float = 1e-05, *, enable_pdl: bool = True) → None

Per-head RMSNorm and sigmoid gate, in place, for head dim 64.

The last dimension is normalized independently, multiplied by sigmoid(gate), and written back to input. The optimized kernel supports contiguous BF16 head rows of width 64 on SM100, SM103, and SM107 GPUs.

Parameters:

• input (torch.Tensor) – BF16 tensor of shape (rows, 64) or (rows, heads, 64), normalized in place. The last dimension must be contiguous, heads must be packed, and the row stride must be a multiple of 8 elements.

• gate (torch.Tensor) – BF16 gate tensor with the same shape and layout constraints as input.

• eps (float) – Epsilon added to the RMS denominator.

• enable_pdl (bool) – Whether to enable programmatic dependent launch. Defaults to True; pass False to disable it.