nvfp4-megamoe-kernel/vllm/patches/kernels/mhc_torch_ops.py

# SPDX-License-Identifier: Apache-2.0
# Pure PyTorch MHC kernels for DeepSeek V4.
# Replaces TileLang kernels to avoid TileLang JIT compilation on Blackwell (SM100).
# Registers torch.ops.vllm.mhc_pre, mhc_post, mhc_fused_post_pre, hc_head_fused_kernel.

import torch
from vllm.utils.torch_utils import direct_register_custom_op


# ── Pure PyTorch implementations ──────────────────────────────────────

def mhc_pre(
    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]:
    hc_mult = residual.shape[-2]
    hidden_size = residual.shape[-1]
    hc_mult2 = hc_mult * hc_mult
    hc_mult3 = hc_mult * 2 + hc_mult2
    hc_hidden_size = hc_mult * hidden_size
    outer_shape = residual.shape[:-2]

    residual_flat = residual.view(-1, hc_mult, hidden_size)
    num_tokens = residual_flat.shape[0]

    x = residual_flat.view(num_tokens, hc_hidden_size).to(torch.float32)
    mixes = torch.matmul(x, fn.t())
    sqrsum = x.square().sum(dim=-1, keepdim=True)
    mixes = mixes * torch.rsqrt(sqrsum / hc_hidden_size + rms_eps)

    pre_logits = mixes[:, :hc_mult] * hc_scale[0] + hc_base[:hc_mult]
    pre_mix = torch.sigmoid(pre_logits) + hc_pre_eps

    post_logits = mixes[:, hc_mult:2 * hc_mult] * hc_scale[1] + hc_base[hc_mult:2 * hc_mult]
    post_mix = torch.sigmoid(post_logits) * hc_post_mult_value

    comb_logits = (mixes[:, 2 * hc_mult:]
                   .view(num_tokens, hc_mult, hc_mult)
                   * hc_scale[2]
                   + hc_base[2 * hc_mult:].view(1, hc_mult, hc_mult))
    comb_mix = torch.softmax(comb_logits, dim=-1) + hc_sinkhorn_eps
    comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)
    for _ in range(sinkhorn_repeat - 1):
        comb_mix = comb_mix / (comb_mix.sum(dim=-1, keepdim=True) + hc_sinkhorn_eps)
        comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)

    layer_input = torch.sum(
        pre_mix.unsqueeze(-1) * residual_flat.to(torch.float32), dim=1
    ).to(torch.bfloat16)

    return (
        post_mix.view(*outer_shape, hc_mult, 1),
        comb_mix.view(*outer_shape, hc_mult, hc_mult),
        layer_input.view(*outer_shape, hidden_size),
    )


def _mhc_pre_fake(
    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]:
    hc_mult = residual.shape[-2]
    hidden_size = residual.shape[-1]
    outer_shape = residual.shape[:-2]
    return (
        torch.empty(*outer_shape, hc_mult, 1, dtype=torch.float32, device=residual.device),
        torch.empty(*outer_shape, hc_mult, hc_mult, dtype=torch.float32, device=residual.device),
        torch.empty(*outer_shape, hidden_size, dtype=torch.bfloat16, device=residual.device),
    )


def mhc_post(
    x: torch.Tensor,
    residual: torch.Tensor,
    post_layer_mix: torch.Tensor,
    comb_res_mix: torch.Tensor,
) -> torch.Tensor:
    mixed_residual = torch.einsum(
        "...ij,...ih->...jh",
        comb_res_mix.to(torch.float32),
        residual.to(torch.float32),
    )
    post_term = post_layer_mix.to(torch.float32) * x.unsqueeze(-2).to(torch.float32)
    return (mixed_residual + post_term).to(residual.dtype)


def _mhc_post_fake(
    x: torch.Tensor,
    residual: torch.Tensor,
    post_layer_mix: torch.Tensor,
    comb_res_mix: torch.Tensor,
) -> torch.Tensor:
    return torch.empty_like(residual)


def mhc_fused_post_pre(
    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]:
    new_residual = mhc_post(x, residual, post_layer_mix, comb_res_mix)
    post_mix, res_mix, layer_input = mhc_pre(
        new_residual, fn, hc_scale, hc_base,
        rms_eps, hc_pre_eps, hc_sinkhorn_eps,
        hc_post_mult_value, sinkhorn_repeat, n_splits,
    )
    return new_residual, post_mix, res_mix, layer_input


def _mhc_fused_post_pre_fake(
    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]:
    hc_mult = residual.shape[-2]
    hidden_size = residual.shape[-1]
    outer_shape = residual.shape[:-2]
    return (
        torch.empty_like(residual),
        torch.empty(*outer_shape, hc_mult, 1, dtype=torch.float32, device=residual.device),
        torch.empty(*outer_shape, hc_mult, hc_mult, dtype=torch.float32, device=residual.device),
        torch.empty(*outer_shape, hidden_size, dtype=torch.bfloat16, device=residual.device),
    )


def hc_head_fused_kernel(
    hs_flat: torch.Tensor,
    fn: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    out: torch.Tensor,
    hidden_size: int,
    rms_eps: float,
    hc_eps: float,
    hc_mult: int,
) -> None:
    x_flat = hs_flat.flatten(-2)
    sqrsum = x_flat.to(torch.float32).square().sum(dim=-1, keepdim=True)
    x_normed = x_flat * torch.rsqrt(sqrsum / x_flat.shape[-1] + rms_eps)
    mixes = torch.nn.functional.linear(x_normed.to(torch.float32), fn)
    pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps
    result = torch.sum(
        pre.unsqueeze(-1) * hs_flat.to(torch.float32), dim=1
    ).to(torch.bfloat16)
    out.copy_(result)


def _hc_head_fused_kernel_fake(
    hs_flat: torch.Tensor,
    fn: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    out: torch.Tensor,
    hidden_size: int,
    rms_eps: float,
    hc_eps: float,
    hc_mult: int,
) -> None:
    return None


# ── Register as torch custom ops ──────────────────────────────────────
# These replace the TileLang-registered ops that the model code calls via
# torch.ops.vllm.mhc_pre, torch.ops.vllm.mhc_post, etc.

direct_register_custom_op(
    op_name="mhc_pre",
    op_func=mhc_pre,
    mutates_args=[],
    fake_impl=_mhc_pre_fake,
)

direct_register_custom_op(
    op_name="mhc_post",
    op_func=mhc_post,
    mutates_args=[],
    fake_impl=_mhc_post_fake,
)

direct_register_custom_op(
    op_name="mhc_fused_post_pre",
    op_func=mhc_fused_post_pre,
    mutates_args=[],
    fake_impl=_mhc_fused_post_pre_fake,
)

direct_register_custom_op(
    op_name="hc_head_fused_kernel",
    op_func=hc_head_fused_kernel,
    mutates_args=["out"],
    fake_impl=_hc_head_fused_kernel_fake,
)