diff --git a/Dockerfile b/Dockerfile index 13e30f3e..904a5021 100644 --- a/Dockerfile +++ b/Dockerfile @@ -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 diff --git a/vllm/patches/kernels/mhc_torch_ops.py b/vllm/patches/kernels/mhc_torch_ops.py deleted file mode 100644 index eee5aba4..00000000 --- a/vllm/patches/kernels/mhc_torch_ops.py +++ /dev/null @@ -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, -) diff --git a/vllm/patches/layers/mhc.py b/vllm/patches/layers/mhc.py index bd484405..e46e3fb5 100644 --- a/vllm/patches/layers/mhc.py +++ b/vllm/patches/layers/mhc.py @@ -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"], +)