mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-01 17:26:25 +08:00
39b086d781
The references were importable and numerically validated, but nothing pointed the model at them: the forward pass still called the compiled kernels, which raise ModuleNotFoundError on a host without TileLang. This adds the routing layer and swaps the call sites. `deepinfer/tk_compat/` exposes the *compiled* signatures -- out-buffers, in-place conventions, argument order -- because that is what the 25 call sites are written against, and adapts to the reference's shape internally. Selection is `DEEPINFER_KERNELS=auto|cuda|ref`; `auto` follows whether TileLang imports. The adapters are not pass-throughs. Each difference was found by running both backends: * `per_token_cast` takes `num_per_channels=n`, the reference takes `block_size=(1, n)`. * `apply_rotary` takes a concatenated cos|sin cache and rotates in place returning None; the reference takes cos and sin separately and returns a new tensor. The adapter splits, calls, and copies back, because the call site relies on the in-place contract. * `engram_gate` takes the pre-multiplied weight and `(eps, clamp_value)`; the reference needs the two weights separately and `(clamp_value, eps)`. Since a product cannot be un-multiplied, `fused_weight` returns the pair on the reference path and the gate unpacks it. * `mhc_pre_big_fuse_fwd` has no single reference: its five steps are composed from the reference's own components. * `mhc_reduce_partials_and_rmsnorm_fwd` likewise; rebuilt from the reference's formula. `moe_topk_gate(force_random=...)` raises rather than silently ignoring the flag. The reference has no equivalent, and dropping it would route tokens differently while looking healthy. Verification, all on the real call signatures: 18 cases compared between `cuda` and `ref` modes, `ALL_MATCH`, 12 bitwise identical including every quantised byte and every routing index; plus a smoke run that executes all 23 APIs with TileLang blocked, `ALL_CALLS_OK`. End to end still has not completed, and the reason is not this change: six GPUs on the host are held by other containers (82 GiB of 139 each), so the 8-way EP run OOMs in allocate_userbuffer and the surviving ranks wait at a barrier. Recorded in the doc along with the command to rerun it.
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 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,
|
|
)
|
|
from deepinfer.tk_compat.mhc import mhc_pre_big_fuse_fwd
|
|
|
|
|
|
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
|