flashinfer.mhc.mhc_pre_big_fuse_with_prenorm

flashinfer.mhc.mhc_pre_big_fuse_with_prenorm(dot_mix: Tensor, residual: Tensor, mhc_scale: Tensor, mhc_base: Tensor, 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, block_size: int = 0) → tuple[Tensor, Tensor, Tensor]

Apply mHC pre-map big-fuse and compute RMS sqrsum from residual.

This matches the Agentic mhc_pre_finalize boundary when no precomputed sqrsum is supplied. dot_mix may be shaped [..., 24] or [1, ..., 24].

Parameters:

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

  • residual (torch.Tensor) – Multi-head residual tensor, shape [..., HC=4, H] (bfloat16). sqrsum is computed from this tensor inside the kernel; 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.

  • 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.

  • 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]