flashinfer.norm.rmsnorm_sigmoid_gate

flashinfer.norm.rmsnorm_sigmoid_gate(input: Tensor, gate: Tensor, eps: float = 1e-05, *, enable_pdl: bool = True, backend: Literal['auto', 'cake', 'cute'] = 'auto') → 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 launch with programmatic dependent launch. Defaults to True; pass False to disable it. The CuTe DSL kernel takes the attribute directly; the Cake package serves the call from its PDL route, whose windows name a PDL build, a late-signal PDL build or a plain build (launched without the attribute where the PDL build measured slower) or leave a token count to the CuTe DSL kernel under PDL (a declared gap). The package ships that route on SM100, SM103 and SM107.

  • backend (str) – "auto" (default) runs the Cake kernels (flashinfer.cake_muse_spark_norm_backend) where the package ships a program for the device, the call’s launch form, head count and token count (the package routes 128, 32 and 16 heads; its windows leave token-count gaps where the CuTe DSL kernel measured ahead) and the tensors sit in 32-byte aligned storage with token strides that are multiples of 16 elements (the package’s 16-head programs on SM107 read and write 256-bit vectors; other storage runs the CuTe DSL kernel), else the CuTe DSL kernel; "cake" and "cute" select one and raise ValueError when it is not available ("cake" at a head count or token count outside the package’s windows, or on storage the package does not serve, raises too).