diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index b6a504d0cf..94ca4f7fe8 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -3299,6 +3299,8 @@ class rocm_aiter_ops: hc_sinkhorn_eps: float, hc_post_mult_value: float, sinkhorn_repeat: int, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Forward pass for mHC pre block. @@ -3313,6 +3315,8 @@ class rocm_aiter_ops: hc_sinkhorn_eps: sinkhorn epsilon hc_post_mult_value: post-mix multiplier value sinkhorn_repeat: number of sinkhorn iterations + norm_weight: optional RMSNorm weight fused into the pre kernel + norm_eps: epsilon for the fused RMSNorm when norm_weight is set Returns: post_mix: shape (..., hc_mult), dtype torch.float32 @@ -3380,6 +3384,8 @@ class rocm_aiter_ops: hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, + norm_weight, + norm_eps, ) return ( post_mix.view(*outer_shape, hc_mult, 1), @@ -3465,5 +3471,95 @@ class rocm_aiter_ops: ) return out.view_as(residual) + @staticmethod + 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, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused mHC post + next mHC pre via AITER. + + Returns residual_cur, post_mix, comb_mix, layer_input in vLLM order + (AITER returns post_mix, comb_mix, layer_input, next_residual). + """ + from aiter.ops.mhc import mhc_fused_post_pre + + assert x.dtype == torch.bfloat16 + assert residual.dtype == torch.bfloat16 + assert fn.dtype == torch.float32 + assert hc_scale.dtype == torch.float32 + assert hc_base.dtype == torch.float32 + + hc_mult = residual.shape[-2] + hidden_size = residual.shape[-1] + outer_shape = residual.shape[:-2] + + residual_flat = residual.view(-1, hc_mult, hidden_size) + num_tokens = residual_flat.shape[0] + x_flat = x.view(num_tokens, hidden_size) + post_flat = post_layer_mix.view(num_tokens, hc_mult, 1) + comb_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult) + + if num_tokens == 0: + return ( + torch.empty_like(residual_flat).view_as(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, + ), + ) + + with torch.device(residual_flat.device): + post_mix, comb_mix, layer_input, next_residual = mhc_fused_post_pre( + x_flat, + residual_flat, + post_flat, + comb_flat, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + norm_weight, + norm_eps, + ) + + return ( + next_residual.view_as(residual), + post_mix.view(*outer_shape, hc_mult, 1), + comb_mix.view(*outer_shape, hc_mult, hc_mult), + layer_input.view(*outer_shape, hidden_size), + ) + rocm_aiter_ops.register_ops_once() diff --git a/vllm/model_executor/kernels/mhc/aiter.py b/vllm/model_executor/kernels/mhc/aiter.py index e844b4191e..ce7690ec1d 100644 --- a/vllm/model_executor/kernels/mhc/aiter.py +++ b/vllm/model_executor/kernels/mhc/aiter.py @@ -16,6 +16,8 @@ def mhc_pre_aiter( 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]: """ Forward pass for mHC pre block. @@ -31,6 +33,8 @@ def mhc_pre_aiter( hc_post_mult_value: post-mix multiplier value sinkhorn_repeat: number of sinkhorn iterations n_splits: split-k factor; + norm_weight: optional RMSNorm weight fused into the pre kernel + norm_eps: epsilon for the fused RMSNorm when norm_weight is set Returns: post_mix: shape (..., hc_mult), dtype torch.float32 @@ -52,6 +56,8 @@ def mhc_pre_aiter( hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, + norm_weight, + norm_eps, ) @@ -66,6 +72,8 @@ def _mhc_pre_aiter_fake( 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]: hc_mult = residual.shape[-2] hidden_size = residual.shape[-1] @@ -124,6 +132,94 @@ def _mhc_post_aiter_fake( return torch.empty_like(residual) +def mhc_fused_post_pre_aiter( + 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]: + """Fused mHC post + next mHC pre on ROCm via AITER. + + Returns residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur. + """ + hidden_size = residual.shape[-1] + assert hidden_size % 256 == 0 + from vllm._aiter_ops import rocm_aiter_ops + + return rocm_aiter_ops.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, + norm_weight, + norm_eps, + ) + + +def _mhc_fused_post_pre_aiter_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, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + hc_mult = residual.shape[-2] + hidden_size = residual.shape[-1] + outer_shape = residual.shape[:-2] + post_mix = torch.empty( + *outer_shape, + hc_mult, + 1, + dtype=torch.float32, + device=residual.device, + ) + comb_mix = torch.empty( + *outer_shape, + hc_mult, + hc_mult, + dtype=torch.float32, + device=residual.device, + ) + layer_input = torch.empty( + *outer_shape, + hidden_size, + dtype=torch.bfloat16, + device=residual.device, + ) + return torch.empty_like(residual), post_mix, comb_mix, layer_input + + direct_register_custom_op( op_name="mhc_pre_aiter", op_func=mhc_pre_aiter, @@ -136,3 +232,9 @@ direct_register_custom_op( mutates_args=[], fake_impl=_mhc_post_aiter_fake, ) +direct_register_custom_op( + op_name="mhc_fused_post_pre_aiter", + op_func=mhc_fused_post_pre_aiter, + mutates_args=[], + fake_impl=_mhc_fused_post_pre_aiter_fake, +) diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py index da6bbf2e00..8f1c22f16a 100644 --- a/vllm/model_executor/layers/mhc.py +++ b/vllm/model_executor/layers/mhc.py @@ -90,7 +90,8 @@ class MHCPreOp(CustomOp): norm_eps: float = 0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: hidden_size = residual.shape[-1] - if HAS_AITER_MHC and hidden_size % 256 == 0: + hc_mult = residual.shape[-2] + if HAS_AITER_MHC and hidden_size % 256 == 0 and hc_mult == 4: return torch.ops.vllm.mhc_pre_aiter( residual, fn, @@ -101,6 +102,9 @@ class MHCPreOp(CustomOp): 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( @@ -222,7 +226,8 @@ class MHCPostOp(CustomOp): comb_res_mix: torch.Tensor, ) -> torch.Tensor: hidden_size = residual.shape[-1] - if HAS_AITER_MHC and hidden_size % 256 == 0: + hc_mult = residual.shape[-2] + if HAS_AITER_MHC and hidden_size % 256 == 0 and hc_mult == 4: return torch.ops.vllm.mhc_post_aiter( x, residual, @@ -444,6 +449,27 @@ class MHCFusedPostPreOp(CustomOp): norm_weight: torch.Tensor | None = None, norm_eps: float = 0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + hidden_size = residual.shape[-1] + hc_mult = residual.shape[-2] + if HAS_AITER_MHC and hidden_size % 256 == 0 and hc_mult == 4: + 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, diff --git a/vllm/models/deepseek_v4/amd/dspark.py b/vllm/models/deepseek_v4/amd/dspark.py index 28056cb550..edca59032a 100644 --- a/vllm/models/deepseek_v4/amd/dspark.py +++ b/vllm/models/deepseek_v4/amd/dspark.py @@ -9,8 +9,8 @@ ROCm port of ``nvidia/dspark.py``. Follows the same nvidia->amd recipe used for attention + MHC CustomOp path) instead of the nvidia one; * route the MHC head through the ``HCHeadOp`` CustomOp dispatcher (aiter / tilelang / triton / torch) instead of calling the tilelang kernels directly, - and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (False on the aiter - path, where the decoder layer already applies hc_post in-layer); + and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (True when AITER + or TileLang fused MHC is available; False only on the torch fallback); * drop the mega-MoE weight path (``make_deepseek_v4_expert_params_mapping`` / ``use_mega_moe`` / ``finalize_mega_moe_weights`` do not exist in amd/model.py). diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 1456408605..f28df04d4b 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -305,6 +305,10 @@ class DeepseekV4MoE(nn.Module): return final_hidden_states.view(org_shape) +# Hidden sizes supported by AITER mhc_pre_big_fuse_rmsnorm. +_AITER_MHC_FUSED_RMSNORM_SIZES = frozenset({1280, 2560, 4096, 7168}) + + class DeepseekV4DecoderLayer(nn.Module): def __init__( self, @@ -384,9 +388,18 @@ class DeepseekV4DecoderLayer(nn.Module): self.mhc_pre = MHCPreOp() self.mhc_post = MHCPostOp() self.mhc_fused_post_pre = MHCFusedPostPreOp() - self.use_fused_mhc = HAS_TILELANG_MHC and not ( - HAS_AITER_MHC and self.hidden_size % 256 == 0 + # AITER mhc kernels (pre/post/fused) require hc_mult == 4. + use_aiter_mhc = ( + HAS_AITER_MHC and self.hidden_size % 256 == 0 and self.hc_mult == 4 ) + # Prefer AITER fused post+pre when eligible; otherwise TileLang. + self.use_fused_mhc = use_aiter_mhc or HAS_TILELANG_MHC + # Fold attn/ffn RMSNorm into MHC only when the active backend's + # fused-rmsnorm path supports this hidden size. + if use_aiter_mhc: + self.fuse_mhc_rmsnorm = self.hidden_size in _AITER_MHC_FUSED_RMSNORM_SIZES + else: + self.fuse_mhc_rmsnorm = HAS_TILELANG_MHC and self.use_fused_mhc def hc_pre( self, @@ -394,7 +407,13 @@ class DeepseekV4DecoderLayer(nn.Module): hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, ): + """Reduce HC residual streams to the next sub-layer input. + + When ``norm_weight`` is set, RMSNorm is fused into the pre kernel. + """ post_mix, res_mix, layer_input = self.mhc_pre( residual=x, fn=hc_fn, @@ -405,6 +424,8 @@ class DeepseekV4DecoderLayer(nn.Module): hc_sinkhorn_eps=self.hc_eps, hc_post_mult_value=self.hc_post_alpha, sinkhorn_repeat=self.hc_sinkhorn_iters, + norm_weight=norm_weight, + norm_eps=norm_eps, ) return layer_input, post_mix, res_mix @@ -426,11 +447,20 @@ class DeepseekV4DecoderLayer(nn.Module): res_mix: torch.Tensor | None = None, residual: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + attn_norm_weight = self.attn_norm.weight if self.fuse_mhc_rmsnorm else None + attn_norm_eps = ( + self.attn_norm.variance_epsilon if self.fuse_mhc_rmsnorm else 0.0 + ) if residual is None: # Run standalone hc_pre on first layer residual = x x, post_mix, res_mix = self.hc_pre( - x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base + x, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + norm_weight=attn_norm_weight, + norm_eps=attn_norm_eps, ) else: residual, post_mix, res_mix, x = self.mhc_fused_post_pre( @@ -446,11 +476,16 @@ class DeepseekV4DecoderLayer(nn.Module): self.hc_eps, self.hc_post_alpha, self.hc_sinkhorn_iters, + norm_weight=attn_norm_weight, + norm_eps=attn_norm_eps, ) - x = self.attn_norm(x) + if not self.fuse_mhc_rmsnorm: + x = self.attn_norm(x) x = self.attn(positions, x, None) + ffn_norm_weight = self.ffn_norm.weight if self.fuse_mhc_rmsnorm else None + ffn_norm_eps = self.ffn_norm.variance_epsilon if self.fuse_mhc_rmsnorm else 0.0 residual, post_mix, res_mix, x = self.mhc_fused_post_pre( x, residual, @@ -464,8 +499,11 @@ class DeepseekV4DecoderLayer(nn.Module): self.hc_eps, self.hc_post_alpha, self.hc_sinkhorn_iters, + norm_weight=ffn_norm_weight, + norm_eps=ffn_norm_eps, ) - x = self.ffn_norm(x) + if not self.fuse_mhc_rmsnorm: + x = self.ffn_norm(x) x = self.ffn(x, input_ids) return x, residual, post_mix, res_mix @@ -678,7 +716,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): residual, ) if (idx + 1) in self.aux_hidden_state_layers: - # On the unfused (aiter) path the layer already applied hc_post, + # On the unfused path the layer already applied hc_post, # so hidden_states is the reconstructed stream; on the fused # path reconstruct it via hc_post before averaging. if layer.use_fused_mhc: