A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 12:15 UTC
[Model Refactoring] Remove unncessary torch op registration for DSv4 (#43891)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
May 28, 2026 at 14:04 UTC
69b8956dcd5ac0bd3762a2dd5e070bde26443ea4
2 files changed
+10
-110
vllm/models/deepseek_v4/attention.py
+9
-59
@@ -25,7 +25,6 @@ from vllm.models.deepseek_v4.common.ops import (
25
fused_q_kv_rmsnorm,
26
)
27
from vllm.utils.deep_gemm import fp8_einsum
28
-from vllm.utils.torch_utils import direct_register_custom_op
28
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import rocm_inv_rope_einsum
29
30
if TYPE_CHECKING:
@@ -292,8 +291,10 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
291
device=hidden_states.device,
292
)
293
295
- # Attention (inside custom op for torch.compile boundary)
296
- torch.ops.vllm.deepseek_v4_attention(
294
+ # @eager_break_during_capture: this is where the breakable
295
+ # cudagraph capture breaks (the attention op runs eagerly between
296
+ # captured graph segments).
297
+ deepseek_v4_attention(
298
hidden_states,
299
positions,
300
o_padded,
@@ -334,14 +335,12 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
335
device=o.device,
336
dtype=torch.bfloat16,
337
)
337
- torch.ops.vllm.deepseek_v4_fp8_einsum(
338
- o_fp8,
339
- o_scale,
340
- wo_a_fp8,
341
- wo_a_scale,
342
- z,
338
+ fp8_einsum(
339
"bhr,hdr->bhd",
344
- list(self._einsum_recipe),
340
+ (o_fp8, o_scale),
341
+ (wo_a_fp8, wo_a_scale),
342
+ z,
343
+ recipe=self._einsum_recipe,
344
)
345
346
return self.wo_b(z.flatten(1))
@@ -554,55 +553,6 @@ def deepseek_v4_attention(
553
self.attention_impl(hidden_states, positions, out)
554
555
557
-def deepseek_v4_attention_fake(
558
- hidden_states: torch.Tensor,
559
- positions: torch.Tensor,
560
- out: torch.Tensor,
561
- layer_name: str,
562
-) -> None:
563
- return None
564
-
565
-
566
-direct_register_custom_op(
567
- op_name="deepseek_v4_attention",
568
- op_func=deepseek_v4_attention,
569
- mutates_args=["out"],
570
- fake_impl=deepseek_v4_attention_fake,
571
-)
572
-
573
-
574
-def deepseek_v4_fp8_einsum(
575
- a: torch.Tensor,
576
- a_scale: torch.Tensor,
577
- b: torch.Tensor,
578
- b_scale: torch.Tensor,
579
- out: torch.Tensor,
580
- equation: str,
581
- recipe: list[int],
582
-) -> None:
583
- fp8_einsum(equation, (a, a_scale), (b, b_scale), out, recipe=tuple(recipe))
584
-
585
-
586
-def deepseek_v4_fp8_einsum_fake(
587
- a: torch.Tensor,
588
- a_scale: torch.Tensor,
589
- b: torch.Tensor,
590
- b_scale: torch.Tensor,
591
- out: torch.Tensor,
592
- equation: str,
593
- recipe: list[int],
594
-) -> None:
595
- return None
596
-
597
-
598
-direct_register_custom_op(
599
- op_name="deepseek_v4_fp8_einsum",
600
- op_func=deepseek_v4_fp8_einsum,
601
- mutates_args=["out"],
602
- fake_impl=deepseek_v4_fp8_einsum_fake,
603
-)
604
-
605
-
556
class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
557
def __init__(
558
self,
vllm/models/deepseek_v4/nvidia/model.py
+1
-51
@@ -15,7 +15,6 @@ from vllm.distributed import (
15
get_tensor_model_parallel_rank,
16
get_tensor_model_parallel_world_size,
17
)
18
-from vllm.forward_context import get_forward_context
18
from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
19
from vllm.model_executor.layers.fused_moe import FusedMoE
20
from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
@@ -60,7 +59,6 @@ from vllm.models.deepseek_v4.attention import (
59
)
60
from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
61
from vllm.sequence import IntermediateTensors
63
-from vllm.utils.torch_utils import direct_register_custom_op
62
63
64
class DeepseekV4MLP(nn.Module):
@@ -209,13 +207,6 @@ class DeepseekV4MegaMoEExperts(nn.Module):
207
self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None
208
self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None
209
212
- # Register in the static forward context so the custom-op wrapper
213
- # can look up this module by name from within a torch.compile graph.
214
- compilation_config = vllm_config.compilation_config
215
- if prefix in compilation_config.static_forward_context:
216
- raise ValueError(f"Duplicate layer name: {prefix}")
217
- compilation_config.static_forward_context[prefix] = self
218
-
210
def _map_global_expert_id(self, expert_id: int) -> int:
211
if expert_id < self.experts_start_idx or expert_id >= self.experts_end_idx:
212
return -1
@@ -349,12 +340,11 @@ class DeepseekV4MegaMoEExperts(nn.Module):
340
f"but the symmetric buffer was sized for {self.max_num_tokens}."
341
)
342
y = torch.empty_like(hidden_states, dtype=torch.bfloat16)
352
- torch.ops.vllm.deepseek_v4_mega_moe_experts(
343
+ self._run_mega_moe(
344
hidden_states,
345
topk_weights,
346
topk_ids,
347
y,
357
- self.prefix,
348
activation_clamp,
349
fast_math,
350
)
@@ -402,46 +392,6 @@ class DeepseekV4MegaMoEExperts(nn.Module):
392
DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined]
393
394
405
-def _deepseek_v4_mega_moe_experts_op(
406
- hidden_states: torch.Tensor,
407
- topk_weights: torch.Tensor,
408
- topk_ids: torch.Tensor,
409
- out: torch.Tensor,
410
- layer_name: str,
411
- activation_clamp: float | None,
412
- fast_math: bool,
413
-) -> None:
414
- self = get_forward_context().no_compile_layers[layer_name]
415
- self._run_mega_moe(
416
- hidden_states,
417
- topk_weights,
418
- topk_ids,
419
- out,
420
- activation_clamp,
421
- fast_math,
422
- )
423
-
424
-
425
-def _deepseek_v4_mega_moe_experts_op_fake(
426
- hidden_states: torch.Tensor,
427
- topk_weights: torch.Tensor,
428
- topk_ids: torch.Tensor,
429
- out: torch.Tensor,
430
- layer_name: str,
431
- activation_clamp: float | None,
432
- fast_math: bool,
433
-) -> None:
434
- return None
435
-
436
-
437
-direct_register_custom_op(
438
- op_name="deepseek_v4_mega_moe_experts",
439
- op_func=_deepseek_v4_mega_moe_experts_op,
440
- mutates_args=["out"],
441
- fake_impl=_deepseek_v4_mega_moe_experts_op_fake,
442
-)
443
-
444
-
395
class DeepseekV4MoE(nn.Module):
396
def __init__(
397
self,