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 14, 05:30 UTC
Remove redundant mutates_args and dispatch_key for direct_register_custom_op (#25512)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Sep 23, 2025 at 18:48 UTC
7361ab379f817727173691224e41e0175423b795
28 files changed
+9
-66
vllm/attention/layer.py
-3
@@ -575,9 +575,7 @@ def unified_attention_fake(
575
direct_register_custom_op(
576
op_name="unified_attention",
577
op_func=unified_attention,
578
- mutates_args=[],
578
fake_impl=unified_attention_fake,
580
- dispatch_key=current_platform.dispatch_key,
579
tags=tag_cudagraph_unsafe,
580
)
581
@@ -628,6 +626,5 @@ direct_register_custom_op(
626
op_func=unified_attention_with_output,
627
mutates_args=["output", "output_block_scale"],
628
fake_impl=unified_attention_with_output_fake,
631
- dispatch_key=current_platform.dispatch_key,
629
tags=tag_cudagraph_unsafe,
630
)
vllm/compilation/collective_fusion.py
-1
@@ -547,7 +547,6 @@ if flashinfer_comm is not None:
547
"scale_out",
548
],
549
fake_impl=call_trtllm_fused_allreduce_norm_fake,
550
- dispatch_key=current_platform.dispatch_key,
550
)
551
flashinfer_trtllm_fused_allreduce_norm = (
552
torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default)
vllm/distributed/device_communicators/pynccl.py
-1
@@ -46,7 +46,6 @@ def register_nccl_symmetric_ops(pynccl_comm):
46
direct_register_custom_op(
47
op_name="all_reduce_symmetric_with_copy",
48
op_func=all_reduce_symmetric_with_copy_impl,
49
- mutates_args=[],
49
fake_impl=all_reduce_symmetric_with_copy_fake,
50
)
51
vllm/distributed/parallel_state.py
-7
@@ -149,29 +149,22 @@ def all_gather_fake(tensor: torch.Tensor, dim: int, world_size: int,
149
150
151
if supports_custom_op():
152
- from vllm.platforms import current_platform
152
direct_register_custom_op(
153
op_name="all_reduce",
154
op_func=all_reduce,
156
- mutates_args=[],
155
fake_impl=all_reduce_fake,
158
- dispatch_key=current_platform.dispatch_key,
156
)
157
158
direct_register_custom_op(
159
op_name="reduce_scatter",
160
op_func=reduce_scatter,
164
- mutates_args=[],
161
fake_impl=reduce_scatter_fake,
166
- dispatch_key=current_platform.dispatch_key,
162
)
163
164
direct_register_custom_op(
165
op_name="all_gather",
166
op_func=all_gather,
172
- mutates_args=[],
167
fake_impl=all_gather_fake,
174
- dispatch_key=current_platform.dispatch_key,
168
)
169
170
vllm/lora/ops/triton_ops/lora_expand_op.py
-2
@@ -11,7 +11,6 @@ import torch
11
12
from vllm.lora.ops.triton_ops.kernel_utils import do_expand_kernel
13
from vllm.lora.ops.triton_ops.utils import _get_lora_b_ptr
14
-from vllm.platforms import current_platform
14
from vllm.triton_utils import tl, triton
15
from vllm.utils import direct_register_custom_op
16
@@ -283,7 +282,6 @@ try:
282
op_func=_lora_expand,
283
mutates_args=["output_tensor"],
284
fake_impl=_lora_expand_fake,
286
- dispatch_key=current_platform.dispatch_key,
285
)
286
lora_expand = torch.ops.vllm.lora_expand
287
vllm/lora/ops/triton_ops/lora_shrink_op.py
-2
@@ -11,7 +11,6 @@ import torch
11
12
from vllm.lora.ops.triton_ops.kernel_utils import do_shrink_kernel
13
from vllm.lora.ops.triton_ops.utils import _get_lora_a_ptr
14
-from vllm.platforms import current_platform
14
from vllm.triton_utils import tl, triton
15
from vllm.utils import direct_register_custom_op
16
@@ -237,7 +236,6 @@ try:
236
op_func=_lora_shrink,
237
mutates_args=["output_tensor"],
238
fake_impl=_lora_shrink_fake,
240
- dispatch_key=current_platform.dispatch_key,
239
)
240
lora_shrink = torch.ops.vllm.lora_shrink
241
vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py
-1
@@ -92,7 +92,6 @@ def flashinfer_fused_moe_blockscale_fp8_fake(
92
direct_register_custom_op(
93
op_name="flashinfer_fused_moe_blockscale_fp8",
94
op_func=flashinfer_fused_moe_blockscale_fp8,
95
- mutates_args=[],
95
fake_impl=flashinfer_fused_moe_blockscale_fp8_fake,
96
tags=(torch.Tag.needs_fixed_stride_order, ),
97
)
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
-1
@@ -235,6 +235,5 @@ def fused_marlin_moe_fake(hidden_states: torch.Tensor,
235
direct_register_custom_op(
236
op_name="fused_marlin_moe",
237
op_func=fused_marlin_moe,
238
- mutates_args=[],
238
fake_impl=fused_marlin_moe_fake,
239
)
vllm/model_executor/layers/fused_moe/fused_moe.py
-1
@@ -1256,7 +1256,6 @@ def outplace_fused_experts_fake(
1256
direct_register_custom_op(
1257
op_name="outplace_fused_experts",
1258
op_func=outplace_fused_experts,
1259
- mutates_args=[],
1259
fake_impl=outplace_fused_experts_fake,
1260
tags=(() if is_torch_equal_or_newer("2.7.0") else
1261
(torch.Tag.needs_fixed_stride_order, )),
vllm/model_executor/layers/fused_moe/layer.py
-2
@@ -2040,7 +2040,6 @@ direct_register_custom_op(
2040
op_func=moe_forward,
2041
mutates_args=["hidden_states"],
2042
fake_impl=moe_forward_fake,
2043
- dispatch_key=current_platform.dispatch_key,
2043
tags=(torch.Tag.needs_fixed_stride_order, ),
2044
)
2045
@@ -2071,7 +2070,6 @@ direct_register_custom_op(
2070
op_func=moe_forward_shared,
2071
mutates_args=["hidden_states"],
2072
fake_impl=moe_forward_shared_fake,
2074
- dispatch_key=current_platform.dispatch_key,
2073
tags=(torch.Tag.needs_fixed_stride_order, ),
2074
)
2075
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
-7
@@ -223,17 +223,13 @@ if current_platform.is_rocm():
223
direct_register_custom_op(
224
op_name="rocm_aiter_asm_moe_tkw1",
225
op_func=rocm_aiter_asm_moe_tkw1_impl,
226
- mutates_args=[],
226
fake_impl=rocm_aiter_asm_moe_tkw1_fake,
228
- dispatch_key=current_platform.dispatch_key,
227
)
228
229
direct_register_custom_op(
230
op_name="rocm_aiter_fused_moe",
231
op_func=rocm_aiter_fused_moe_impl,
234
- mutates_args=[],
232
fake_impl=rocm_aiter_fused_moe_fake,
236
- dispatch_key=current_platform.dispatch_key,
233
)
234
235
direct_register_custom_op(
@@ -241,7 +237,6 @@ if current_platform.is_rocm():
237
op_func=rocm_aiter_topk_softmax_impl,
238
mutates_args=["topk_weights", "topk_indices", "token_expert_indices"],
239
fake_impl=rocm_aiter_topk_softmax_fake,
244
- dispatch_key=current_platform.dispatch_key,
240
)
241
242
direct_register_custom_op(
@@ -249,7 +244,6 @@ if current_platform.is_rocm():
244
op_func=rocm_aiter_biased_grouped_topk_impl,
245
mutates_args=["topk_weights", "topk_ids"],
246
fake_impl=rocm_aiter_biased_grouped_topk_fake,
252
- dispatch_key=current_platform.dispatch_key,
247
)
248
249
direct_register_custom_op(
@@ -257,7 +251,6 @@ if current_platform.is_rocm():
251
op_func=rocm_aiter_grouped_topk_impl,
252
mutates_args=["topk_weights", "topk_ids"],
253
fake_impl=rocm_aiter_grouped_topk_fake,
260
- dispatch_key=current_platform.dispatch_key,
254
)
255
256
vllm/model_executor/layers/layernorm.py
-4
@@ -103,17 +103,13 @@ if current_platform.is_rocm():
103
direct_register_custom_op(
104
op_name="rocm_aiter_rms_norm",
105
op_func=rocm_aiter_rms_norm_impl,
106
- mutates_args=[],
106
fake_impl=rocm_aiter_rms_norm_fake,
108
- dispatch_key=current_platform.dispatch_key,
107
)
108
109
direct_register_custom_op(
110
op_name="rocm_aiter_rmsnorm2d_fwd_with_add",
111
op_func=rocm_aiter_rmsnorm2d_fwd_with_add_impl,
114
- mutates_args=[],
112
fake_impl=rocm_aiter_rmsnorm2d_fwd_with_add_fake,
116
- dispatch_key=current_platform.dispatch_key,
113
)
114
115
vllm/model_executor/layers/mamba/linear_attn.py
-2
@@ -31,7 +31,6 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
31
MambaStateDtypeCalculator, MambaStateShapeCalculator)
32
from vllm.model_executor.layers.quantization.base_config import (
33
QuantizationConfig)
34
-from vllm.platforms import current_platform
34
from vllm.utils import direct_register_custom_op
35
from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
36
@@ -401,5 +400,4 @@ direct_register_custom_op(
400
op_func=linear_attention,
401
mutates_args=["output"],
402
fake_impl=linear_attention_fake,
404
- dispatch_key=current_platform.dispatch_key,
403
)
vllm/model_executor/layers/mamba/mamba_mixer.py
-2
@@ -27,7 +27,6 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
27
from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
28
selective_scan_fn, selective_state_update)
29
from vllm.model_executor.utils import set_weight_attrs
30
-from vllm.platforms import current_platform
30
from vllm.utils import direct_register_custom_op
31
from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionMetadata
32
@@ -464,5 +463,4 @@ direct_register_custom_op(
463
op_func=mamba_mixer,
464
mutates_args=["output"],
465
fake_impl=mamba_mixer_fake,
467
- dispatch_key=current_platform.dispatch_key,
466
)
vllm/model_executor/layers/mamba/mamba_mixer2.py
-2
@@ -34,7 +34,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
34
from vllm.model_executor.model_loader.weight_utils import (
35
LoaderFunction, composed_weight_loader, sharded_weight_loader)
36
from vllm.model_executor.utils import set_weight_attrs
37
-from vllm.platforms import current_platform
37
from vllm.utils import direct_register_custom_op
38
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
39
@@ -765,5 +764,4 @@ direct_register_custom_op(
764
op_func=mamba_mixer2,
765
mutates_args=["output"],
766
fake_impl=mamba_mixer2_fake,
768
- dispatch_key=current_platform.dispatch_key,
767
)
vllm/model_executor/layers/mamba/short_conv.py
-2
@@ -21,7 +21,6 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
21
MambaStateDtypeCalculator, MambaStateShapeCalculator)
22
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
23
causal_conv1d_fn, causal_conv1d_update)
24
-from vllm.platforms import current_platform
24
from vllm.utils import direct_register_custom_op
25
from vllm.v1.attention.backends.short_conv_attn import (
26
ShortConvAttentionMetadata)
@@ -251,5 +250,4 @@ direct_register_custom_op(
250
op_func=short_conv,
251
mutates_args=["output"],
252
fake_impl=short_conv_fake,
254
- dispatch_key=current_platform.dispatch_key,
253
)
vllm/model_executor/layers/quantization/deepgemm.py
-3
@@ -4,7 +4,6 @@ import logging
4
5
import torch
6
7
-from vllm.platforms import current_platform
7
from vllm.triton_utils import triton
8
from vllm.utils import direct_register_custom_op
9
from vllm.utils.deep_gemm import fp8_gemm_nt
@@ -75,7 +74,5 @@ def w8a8_deepgemm_block_scaled_mm_fake(
74
direct_register_custom_op(
75
op_name="w8a8_deepgemm_block_scaled_mm",
76
op_func=w8a8_deepgemm_block_scaled_mm,
78
- mutates_args=[],
77
fake_impl=w8a8_deepgemm_block_scaled_mm_fake,
80
- dispatch_key=current_platform.dispatch_key,
78
)
vllm/model_executor/layers/quantization/gguf.py
-3
@@ -161,7 +161,6 @@ try:
161
direct_register_custom_op(
162
op_name="_fused_mul_mat_gguf",
163
op_func=_fused_mul_mat_gguf,
164
- mutates_args=[],
164
fake_impl=_fused_mul_mat_gguf_fake,
165
)
166
fused_mul_mat_gguf = torch.ops.vllm._fused_mul_mat_gguf
@@ -273,7 +272,6 @@ try:
272
direct_register_custom_op(
273
op_name="_fused_moe_gguf",
274
op_func=_fused_moe_gguf,
276
- mutates_args=[],
275
fake_impl=_fused_moe_gguf_fake,
276
)
277
fused_moe_gguf = torch.ops.vllm._fused_moe_gguf
@@ -319,7 +317,6 @@ try:
317
direct_register_custom_op(
318
op_name="_apply_gguf_embedding",
319
op_func=_apply_gguf_embedding,
322
- mutates_args=[],
320
fake_impl=_apply_gguf_embedding_fake,
321
)
322
apply_gguf_embedding = torch.ops.vllm._apply_gguf_embedding
vllm/model_executor/layers/quantization/kernels/scaled_mm/aiter.py
-2
@@ -51,9 +51,7 @@ if current_platform.is_rocm():
51
direct_register_custom_op(
52
op_name="rocm_aiter_gemm_w8a8",
53
op_func=rocm_aiter_gemm_w8a8_impl,
54
- mutates_args=[],
54
fake_impl=rocm_aiter_gemm_w8a8_fake,
56
- dispatch_key=current_platform.dispatch_key,
55
)
56
57
vllm/model_executor/layers/quantization/utils/fp8_utils.py
-3
@@ -91,9 +91,7 @@ if current_platform.is_rocm():
91
direct_register_custom_op(
92
op_name="rocm_aiter_gemm_w8a8_blockscale",
93
op_func=rocm_aiter_gemm_w8a8_blockscale_impl,
94
- mutates_args=[],
94
fake_impl=rocm_aiter_gemm_w8a8_blockscale_fake,
96
- dispatch_key=current_platform.dispatch_key,
95
)
96
if (envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_LINEAR
97
and current_platform.is_fp8_fnuz()):
@@ -135,7 +133,6 @@ def _w8a8_triton_block_scaled_mm_fake(
133
direct_register_custom_op(
134
"w8a8_triton_block_scaled_mm_func",
135
_w8a8_triton_block_scaled_mm_func,
138
- mutates_args=[],
136
fake_impl=_w8a8_triton_block_scaled_mm_fake,
137
dispatch_key="CUDA",
138
)
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
-2
@@ -113,7 +113,6 @@ try:
113
direct_register_custom_op(
114
op_name="dequant_mxfp4",
115
op_func=_dequant_mxfp4,
116
- mutates_args=[],
116
fake_impl=_dequant_mxfp4_fake,
117
)
118
dequant_mxfp4 = torch.ops.vllm.dequant_mxfp4
@@ -124,7 +123,6 @@ try:
123
direct_register_custom_op(
124
op_name="quant_dequant_mxfp4",
125
op_func=_quant_dequant_mxfp4,
127
- mutates_args=[],
126
fake_impl=_quant_dequant_mxfp4_fake,
127
)
128
quant_dequant_mxfp4 = torch.ops.vllm.quant_dequant_mxfp4
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
-2
@@ -218,9 +218,7 @@ def rocm_per_tensor_w8a8_scaled_mm(*, qinput: torch.Tensor,
218
direct_register_custom_op(
219
op_name="rocm_per_tensor_w8a8_scaled_mm_impl",
220
op_func=rocm_per_tensor_w8a8_scaled_mm_impl,
221
- mutates_args=[],
221
fake_impl=rocm_per_tensor_w8a8_scaled_mm_fake,
223
- dispatch_key=current_platform.dispatch_key,
222
)
223
224
vllm/model_executor/layers/rotary_embedding/common.py
-1
@@ -147,5 +147,4 @@ direct_register_custom_op(
147
op_func=_flashinfer_rotary_embedding,
148
mutates_args=["query", "key"], # These tensors are modified in-place
149
fake_impl=_flashinfer_rotary_embedding_fake,
150
- dispatch_key=current_platform.dispatch_key,
150
)
vllm/model_executor/layers/utils.py
-2
@@ -136,9 +136,7 @@ def rocm_unquantized_gemm(layer: torch.nn.Module,
136
direct_register_custom_op(
137
op_name="rocm_unquantized_gemm_impl",
138
op_func=rocm_unquantized_gemm_impl,
139
- mutates_args=[],
139
fake_impl=rocm_unquantized_gemm_impl_fake,
141
- dispatch_key=current_platform.dispatch_key,
140
)
141
142
vllm/model_executor/models/deepseek_v2.py
-3
@@ -56,7 +56,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
56
ParallelLMHead, VocabParallelEmbedding)
57
from vllm.model_executor.model_loader.weight_utils import (
58
default_weight_loader, maybe_remap_kv_scale_name)
59
-from vllm.platforms import current_platform
59
from vllm.sequence import IntermediateTensors
60
from vllm.utils import cdiv, direct_register_custom_op
61
@@ -141,9 +140,7 @@ def sequence_parallel_chunk_fake(x: torch.Tensor) -> torch.Tensor:
140
direct_register_custom_op(
141
op_name="sequence_parallel_chunk",
142
op_func=sequence_parallel_chunk,
144
- mutates_args=[],
143
fake_impl=sequence_parallel_chunk_fake,
146
- dispatch_key=current_platform.dispatch_key,
144
tags=(torch.Tag.needs_fixed_stride_order, ),
145
)
146
vllm/model_executor/models/plamo2.py
-2
@@ -48,7 +48,6 @@ from vllm.model_executor.models.utils import (
48
is_pp_missing_parameter, make_empty_intermediate_tensors_factory,
49
make_layers, maybe_prefix)
50
from vllm.model_executor.utils import set_weight_attrs
51
-from vllm.platforms import current_platform
51
from vllm.sequence import IntermediateTensors
52
from vllm.utils import direct_register_custom_op
53
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
@@ -490,7 +489,6 @@ direct_register_custom_op(
489
op_func=plamo2_mamba_mixer,
490
mutates_args=["output"],
491
fake_impl=plamo2_mamba_mixer_fake,
493
- dispatch_key=current_platform.dispatch_key,
492
)
493
494
vllm/model_executor/models/qwen3_next.py
-1
@@ -1225,7 +1225,6 @@ direct_register_custom_op(
1225
op_func=gdn_attention,
1226
mutates_args=["output"],
1227
fake_impl=gdn_attention_fake,
1228
- dispatch_key=current_platform.dispatch_key,
1228
)
1229
1230
vllm/utils/__init__.py
+9
-2
@@ -2546,10 +2546,10 @@ vllm_lib = Library("vllm", "FRAGMENT") # noqa
2546
def direct_register_custom_op(
2547
op_name: str,
2548
op_func: Callable,
2549
- mutates_args: list[str],
2549
+ mutates_args: Optional[list[str]] = None,
2550
fake_impl: Optional[Callable] = None,
2551
target_lib: Optional[Library] = None,
2552
- dispatch_key: str = "CUDA",
2552
+ dispatch_key: Optional[str] = None,
2553
tags: tuple[torch.Tag, ...] = (),
2554
):
2555
"""
@@ -2577,6 +2577,13 @@ def direct_register_custom_op(
2577
"the required dependencies.")
2578
return
2579
2580
+ if mutates_args is None:
2581
+ mutates_args = []
2582
+
2583
+ if dispatch_key is None:
2584
+ from vllm.platforms import current_platform
2585
+ dispatch_key = current_platform.dispatch_key
2586
+
2587
import torch.library
2588
if hasattr(torch.library, "infer_schema"):
2589
schema_str = torch.library.infer_schema(op_func,