flashinfer.mhc.mhc_pre_big_fuse

flashinfer.mhc.mhc_pre_big_fuse(dot_mix: Tensor, sqrsum: Tensor, residual: Tensor, mhc_scale: Tensor, mhc_base: Tensor, k: int, rms_eps: float = 1e-06, mhc_pre_eps: float = 1e-06, mhc_sinkhorn_eps: float = 1e-06, mhc_post_mult_value: float = 1.0, sinkhorn_repeat: int = 20, num_splits: int = 1, block_size: int = 0) → tuple[Tensor, Tensor, Tensor]

Apply mHC pre-map big-fuse using external projection and sqrsum.

dot_mix contains raw projection logits laid out as [pre(4), post(4), comb(16)]. sqrsum contains the corresponding residual-square sums used for RMS normalization. When num_splits > 1, both inputs have a leading split dimension that is reduced inside the CUDA kernel.

Parameters:

  • dot_mix (torch.Tensor) – Raw projection logits, shape [..., 24] when num_splits=1 or [num_splits, ..., 24] when num_splits > 1. The trailing 24 = 4 (pre) + 4 (post) + 16 (comb) slots are partitioned and re-shaped to per-head pre / post / comb factors inside the kernel. See mHC paper / reference implementation for the projection derivation.

  • sqrsum (torch.Tensor) – Pre-computed per-token residual square-sum used for RMS normalization. Shape [...] matching dot_mix outer dims, or [num_splits, ...] when num_splits > 1 (the kernel reduces the split axis internally). dtype is implementation-defined.

  • residual (torch.Tensor) – Multi-head residual tensor, shape [..., HC=4, H] (bfloat16). Same layout as mhc_post()’s residual.

  • mhc_scale (torch.Tensor) – mHC scaling tensor consumed by the Sinkhorn / mix step. Shape and dtype are implementation-defined; see mHC paper / reference implementation for the exact semantics.

  • mhc_base (torch.Tensor) – mHC bias / base tensor consumed by the Sinkhorn / mix step. Shape and dtype are implementation-defined; see mHC paper / reference implementation for the exact semantics.

  • k (int) – mHC algorithm parameter k passed to the CUDA kernel. See mHC paper / reference implementation for the precise semantics.

  • rms_eps (float) – RMSNorm epsilon for numerical stability. Default 1e-6. Must be strictly positive.

  • mhc_pre_eps (float) – Numerical-stability epsilon used in the mHC pre-map step. Default 1e-6. Must be strictly positive.

  • mhc_sinkhorn_eps (float) – Numerical-stability epsilon used in the Sinkhorn iteration step. Default 1e-6. Must be strictly positive.

  • mhc_post_mult_value (float) – Post-mix multiplicative factor applied to the mHC post outputs. Default 1.0.

  • sinkhorn_repeat (int) – Number of Sinkhorn iterations. Default 20.

  • num_splits (int) – Split factor along the leading dimension. Must be one of {1, 2, 4, 8, 16}. When > 1, dot_mix and sqrsum carry a leading split axis that is reduced inside the kernel. Default 1.

  • block_size (int) – CUDA block-size hint forwarded to the kernel (0 selects the kernel’s default). Default 0.

Returns:

(post_mix, comb_mix, layer_input) tensors with shapes [..., HC=4, 1], [..., HC=4, HC=4] and [..., H] respectively. See mHC paper / reference implementation for how downstream layers consume them.

Return type:

Tuple[torch.Tensor, torch.Tensor, torch.Tensor]