diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index d052b41fc5..4fae5dc052 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -25,7 +25,6 @@ from vllm.models.deepseek_v4.common.ops import ( fused_q_kv_rmsnorm, ) from vllm.utils.deep_gemm import fp8_einsum -from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.ops.rocm_aiter_mla_sparse import rocm_inv_rope_einsum if TYPE_CHECKING: @@ -292,8 +291,10 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): device=hidden_states.device, ) - # Attention (inside custom op for torch.compile boundary) - torch.ops.vllm.deepseek_v4_attention( + # @eager_break_during_capture: this is where the breakable + # cudagraph capture breaks (the attention op runs eagerly between + # captured graph segments). + deepseek_v4_attention( hidden_states, positions, o_padded, @@ -334,14 +335,12 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): device=o.device, dtype=torch.bfloat16, ) - torch.ops.vllm.deepseek_v4_fp8_einsum( - o_fp8, - o_scale, - wo_a_fp8, - wo_a_scale, - z, + fp8_einsum( "bhr,hdr->bhd", - list(self._einsum_recipe), + (o_fp8, o_scale), + (wo_a_fp8, wo_a_scale), + z, + recipe=self._einsum_recipe, ) return self.wo_b(z.flatten(1)) @@ -554,55 +553,6 @@ def deepseek_v4_attention( self.attention_impl(hidden_states, positions, out) -def deepseek_v4_attention_fake( - hidden_states: torch.Tensor, - positions: torch.Tensor, - out: torch.Tensor, - layer_name: str, -) -> None: - return None - - -direct_register_custom_op( - op_name="deepseek_v4_attention", - op_func=deepseek_v4_attention, - mutates_args=["out"], - fake_impl=deepseek_v4_attention_fake, -) - - -def deepseek_v4_fp8_einsum( - a: torch.Tensor, - a_scale: torch.Tensor, - b: torch.Tensor, - b_scale: torch.Tensor, - out: torch.Tensor, - equation: str, - recipe: list[int], -) -> None: - fp8_einsum(equation, (a, a_scale), (b, b_scale), out, recipe=tuple(recipe)) - - -def deepseek_v4_fp8_einsum_fake( - a: torch.Tensor, - a_scale: torch.Tensor, - b: torch.Tensor, - b_scale: torch.Tensor, - out: torch.Tensor, - equation: str, - recipe: list[int], -) -> None: - return None - - -direct_register_custom_op( - op_name="deepseek_v4_fp8_einsum", - op_func=deepseek_v4_fp8_einsum, - mutates_args=["out"], - fake_impl=deepseek_v4_fp8_einsum_fake, -) - - class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase): def __init__( self, diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index a5f4bae842..6ade4caf9d 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -15,7 +15,6 @@ from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) -from vllm.forward_context import get_forward_context from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( @@ -60,7 +59,6 @@ from vllm.models.deepseek_v4.attention import ( ) from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs from vllm.sequence import IntermediateTensors -from vllm.utils.torch_utils import direct_register_custom_op class DeepseekV4MLP(nn.Module): @@ -209,13 +207,6 @@ class DeepseekV4MegaMoEExperts(nn.Module): self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None - # Register in the static forward context so the custom-op wrapper - # can look up this module by name from within a torch.compile graph. - compilation_config = vllm_config.compilation_config - if prefix in compilation_config.static_forward_context: - raise ValueError(f"Duplicate layer name: {prefix}") - compilation_config.static_forward_context[prefix] = self - def _map_global_expert_id(self, expert_id: int) -> int: if expert_id < self.experts_start_idx or expert_id >= self.experts_end_idx: return -1 @@ -349,12 +340,11 @@ class DeepseekV4MegaMoEExperts(nn.Module): f"but the symmetric buffer was sized for {self.max_num_tokens}." ) y = torch.empty_like(hidden_states, dtype=torch.bfloat16) - torch.ops.vllm.deepseek_v4_mega_moe_experts( + self._run_mega_moe( hidden_states, topk_weights, topk_ids, y, - self.prefix, activation_clamp, fast_math, ) @@ -402,46 +392,6 @@ class DeepseekV4MegaMoEExperts(nn.Module): DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] -def _deepseek_v4_mega_moe_experts_op( - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - out: torch.Tensor, - layer_name: str, - activation_clamp: float | None, - fast_math: bool, -) -> None: - self = get_forward_context().no_compile_layers[layer_name] - self._run_mega_moe( - hidden_states, - topk_weights, - topk_ids, - out, - activation_clamp, - fast_math, - ) - - -def _deepseek_v4_mega_moe_experts_op_fake( - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - out: torch.Tensor, - layer_name: str, - activation_clamp: float | None, - fast_math: bool, -) -> None: - return None - - -direct_register_custom_op( - op_name="deepseek_v4_mega_moe_experts", - op_func=_deepseek_v4_mega_moe_experts_op, - mutates_args=["out"], - fake_impl=_deepseek_v4_mega_moe_experts_op_fake, -) - - class DeepseekV4MoE(nn.Module): def __init__( self,