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 11, 12:30 UTC
set assume_32bit_indexing and pass unbacked hints (#30459)
Signed-off-by: Laith Sakka <lsakka@meta.com>
Laith Sakka committed
Dec 13, 2025 at 18:36 UTC
763963aa7358e19d627f1bf614a00f415a4ef6b3
1 file changed
+22
-3
vllm/compilation/decorators.py
+22
-3
@@ -28,7 +28,7 @@ from vllm.config.compilation import DynamicShapesType
28
from vllm.logger import init_logger
29
from vllm.sequence import IntermediateTensors
30
from vllm.utils.import_utils import resolve_obj_by_qualname
31
-from vllm.utils.torch_utils import supports_dynamo
31
+from vllm.utils.torch_utils import is_torch_equal_or_newer, supports_dynamo
32
33
from .monitor import start_monitoring_torch_compile
34
@@ -316,7 +316,13 @@ def _support_torch_compile(
316
def _mark_dynamic_inputs(mod, type, *args, **kwargs):
317
def mark_dynamic(arg, dims):
318
if type == DynamicShapesType.UNBACKED:
319
- torch._dynamo.decorators.mark_unbacked(arg, dims)
319
+ if is_torch_equal_or_newer("2.10.0.dev"):
320
+ for dim in dims:
321
+ torch._dynamo.decorators.mark_unbacked(
322
+ arg, dim, hint_override=arg.size()[dim]
323
+ )
324
+ else:
325
+ torch._dynamo.decorators.mark_unbacked(arg, dims)
326
else:
327
torch._dynamo.mark_dynamic(arg, dims)
328
@@ -350,7 +356,13 @@ def _support_torch_compile(
356
if isinstance(arg, torch.Tensor):
357
# In case dims is specified with negative indexing
358
dims = [arg.ndim + dim if dim < 0 else dim for dim in dims]
353
- torch._dynamo.decorators.mark_unbacked(arg, dims)
359
+ if is_torch_equal_or_newer("2.10.0.dev"):
360
+ for dim in dims:
361
+ torch._dynamo.decorators.mark_unbacked(
362
+ arg, dim, hint_override=arg.size()[dim]
363
+ )
364
+ else:
365
+ torch._dynamo.decorators.mark_unbacked(arg, dims)
366
367
def __call__(self, *args, **kwargs):
368
# torch.compiler.is_compiling() means we are inside the compilation
@@ -488,6 +500,12 @@ def _support_torch_compile(
500
if ds_type == DynamicShapesType.BACKED_SIZE_OBLIVIOUS:
501
fx_config_patches["backed_size_oblivious"] = True
502
503
+ # Prepare inductor config patches
504
+ # assume_32bit_indexing is only available in torch 2.10.0.dev+
505
+ inductor_config_patches = {}
506
+ if is_torch_equal_or_newer("2.10.0.dev"):
507
+ inductor_config_patches["assume_32bit_indexing"] = True
508
+
509
with (
510
patch.object(
511
InliningInstructionTranslator, "inline_call_", patched_inline_call
@@ -496,6 +514,7 @@ def _support_torch_compile(
514
maybe_use_cudagraph_partition_wrapper(self.vllm_config),
515
torch.fx.experimental._config.patch(**fx_config_patches),
516
_torch27_patch_tensor_subclasses(),
517
+ torch._inductor.config.patch(**inductor_config_patches),
518
):
519
if envs.VLLM_USE_AOT_COMPILE:
520
self.aot_compiled_fn = self.aot_compile(*args, **kwargs)