flashinfer.mhc.mhc_post

flashinfer.mhc.mhc_post(x: Tensor, residual: Tensor, post_layer_mix: Tensor, comb_res_mix: Tensor) → Tensor

Apply mHC post mapping for HC=4.

out[..., new_hc, h] = x[..., h] * post_layer_mix[..., new_hc] + sum_old residual[..., old_hc, h] * comb_res_mix[..., old_hc, new_hc]

Parameters:

  • x (torch.Tensor) – Layer output tensor, shape [..., H] (matches residual’s outer dimensions and H = residual.shape[-1]).

  • residual (torch.Tensor) – Multi-head residual tensor, shape [..., HC=4, H]. The HC=4 axis must equal 4 (mHC is currently hard-wired to 4 sub-heads).

  • post_layer_mix (torch.Tensor) – Layer-output mixing weights, shape [..., HC=4] or [..., HC=4, 1] (both accepted). Provides the per-new-head coefficient applied to x (the post_layer_mix factor in the formula above). See mHC paper / reference implementation for the exact training-time derivation.

  • comb_res_mix (torch.Tensor) – Residual combination matrix, shape [..., HC=4, HC=4]. Each [old_hc, new_hc] entry is the coefficient mixing residual[..., old_hc, :] into the new head new_hc (the comb_res_mix factor above). See mHC paper / reference implementation for the training-time derivation.

Returns:

Output tensor with the same shape as residual ([..., HC=4, H]).

Return type:

torch.Tensor