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 7, 22:45 UTC
[Model Bash][DSR1] Add selective dynamic shape marking for CustomOp (#34900)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Vadim Gimpelson committed
Feb 22, 2026 at 04:28 UTC
74d90b1ce49e5984ccf054d6e918c8efbafce3c1
2 files changed
+46
-7
vllm/model_executor/custom_op.py
+42
-6
@@ -1,5 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import functools
4
+import inspect
5
+
6
import torch
7
import torch.nn as nn
8
@@ -205,9 +208,9 @@ class CustomOp(nn.Module):
208
NOTE: this does not enable fusion across ops, so opaque custom ops
209
should still be unwrapped wherever possible.
210
"""
208
- # Do not compile if compilation disabled
211
from vllm.config.compilation import CompilationMode
212
213
+ # Do not compile if compilation disabled
214
if not enable:
215
return fn
216
@@ -220,14 +223,42 @@ class CustomOp(nn.Module):
223
if compilation_config.backend == "eager":
224
return fn
225
226
+ compile_options = maybe_disable_graph_partition(
227
+ current_platform.simple_compile_backend
228
+ )
229
+ backend = current_platform.simple_compile_backend
230
+
231
+ dynamic_arg_dims = getattr(self.__class__, "_dynamic_arg_dims", None)
232
+ if dynamic_arg_dims is not None:
233
+ compiled_fn = torch.compile(
234
+ fn,
235
+ dynamic=False,
236
+ backend=backend,
237
+ options=compile_options,
238
+ )
239
+ sig = inspect.signature(fn)
240
+
241
+ @functools.wraps(fn)
242
+ def wrapper(*args, **kwargs):
243
+ bound = sig.bind(*args, **kwargs)
244
+ bound.apply_defaults()
245
+ for name, dims in dynamic_arg_dims.items():
246
+ arg = bound.arguments.get(name)
247
+ if arg is not None and isinstance(arg, torch.Tensor):
248
+ dims_list = [dims] if isinstance(dims, int) else dims
249
+ for d in dims_list:
250
+ real_d = arg.ndim + d if d < 0 else d
251
+ torch._dynamo.mark_dynamic(arg, real_d)
252
+ return compiled_fn(*args, **kwargs)
253
+
254
+ return wrapper
255
+
256
# dynamic=True to avoid recompilations
257
return torch.compile(
258
fn,
259
dynamic=True,
227
- backend=current_platform.simple_compile_backend,
228
- options=maybe_disable_graph_partition(
229
- current_platform.simple_compile_backend
230
- ),
260
+ backend=backend,
261
+ options=compile_options,
262
)
263
264
@classmethod
@@ -267,10 +298,15 @@ class CustomOp(nn.Module):
298
299
# Decorator to register custom ops.
300
@classmethod
270
- def register(cls, name: str):
301
+ def register(
302
+ cls,
303
+ name: str,
304
+ dynamic_arg_dims: dict[str, int | list[int]] | None = None,
305
+ ):
306
def decorator(op_cls):
307
assert name not in op_registry, f"Duplicate op name: {name}"
308
op_cls.name = name
309
+ op_cls._dynamic_arg_dims = dynamic_arg_dims
310
op_registry[name] = op_cls
311
return op_cls
312
vllm/model_executor/layers/attention/mla_attention.py
+4
-1
@@ -950,7 +950,10 @@ def dynamic_per_batched_tensor_quant(
950
logger = init_logger(__name__)
951
952
953
-@CustomOp.register("mla_decode_concat_quant_fp8")
953
+@CustomOp.register(
954
+ "mla_decode_concat_quant_fp8",
955
+ dynamic_arg_dims={"decode_ql_nope": 0, "decode_q_pe": 0},
956
+)
957
class _DecodeConcatQuantFP8(QuantFP8):
958
"""
959
QuantFP8 variant that concatenates decode_ql_nope and decode_q_pe before