Fangzhou AiandCursor
00972dfd72
Signed-off-by: Fangzhou Ai <fangzhou@semianalysis.ai> Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
977 lines
29 KiB
Python
977 lines
29 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import inspect
|
|
|
|
import torch
|
|
|
|
# this import will also register the custom ops
|
|
# import vllm.model_executor.kernels.mhc # noqa: F401
|
|
import vllm.model_executor.kernels.mhc as mhc_kernels
|
|
from vllm._aiter_ops import is_aiter_found_and_supported
|
|
from vllm.model_executor.custom_op import CustomOp
|
|
from vllm.platforms import current_platform
|
|
from vllm.utils.import_utils import has_tilelang
|
|
|
|
|
|
def _has_tilelang_mhc() -> bool:
|
|
if not has_tilelang():
|
|
return False
|
|
if current_platform.is_cuda():
|
|
return True
|
|
if current_platform.is_rocm():
|
|
from vllm.platforms.rocm import on_gfx942
|
|
|
|
# TileLang MHC currently produces incorrect results on gfx942. Keep
|
|
# gfx942 on the existing torch/triton fallbacks until that path is fixed.
|
|
return not on_gfx942()
|
|
return False
|
|
|
|
|
|
HAS_TILELANG_MHC = _has_tilelang_mhc()
|
|
HAS_AITER_MHC = is_aiter_found_and_supported()
|
|
|
|
|
|
def _has_aiter_mhc_fused() -> bool:
|
|
if not HAS_AITER_MHC:
|
|
return False
|
|
try:
|
|
from aiter.ops.mhc import mhc_fused_post_pre
|
|
except Exception:
|
|
return False
|
|
return callable(mhc_fused_post_pre)
|
|
|
|
|
|
def _aiter_mhc_op_accepts_norm(op_name: str) -> bool:
|
|
if not HAS_AITER_MHC:
|
|
return False
|
|
try:
|
|
from aiter.ops import mhc as aiter_mhc
|
|
|
|
for candidate_name in (op_name, f"{op_name}_fake"):
|
|
op = getattr(aiter_mhc, candidate_name, None)
|
|
if op is None:
|
|
continue
|
|
parameters = inspect.signature(op).parameters
|
|
if "norm_weight" in parameters and "norm_eps" in parameters:
|
|
return True
|
|
except (AttributeError, ImportError, TypeError, ValueError):
|
|
pass
|
|
return False
|
|
|
|
|
|
HAS_AITER_MHC_FUSED = _has_aiter_mhc_fused()
|
|
HAS_AITER_MHC_PRE_NORM = _aiter_mhc_op_accepts_norm("mhc_pre")
|
|
HAS_AITER_MHC_FUSED_NORM = _aiter_mhc_op_accepts_norm("mhc_fused_post_pre")
|
|
|
|
|
|
def _aiter_mhc_supported(
|
|
residual: torch.Tensor,
|
|
norm_weight: torch.Tensor | None,
|
|
*,
|
|
supports_norm: bool,
|
|
) -> bool:
|
|
hidden_size = residual.shape[-1]
|
|
hc_mult = residual.shape[-2]
|
|
return (
|
|
HAS_AITER_MHC
|
|
and hidden_size % 256 == 0
|
|
and hc_mult == 4
|
|
and (norm_weight is None or supports_norm)
|
|
)
|
|
|
|
|
|
def _apply_mhc_norm(
|
|
layer_input: torch.Tensor,
|
|
norm_weight: torch.Tensor | None,
|
|
norm_eps: float,
|
|
) -> torch.Tensor:
|
|
if norm_weight is None:
|
|
return layer_input
|
|
from vllm import ir
|
|
|
|
return ir.ops.rms_norm(layer_input, norm_weight, norm_eps)
|
|
|
|
|
|
# --8<-- [start:mhc_pre]
|
|
@CustomOp.register("mhc_pre")
|
|
class MHCPreOp(CustomOp):
|
|
"""MHC pre block.
|
|
|
|
Computes mix logits from RMS-normalized HC residual streams, then
|
|
returns post_mix, comb_mix, and
|
|
layer_input = sum_i pre_mix_i * residual_i.
|
|
"""
|
|
|
|
# --8<-- [end:mhc_pre]
|
|
@classmethod
|
|
def enabled(cls) -> bool:
|
|
return True
|
|
|
|
def forward_cuda(
|
|
self,
|
|
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]:
|
|
return torch.ops.vllm.mhc_pre_tilelang(
|
|
residual,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
n_splits,
|
|
norm_weight,
|
|
norm_eps,
|
|
)
|
|
|
|
def forward_hip(
|
|
self,
|
|
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]:
|
|
if _aiter_mhc_supported(
|
|
residual,
|
|
norm_weight,
|
|
supports_norm=HAS_AITER_MHC_PRE_NORM,
|
|
):
|
|
return torch.ops.vllm.mhc_pre_aiter(
|
|
residual,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
n_splits,
|
|
norm_weight,
|
|
norm_eps,
|
|
)
|
|
elif HAS_TILELANG_MHC:
|
|
return torch.ops.vllm.mhc_pre_tilelang(
|
|
residual,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
n_splits,
|
|
norm_weight,
|
|
norm_eps,
|
|
)
|
|
else:
|
|
post_mix, comb_mix, layer_input = self.forward_native(
|
|
residual,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
n_splits,
|
|
norm_weight,
|
|
norm_eps,
|
|
)
|
|
return (
|
|
post_mix,
|
|
comb_mix,
|
|
_apply_mhc_norm(layer_input, norm_weight, norm_eps),
|
|
)
|
|
|
|
def forward_native(
|
|
self,
|
|
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]:
|
|
return mhc_kernels.mhc_pre_torch(
|
|
residual,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
)
|
|
|
|
def forward_xpu(
|
|
self,
|
|
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]:
|
|
return torch.ops._xpu_C.mhc_pre(
|
|
residual,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
)
|
|
|
|
def forward_cpu(
|
|
self,
|
|
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]:
|
|
return mhc_kernels.mhc_pre_cpu(
|
|
residual,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
n_splits,
|
|
norm_weight,
|
|
norm_eps,
|
|
)
|
|
|
|
|
|
# --8<-- [start:mhc_pre_delayed]
|
|
@CustomOp.register("mhc_pre_delayed")
|
|
class MHCPreDelayedOp(CustomOp):
|
|
"""MHC pre block using the pre-mix carried from the previous sublayer.
|
|
|
|
Same gates as :class:`MHCPreOp`, but the stream collapse applies the
|
|
caller's ``pre_mix`` and this sublayer's pre-mix is returned for the next
|
|
sublayer seam.
|
|
|
|
Passing ``sublayer_out`` / ``post_layer_mix`` / ``comb_res_mix`` also
|
|
applies the preceding post block, which lets AITER fold it into the pre
|
|
projection and saves a launch. Returns residual, post_mix, comb_mix,
|
|
layer_input, next_pre_mix; ``residual`` is the caller's own tensor when
|
|
no post is requested.
|
|
"""
|
|
|
|
# --8<-- [end:mhc_pre_delayed]
|
|
@classmethod
|
|
def enabled(cls) -> bool:
|
|
return True
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
# Built here, not lazily: CustomOp resolves dispatch against the
|
|
# current vLLM config, which is only set during model construction.
|
|
self._post = MHCPostOp()
|
|
|
|
def _maybe_post(
|
|
self,
|
|
residual: torch.Tensor,
|
|
sublayer_out: torch.Tensor | None,
|
|
post_layer_mix: torch.Tensor | None,
|
|
comb_res_mix: torch.Tensor | None,
|
|
) -> torch.Tensor:
|
|
if sublayer_out is None:
|
|
return residual
|
|
assert post_layer_mix is not None and comb_res_mix is not None
|
|
return self._post(sublayer_out, residual, post_layer_mix, comb_res_mix)
|
|
|
|
def forward_cuda(
|
|
self,
|
|
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,
|
|
pre_mix: torch.Tensor | None = None,
|
|
x: torch.Tensor | None = None,
|
|
norm_weight: torch.Tensor | None = None,
|
|
norm_eps: float = 1e-6,
|
|
sublayer_out: torch.Tensor | None = None,
|
|
post_layer_mix: torch.Tensor | None = None,
|
|
comb_res_mix: torch.Tensor | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
residual = self._maybe_post(
|
|
residual, sublayer_out, post_layer_mix, comb_res_mix
|
|
)
|
|
return residual, *torch.ops.vllm.mhc_pre_delayed_tilelang(
|
|
residual,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
pre_mix,
|
|
x,
|
|
norm_weight,
|
|
norm_eps,
|
|
)
|
|
|
|
def forward_hip(
|
|
self,
|
|
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,
|
|
pre_mix: torch.Tensor | None = None,
|
|
x: torch.Tensor | None = None,
|
|
norm_weight: torch.Tensor | None = None,
|
|
norm_eps: float = 1e-6,
|
|
sublayer_out: torch.Tensor | None = None,
|
|
post_layer_mix: torch.Tensor | None = None,
|
|
comb_res_mix: torch.Tensor | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
# The AITER delayed path drives mhc_pre_gemm_sqrsum against `residual`
|
|
# and folds no RMSNorm, so it cannot serve the model-entry broadcast
|
|
# (which projects a narrower `x`) or a fused norm. Both are handled by
|
|
# TileLang, or by the reference when TileLang is unavailable.
|
|
if (
|
|
x is None
|
|
and norm_weight is None
|
|
and _aiter_mhc_supported(residual, None, supports_norm=False)
|
|
):
|
|
from vllm._aiter_ops import rocm_aiter_ops
|
|
|
|
num_tokens = residual.numel() // (residual.shape[-1] * residual.shape[-2])
|
|
# Folding the post in only pays while the residual still fits in
|
|
# cache; past that AITER's heuristic wants the separate kernels,
|
|
# which can ask for a non-temporal store and a large-m split-k.
|
|
if sublayer_out is not None and (
|
|
rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(num_tokens)
|
|
):
|
|
residual = self._maybe_post(
|
|
residual, sublayer_out, post_layer_mix, comb_res_mix
|
|
)
|
|
sublayer_out = None
|
|
# The op writes the folded post's residual into this buffer rather
|
|
# than returning it, since on the unfused path it has none of its
|
|
# own and an op must not hand back one of its inputs.
|
|
residual_out = (
|
|
torch.empty_like(residual) if sublayer_out is not None else None
|
|
)
|
|
rest = torch.ops.vllm.mhc_pre_delayed_aiter(
|
|
residual,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
pre_mix,
|
|
sublayer_out,
|
|
post_layer_mix,
|
|
comb_res_mix,
|
|
residual_out,
|
|
)
|
|
return (residual if residual_out is None else residual_out), *rest
|
|
residual = self._maybe_post(
|
|
residual, sublayer_out, post_layer_mix, comb_res_mix
|
|
)
|
|
if HAS_TILELANG_MHC:
|
|
return residual, *torch.ops.vllm.mhc_pre_delayed_tilelang(
|
|
residual,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
pre_mix,
|
|
x,
|
|
norm_weight,
|
|
norm_eps,
|
|
)
|
|
# The post is already applied, so it must not be requested again.
|
|
return self.forward_native(
|
|
residual,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
pre_mix,
|
|
x,
|
|
norm_weight,
|
|
norm_eps,
|
|
)
|
|
|
|
def forward_native(
|
|
self,
|
|
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,
|
|
pre_mix: torch.Tensor | None = None,
|
|
x: torch.Tensor | None = None,
|
|
norm_weight: torch.Tensor | None = None,
|
|
norm_eps: float = 1e-6,
|
|
sublayer_out: torch.Tensor | None = None,
|
|
post_layer_mix: torch.Tensor | None = None,
|
|
comb_res_mix: torch.Tensor | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
residual = self._maybe_post(
|
|
residual, sublayer_out, post_layer_mix, comb_res_mix
|
|
)
|
|
post_mix, comb_mix, layer_input, next_pre_mix = (
|
|
mhc_kernels.mhc_pre_delayed_torch(
|
|
residual,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
pre_mix=pre_mix,
|
|
x=x,
|
|
)
|
|
)
|
|
return (
|
|
residual,
|
|
post_mix,
|
|
comb_mix,
|
|
_apply_mhc_norm(layer_input, norm_weight, norm_eps),
|
|
next_pre_mix,
|
|
)
|
|
|
|
|
|
# --8<-- [start:mhc_post]
|
|
@CustomOp.register("mhc_post")
|
|
class MHCPostOp(CustomOp):
|
|
"""MHC post block.
|
|
|
|
Combines the layer output with the HC residual streams:
|
|
out_j = post_layer_mix_j * x + sum_i comb_res_mix_ij * residual_i.
|
|
"""
|
|
|
|
# --8<-- [end:mhc_post]
|
|
|
|
@classmethod
|
|
def enabled(cls) -> bool:
|
|
return True
|
|
|
|
def forward_cuda(
|
|
self,
|
|
x: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
post_layer_mix: torch.Tensor,
|
|
comb_res_mix: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
return torch.ops.vllm.mhc_post_tilelang(
|
|
x, residual, post_layer_mix, comb_res_mix
|
|
)
|
|
|
|
def forward_hip(
|
|
self,
|
|
x: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
post_layer_mix: torch.Tensor,
|
|
comb_res_mix: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
if _aiter_mhc_supported(residual, None, supports_norm=True):
|
|
return torch.ops.vllm.mhc_post_aiter(
|
|
x,
|
|
residual,
|
|
post_layer_mix,
|
|
comb_res_mix,
|
|
)
|
|
if HAS_TILELANG_MHC:
|
|
return torch.ops.vllm.mhc_post_tilelang(
|
|
x, residual, post_layer_mix, comb_res_mix
|
|
)
|
|
else:
|
|
return self.forward_native(x, residual, post_layer_mix, comb_res_mix)
|
|
|
|
def forward_native(
|
|
self,
|
|
x: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
post_layer_mix: torch.Tensor,
|
|
comb_res_mix: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
return mhc_kernels.mhc_post_torch(
|
|
x,
|
|
residual,
|
|
post_layer_mix,
|
|
comb_res_mix,
|
|
)
|
|
|
|
def forward_xpu(
|
|
self,
|
|
x: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
post_layer_mix: torch.Tensor,
|
|
comb_res_mix: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
return torch.ops._xpu_C.mhc_post(
|
|
x,
|
|
residual,
|
|
post_layer_mix,
|
|
comb_res_mix,
|
|
)
|
|
|
|
def forward_cpu(
|
|
self,
|
|
x: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
post_layer_mix: torch.Tensor,
|
|
comb_res_mix: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
return mhc_kernels.mhc_post_cpu(x, residual, post_layer_mix, comb_res_mix)
|
|
|
|
|
|
# --8<-- [start:hc_head]
|
|
@CustomOp.register("hc_head")
|
|
class HCHeadOp(CustomOp):
|
|
"""HC head reduction for DeepSeek V4.
|
|
|
|
Computes gates from the RMS-normalized flattened HC residual and
|
|
returns out = sum_i gate_i * residual_i, collapsing hc_mult streams
|
|
to one.
|
|
"""
|
|
|
|
# --8<-- [end:hc_head]
|
|
@classmethod
|
|
def enabled(cls) -> bool:
|
|
return True
|
|
|
|
def forward_cuda(
|
|
self,
|
|
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:
|
|
hc_mult, hidden_size = hidden_states.shape[-2:]
|
|
outer_shape = hidden_states.shape[:-2]
|
|
hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
|
|
out = torch.ops.vllm.hc_head_fused_kernel_tilelang(
|
|
hs_flat,
|
|
hc_fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_norm_eps,
|
|
hc_eps,
|
|
)
|
|
return out.view(*outer_shape, hidden_size)
|
|
|
|
def forward_hip(
|
|
self,
|
|
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:
|
|
hc_mult, hidden_size = hidden_states.shape[-2:]
|
|
outer_shape = hidden_states.shape[:-2]
|
|
hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
|
|
|
|
if HAS_TILELANG_MHC:
|
|
out = torch.ops.vllm.hc_head_fused_kernel_tilelang(
|
|
hs_flat,
|
|
hc_fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_norm_eps,
|
|
hc_eps,
|
|
)
|
|
else:
|
|
num_tokens = hs_flat.shape[0]
|
|
out = torch.empty(
|
|
num_tokens,
|
|
hidden_size,
|
|
dtype=torch.bfloat16,
|
|
device=hidden_states.device,
|
|
)
|
|
torch.ops.vllm.hc_head_triton(
|
|
hs_flat,
|
|
hc_fn,
|
|
hc_scale,
|
|
hc_base,
|
|
out,
|
|
hidden_size,
|
|
rms_norm_eps,
|
|
hc_eps,
|
|
hc_mult,
|
|
)
|
|
|
|
return out.view(*outer_shape, hidden_size)
|
|
|
|
def forward_native(self, *args, **kwargs):
|
|
raise NotImplementedError("Native implementation of hc_head is not available")
|
|
|
|
def forward_xpu(
|
|
self,
|
|
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:
|
|
hc_mult, hidden_size = hidden_states.shape[-2:]
|
|
outer_shape = hidden_states.shape[:-2]
|
|
hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
|
|
num_tokens = hs_flat.shape[0]
|
|
|
|
out = torch.empty(
|
|
num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
|
|
)
|
|
torch.ops._xpu_C.hc_head_fused(
|
|
hs_flat, hc_fn, hc_scale, hc_base, out, rms_norm_eps, hc_eps
|
|
)
|
|
return out.view(*outer_shape, hidden_size)
|
|
|
|
def forward_cpu(
|
|
self,
|
|
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:
|
|
return mhc_kernels.hc_head_fused_cpu(
|
|
hidden_states, hc_fn, hc_scale, hc_base, rms_norm_eps, hc_eps
|
|
)
|
|
|
|
|
|
# --8<-- [start:mhc_fused_post_pre]
|
|
@CustomOp.register("mhc_fused_post_pre")
|
|
class MHCFusedPostPreOp(CustomOp):
|
|
"""Fused MHC post block followed by the next MHC pre block.
|
|
|
|
Equivalent to applying MHCPostOp and then MHCPreOp to the updated
|
|
residual streams, returning residual_cur, post_mix_cur, comb_mix_cur,
|
|
and layer_input_cur.
|
|
"""
|
|
|
|
# --8<-- [end:mhc_fused_post_pre]
|
|
@classmethod
|
|
def enabled(cls) -> bool:
|
|
return True
|
|
|
|
def forward_cuda(
|
|
self,
|
|
x: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
post_layer_mix: torch.Tensor,
|
|
comb_res_mix: 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,
|
|
tile_n: int = 1,
|
|
norm_weight: torch.Tensor | None = None,
|
|
norm_eps: float = 0.0,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
return torch.ops.vllm.mhc_fused_post_pre_tilelang(
|
|
x,
|
|
residual,
|
|
post_layer_mix,
|
|
comb_res_mix,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
n_splits,
|
|
tile_n,
|
|
norm_weight,
|
|
norm_eps,
|
|
)
|
|
|
|
def forward_hip(
|
|
self,
|
|
x: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
post_layer_mix: torch.Tensor,
|
|
comb_res_mix: 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,
|
|
tile_n: int = 1,
|
|
norm_weight: torch.Tensor | None = None,
|
|
norm_eps: float = 0.0,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
if HAS_AITER_MHC_FUSED and _aiter_mhc_supported(
|
|
residual,
|
|
norm_weight,
|
|
supports_norm=HAS_AITER_MHC_FUSED_NORM,
|
|
):
|
|
return torch.ops.vllm.mhc_fused_post_pre_aiter(
|
|
x,
|
|
residual,
|
|
post_layer_mix,
|
|
comb_res_mix,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
n_splits,
|
|
tile_n,
|
|
norm_weight,
|
|
norm_eps,
|
|
)
|
|
if HAS_TILELANG_MHC:
|
|
return torch.ops.vllm.mhc_fused_post_pre_tilelang(
|
|
x,
|
|
residual,
|
|
post_layer_mix,
|
|
comb_res_mix,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
n_splits,
|
|
tile_n,
|
|
norm_weight,
|
|
norm_eps,
|
|
)
|
|
residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur = self.forward_native(
|
|
x,
|
|
residual,
|
|
post_layer_mix,
|
|
comb_res_mix,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
n_splits,
|
|
tile_n,
|
|
norm_weight,
|
|
norm_eps,
|
|
)
|
|
return (
|
|
residual_cur,
|
|
post_mix_cur,
|
|
comb_mix_cur,
|
|
_apply_mhc_norm(layer_input_cur, norm_weight, norm_eps),
|
|
)
|
|
|
|
def forward_native(
|
|
self,
|
|
x: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
post_layer_mix: torch.Tensor,
|
|
comb_res_mix: 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,
|
|
tile_n: int = 1,
|
|
norm_weight: torch.Tensor | None = None,
|
|
norm_eps: float = 0.0,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
# Decompose into post + pre (no fused kernel available).
|
|
residual_cur = mhc_kernels.mhc_post_torch(
|
|
x, residual, post_layer_mix, comb_res_mix
|
|
)
|
|
post_mix_cur, comb_mix_cur, layer_input_cur = mhc_kernels.mhc_pre_torch(
|
|
residual_cur,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
)
|
|
return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur
|
|
|
|
def forward_xpu(
|
|
self,
|
|
x: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
post_layer_mix: torch.Tensor,
|
|
comb_res_mix: 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,
|
|
tile_n: int = 1,
|
|
norm_weight: torch.Tensor | None = None,
|
|
norm_eps: float = 0.0,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
return torch.ops._xpu_C.mhc_fused_post_pre(
|
|
x,
|
|
residual,
|
|
post_layer_mix,
|
|
comb_res_mix,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
)
|
|
|
|
def forward_cpu(
|
|
self,
|
|
x: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
post_layer_mix: torch.Tensor,
|
|
comb_res_mix: 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,
|
|
tile_n: int = 1,
|
|
norm_weight: torch.Tensor | None = None,
|
|
norm_eps: float = 0.0,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
# Decompose into post + pre (no fused kernel available).
|
|
residual_cur = mhc_kernels.mhc_post_cpu(
|
|
x, residual, post_layer_mix, comb_res_mix
|
|
)
|
|
post_mix_cur, comb_mix_cur, layer_input_cur = mhc_kernels.mhc_pre_cpu(
|
|
residual_cur,
|
|
fn,
|
|
hc_scale,
|
|
hc_base,
|
|
rms_eps,
|
|
hc_pre_eps,
|
|
hc_sinkhorn_eps,
|
|
hc_post_mult_value,
|
|
sinkhorn_repeat,
|
|
)
|
|
return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur
|
|
|
|
|
|
def hc_expand(x: torch.Tensor, n: int) -> torch.Tensor:
|
|
"""[s, hidden_size] -> [s, n * hidden_size] by replication."""
|
|
return x.unsqueeze(1).expand(-1, n, -1).contiguous()
|
|
|
|
|
|
def hc_contract(x: torch.Tensor, n: int) -> torch.Tensor:
|
|
"""[s, n * hidden_size] -> [s, hidden_size] by averaging."""
|
|
return x.mean(dim=1)
|