97 lines
3.3 KiB
Python
97 lines
3.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import torch
|
|
|
|
from vllm import _custom_ops as ops
|
|
|
|
|
|
def mhc_pre_cpu(
|
|
residual: torch.Tensor,
|
|
fn: torch.Tensor,
|
|
hc_scale: torch.Tensor,
|
|
hc_base: torch.Tensor,
|
|
rms_eps: float,
|
|
hc_pre_eps: float,
|
|
hc_sinkhorn_eps: float,
|
|
hc_post_mult_value: float,
|
|
sinkhorn_repeat: int,
|
|
n_splits: int = 1,
|
|
norm_weight: torch.Tensor | None = None,
|
|
norm_eps: float = 0.0,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""CPU-ported mHC pre block (see `mhc_pre_torch` for the eager reference).
|
|
|
|
The ported kernel (`hc_pre_fused_cpu`) only exposes one merged `hc_eps`
|
|
(used for both `hc_pre_eps`/`hc_sinkhorn_eps`) and hardcodes the post-mix
|
|
multiplier to 2.0 -- true of every real call site in
|
|
`models/deepseek_v4/cpu/model.py`, so this is not a capability loss here.
|
|
"""
|
|
assert n_splits == 1, "mhc_pre_cpu does not support n_splits != 1"
|
|
assert hc_pre_eps == hc_sinkhorn_eps, (
|
|
"mhc_pre_cpu requires hc_pre_eps == hc_sinkhorn_eps (single merged hc_eps)"
|
|
)
|
|
assert hc_post_mult_value == 2.0, "mhc_pre_cpu hardcodes the post-mix multiplier"
|
|
|
|
hc_mult, hidden_size = residual.shape[-2:]
|
|
outer_shape = residual.shape[:-2]
|
|
x_flat = residual.reshape(-1, hc_mult, hidden_size)
|
|
|
|
layer_input, post, comb = ops.hc_pre_fused_cpu(
|
|
x_flat,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
hc_mult,
|
|
sinkhorn_repeat,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
)
|
|
|
|
post_mix = post.view(*outer_shape, hc_mult, 1)
|
|
comb_mix = comb.view(*outer_shape, hc_mult, hc_mult)
|
|
layer_input = layer_input.view(*outer_shape, hidden_size)
|
|
return post_mix, comb_mix, layer_input
|
|
|
|
|
|
def mhc_post_cpu(
|
|
x: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
post_layer_mix: torch.Tensor,
|
|
comb_res_mix: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""CPU-ported mHC post block (see `mhc_post_torch` for the eager reference)."""
|
|
hc_mult, hidden_size = residual.shape[-2:]
|
|
outer_shape = residual.shape[:-2]
|
|
|
|
x_flat = x.reshape(-1, hidden_size)
|
|
residual_flat = residual.reshape(-1, hc_mult, hidden_size)
|
|
post_flat = post_layer_mix.reshape(-1, hc_mult).float()
|
|
comb_flat = comb_res_mix.reshape(-1, hc_mult, hc_mult).float()
|
|
|
|
out = ops.hc_post_fused_cpu(x_flat, residual_flat, post_flat, comb_flat)
|
|
return out.view(*outer_shape, hc_mult, hidden_size)
|
|
|
|
|
|
def hc_head_fused_cpu(
|
|
hidden_states: torch.Tensor,
|
|
hc_fn: torch.Tensor,
|
|
hc_scale: torch.Tensor,
|
|
hc_base: torch.Tensor,
|
|
rms_norm_eps: float,
|
|
hc_eps: float,
|
|
) -> torch.Tensor:
|
|
"""CPU-ported HC head reduction (see `test_hc_head_cpu` in
|
|
tests/kernels/test_mhc_kernels.py for the eager reference this is tested
|
|
against).
|
|
|
|
The ported kernel's C++ signature takes `(hc_eps, norm_eps)` -- the
|
|
opposite order from this wrapper's `(rms_norm_eps, hc_eps)`, which matches
|
|
the real call site in `models/deepseek_v4/cpu/model.py`.
|
|
"""
|
|
hc_mult, hidden_size = hidden_states.shape[-2:]
|
|
outer_shape = hidden_states.shape[:-2]
|
|
hs_flat = hidden_states.reshape(-1, hc_mult, hidden_size)
|
|
|
|
out = ops.hc_head_fused_cpu(hs_flat, hc_fn, hc_scale, hc_base, hc_eps, rms_norm_eps)
|
|
return out.view(*outer_shape, hidden_size)
|