biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

MHC: replace monolithic layers/mhc.py with pure PyTorch

Browse Source 
The nightly vLLM image puts ALL MHC code in layers/mhc.py (not kernels/mhc/).
It imports tilelang at top level and JIT-compiles kernels.

Replace the entire file with pure PyTorch implementations using
direct_register_custom_op for mhc_pre, mhc_post, mhc_fused_post_pre,
and hc_head_fused_kernel. No tilelang dependency at all.

Also removes the separate mhc_torch_ops.py and kernels/mhc/ patches
which don't apply to the nightly image layout.
This commit is contained in:
biondizzle
2026-05-19 05:41:55 +00:00
parent dfd9c10ae9
commit 918342feeb
3 changed files with 180 additions and 320 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
Dockerfile
View File

@@ -40,14 +40,9 @@ COPY vllm/patches/deepseek_v4_attention.py ${VLLM_LAYERS_DIR}/deepseek_v4_attent
COPY vllm/patches/layers/deepseek_compressor.py ${VLLM_LAYERS_DIR}/deepseek_compressor.py
# Replace MHC TileLang kernels with pure PyTorch (avoids TileLang JIT on Blackwell)
# 1. Patch layers/mhc.py — CustomOp dispatch uses torch impls instead of tilelang
# 2. Install our torch op registrations (mhc_torch_ops.py)
# 3. Patch kernels/mhc/__init__.py to not import tilelang/aiter
ARG VLLM_MHC_KERNELS_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/kernels/mhc
# The nightly image has all MHC in layers/mhc.py (imports tilelang at top level).
# Our replacement is pure PyTorch — no tilelang dependency at all.
COPY vllm/patches/layers/mhc.py ${VLLM_LAYERS_DIR}/mhc.py
COPY vllm/patches/kernels/mhc_torch_ops.py ${VLLM_MHC_KERNELS_DIR}/mhc_torch_ops.py
# Rewrite __init__.py: import torch impls + our custom ops, skip tilelang/aiter
RUN printf 'from .torch import mhc_pre_torch, mhc_post_torch\nfrom .mhc_torch_ops import *\n' > ${VLLM_MHC_KERNELS_DIR}/__init__.py
# CuTeDSL NVFP4 linear kernel (registered as NvFp4LinearKernel)
ARG VLLM_NVFP4_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/kernels/linear/nvfp4

230
vllm/patches/kernels/mhc_torch_ops.py
View File

@@ -1,230 +0,0 @@
# 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,
)

261
vllm/patches/layers/mhc.py
View File

@@ -1,100 +1,195 @@
# SPDX-License-Identifier: Apache-2.0
# Patched MHC layer — replaces TileLang kernels with pure PyTorch.
# This avoids TileLang JIT compilation on Blackwell (SM100).
# Patched MHC layer — pure PyTorch, no TileLang.
# Replaces the TileLang-based mhc.py from the vLLM nightly image.
# The original imports tilelang at the top level and JIT-compiles kernels
# which don't work correctly on Blackwell (SM100).
import torch
from vllm.model_executor.custom_op import CustomOp
# Import our torch implementations (registers torch.ops.vllm.mhc_pre, etc.)
import vllm.model_executor.kernels.mhc.mhc_torch_ops as _mhc_torch # noqa: F401
from vllm.utils.torch_utils import direct_register_custom_op
@CustomOp.register("mhc_pre")
class MHCPreOp(CustomOp):
@classmethod
def enabled(cls) -> bool:
return True
# ── Pure PyTorch MHC implementations ──────────────────────────────────
def forward_cuda(self, residual, fn, hc_scale, hc_base,
rms_eps, hc_pre_eps, hc_sinkhorn_eps,
hc_post_mult_value, sinkhorn_repeat, n_splits=1):
return _mhc_torch.mhc_pre(
residual, fn, hc_scale, hc_base,
rms_eps, hc_pre_eps, hc_sinkhorn_eps,
hc_post_mult_value, sinkhorn_repeat, n_splits,
)
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]:
assert residual.dtype == torch.bfloat16
assert fn.dtype == torch.float32
assert hc_scale.dtype == torch.float32
assert hc_base.dtype == torch.float32
def forward_hip(self, *args, **kwargs):
return self.forward_cuda(*args, **kwargs)
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]
def forward_native(self, *args, **kwargs):
return self.forward_cuda(*args, **kwargs)
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),
)
@CustomOp.register("mhc_post")
class MHCPostOp(CustomOp):
@classmethod
def enabled(cls) -> bool:
return True
def forward_cuda(self, x, residual, post_layer_mix, comb_res_mix):
return _mhc_torch.mhc_post(x, residual, post_layer_mix, comb_res_mix)
def forward_hip(self, *args, **kwargs):
return self.forward_cuda(*args, **kwargs)
def forward_native(self, *args, **kwargs):
return self.forward_cuda(*args, **kwargs)
def _mhc_pre_fake(
residual, fn, hc_scale, hc_base, rms_eps, hc_pre_eps,
hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, n_splits=1,
):
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),
)
@CustomOp.register("hc_head")
class HCHeadOp(CustomOp):
@classmethod
def enabled(cls) -> bool:
return True
def forward_cuda(self, hidden_states, hc_fn, hc_scale, hc_base,
rms_norm_eps, hc_eps):
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.empty(
hs_flat.shape[0], hidden_size,
dtype=torch.bfloat16, device=hidden_states.device,
)
_mhc_torch.hc_head_fused_kernel(
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, *args, **kwargs):
return self.forward_cuda(*args, **kwargs)
def forward_native(self, *args, **kwargs):
return self.forward_cuda(*args, **kwargs)
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)
@CustomOp.register("mhc_fused_post_pre")
class MHCFusedPostPreOp(CustomOp):
@classmethod
def enabled(cls) -> bool:
return True
def _mhc_post_fake(x, residual, post_layer_mix, comb_res_mix):
return torch.empty_like(residual)
def forward_cuda(self, 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=1, tile_n=1):
return _mhc_torch.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, n_splits, tile_n,
)
def forward_hip(self, *args, **kwargs):
return self.forward_cuda(*args, **kwargs)
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 forward_native(self, *args, **kwargs):
return self.forward_cuda(*args, **kwargs)
def _mhc_fused_post_pre_fake(
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=1, tile_n=1,
):
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:
if hs_flat.shape[0] == 0:
return
x_flat = hs_flat.reshape(hs_flat.shape[0], hc_mult * hidden_size).to(torch.float32)
mixes = torch.matmul(x_flat, fn.t())
sqrsum = x_flat.square().sum(dim=-1, keepdim=True)
rsqrt = torch.rsqrt(sqrsum / (hc_mult * hidden_size) + rms_eps)
pre_mix = torch.sigmoid(mixes * rsqrt * hc_scale[0] + hc_base) + hc_eps
result = torch.sum(pre_mix.unsqueeze(-1) * hs_flat.to(torch.float32), dim=1).to(out.dtype)
out.copy_(result)
# ── Register as torch custom ops ──────────────────────────────────────
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"],
)