mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
143 lines
4.8 KiB
Python
143 lines
4.8 KiB
Python
"""MHC modeling helpers that only DeepInfer needs.
|
|
|
|
`mhc_head` and the fused inference path `mhc_pre_big_fuse` combine the MHC operators
|
|
from `mhc_ops`; both compositions are inference-only and have no other caller.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import torch.nn.functional as F # noqa: N812
|
|
from tile_kernels.mhc.pre_big_fuse_kernel import mhc_pre_big_fuse_fwd
|
|
|
|
from deepinfer.model_executor.models.v4_modules.mhc_ops import (
|
|
default_n_splits,
|
|
mhc_head_compute_mix,
|
|
mhc_pre_apply_mix,
|
|
mhc_pre_norm_fn,
|
|
mhc_prenorm_gemm,
|
|
)
|
|
|
|
|
|
def mhc_head(
|
|
residual: torch.Tensor,
|
|
fn: torch.Tensor,
|
|
scale: torch.Tensor,
|
|
base: torch.Tensor,
|
|
*,
|
|
norm_weight: torch.Tensor | None = None,
|
|
norm_eps: float = 1e-6,
|
|
mhc_mult: int = 4,
|
|
pre_eps: float = 1e-6,
|
|
n_splits: int | None = None,
|
|
) -> torch.Tensor:
|
|
"""MHC head processing for the final language model head.
|
|
|
|
Combines pre_norm_fn + head_compute_mix + pre_apply_mix into a single call.
|
|
The fn parameter follows the same convention as block-level fn:
|
|
[mhc_mult, mhc_mult * hidden_size], which is padded to [mhc_mult * (mhc_mult + 2), ...]
|
|
internally to reuse the pre_norm_fn kernel.
|
|
|
|
Args:
|
|
residual: MHC residual tensor of shape [..., mhc_mult, hidden_size]
|
|
fn: weight matrix of shape [mhc_mult, mhc_mult * hidden_size]
|
|
scale: sigmoid scaling of shape [1] or scalar
|
|
base: mix biases of shape [mhc_mult]
|
|
norm_weight: optional RMSNorm weight of shape [mhc_mult * hidden_size]
|
|
norm_eps: epsilon for RMSNorm
|
|
mhc_mult: number of hyper-connection heads (currently only 4 is guaranteed to work)
|
|
pre_eps: epsilon for pre-layer mix sigmoid
|
|
n_splits: number of splits for the split-K GEMM; the platform default applies
|
|
|
|
Returns:
|
|
layer_input: tensor of shape [..., hidden_size], input to lm_head
|
|
"""
|
|
n_splits = default_n_splits() if n_splits is None else n_splits
|
|
mhc_mult3 = mhc_mult * (2 + mhc_mult)
|
|
|
|
if fn.shape[0] < mhc_mult3:
|
|
fn = F.pad(fn, (0, 0, 0, mhc_mult3 - fn.shape[0]))
|
|
|
|
mixes = mhc_pre_norm_fn(residual, fn, norm_weight, norm_eps, n_splits=n_splits)
|
|
|
|
mixes = mixes[..., :mhc_mult].contiguous()
|
|
|
|
if scale.numel() == 1:
|
|
scale = scale.reshape(1)
|
|
mix = mhc_head_compute_mix(mixes, scale, base, pre_eps)
|
|
|
|
return mhc_pre_apply_mix(residual, mix.unsqueeze(-1))
|
|
|
|
|
|
def mhc_pre_big_fuse(
|
|
residual: torch.Tensor,
|
|
fn: torch.Tensor,
|
|
mhc_scale: torch.Tensor,
|
|
mhc_base: torch.Tensor,
|
|
rms_eps: float,
|
|
mhc_pre_eps: float,
|
|
mhc_sinkhorn_eps: float,
|
|
mhc_post_mult_value: float,
|
|
sinkhorn_repeat: int,
|
|
n_splits: int | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""Fused inference path: prenorm GEMM + reduce/rmsnorm + split mixes + sinkhorn + apply mix.
|
|
|
|
Only used with gradients disabled: the GEMM is called directly and the epilogue stays
|
|
in the tile-kernels operator. The FP32 fn matrix already includes the norm weight.
|
|
"""
|
|
n_splits = default_n_splits() if n_splits is None else n_splits
|
|
assert residual.dtype == torch.bfloat16
|
|
assert fn.dtype == torch.float32
|
|
assert mhc_scale.dtype == torch.float32
|
|
assert mhc_base.dtype == torch.float32
|
|
|
|
mhc_mult = residual.shape[-2]
|
|
hidden_size = residual.shape[-1]
|
|
mhc_mult2 = mhc_mult * mhc_mult
|
|
mhc_mult3 = mhc_mult * 2 + mhc_mult2
|
|
|
|
mhc_hidden_size = mhc_mult * hidden_size
|
|
assert fn.shape[0] == mhc_mult3
|
|
assert fn.shape[1] == mhc_hidden_size
|
|
assert mhc_scale.shape == (3,)
|
|
assert mhc_base.shape == (mhc_mult3,)
|
|
|
|
outer_shape = residual.shape[:-2]
|
|
|
|
residual_flat = residual.view(-1, mhc_mult, hidden_size)
|
|
num_tokens = residual_flat.shape[0]
|
|
|
|
post_mix = torch.empty(num_tokens, mhc_mult, dtype=torch.float32, device=residual.device)
|
|
comb_mix = torch.empty(num_tokens, mhc_mult2, dtype=torch.float32, device=residual.device)
|
|
layer_input = torch.empty(num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device)
|
|
|
|
gemm_out_mul, gemm_out_sqrsum = mhc_prenorm_gemm(residual_flat, fn, n_splits)
|
|
|
|
actual_n_splits = gemm_out_mul.shape[0]
|
|
gemm_out_mul = gemm_out_mul.view(actual_n_splits, num_tokens, mhc_mult3)
|
|
gemm_out_sqrsum = gemm_out_sqrsum.view(actual_n_splits, num_tokens)
|
|
|
|
mhc_pre_big_fuse_fwd(
|
|
gemm_out_mul,
|
|
gemm_out_sqrsum,
|
|
mhc_scale,
|
|
mhc_base,
|
|
residual_flat,
|
|
post_mix,
|
|
comb_mix,
|
|
layer_input,
|
|
rms_eps,
|
|
mhc_pre_eps,
|
|
mhc_sinkhorn_eps,
|
|
mhc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
n_splits=n_splits,
|
|
)
|
|
|
|
post_mix = post_mix.view(*outer_shape, mhc_mult, 1)
|
|
comb_mix = comb_mix.view(*outer_shape, mhc_mult, mhc_mult)
|
|
layer_input = layer_input.view(*outer_shape, hidden_size)
|
|
|
|
return post_mix, comb_mix, layer_input
|