292 lines
8.2 KiB
Python
292 lines
8.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import torch
|
|
|
|
# this import will also register the custom ops
|
|
import vllm.model_executor.kernels.mhc as mhc_kernels
|
|
from vllm.model_executor.custom_op import CustomOp
|
|
|
|
|
|
# --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,
|
|
) -> 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,
|
|
)
|
|
|
|
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,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
# TODO: Reenable aiter after we are at the aiter
|
|
# version that has this bugfix
|
|
# https://github.com/ROCm/aiter/commit/b639cb63bcac4672dce33a731fad042a65cb3649
|
|
# It has accuracy problem at large number of tokens.
|
|
# hidden_size = residual.shape[-1]
|
|
# if hidden_size % 256 == 0:
|
|
# 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,
|
|
# )
|
|
# else:
|
|
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_native(self, *args, **kwargs):
|
|
raise NotImplementedError("Native implementation of mhc_pre is not available")
|
|
|
|
|
|
# --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:
|
|
# TODO: Reenable aiter after we are at the aiter
|
|
# version that has this bugfix
|
|
# https://github.com/ROCm/aiter/commit/b639cb63bcac4672dce33a731fad042a65cb3649
|
|
# It has accuracy problem at large number of tokens.
|
|
# hidden_size = residual.shape[-1]
|
|
# if hidden_size % 256 == 0:
|
|
# return torch.ops.vllm.mhc_post_aiter(
|
|
# x,
|
|
# residual,
|
|
# post_layer_mix,
|
|
# comb_res_mix,
|
|
# )
|
|
# else:
|
|
return mhc_kernels.mhc_post_torch(
|
|
x,
|
|
residual,
|
|
post_layer_mix,
|
|
comb_res_mix,
|
|
)
|
|
|
|
def forward_native(self, *args, **kwargs):
|
|
raise NotImplementedError("Native implementation of mhc_post is not available")
|
|
|
|
|
|
# --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)
|
|
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_fused_kernel_tilelang(
|
|
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_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)
|
|
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")
|
|
|
|
|
|
# --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,
|
|
) -> 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,
|
|
)
|
|
|
|
def forward_hip(self, *args, **kwargs):
|
|
raise NotImplementedError(
|
|
"Hip implementation of mhc_fused_post_pre is not available"
|
|
)
|
|
|
|
def forward_native(self, *args, **kwargs):
|
|
raise NotImplementedError(
|
|
"Native implementation of mhc_fused_post_pre is not available"
|
|
)
|