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 10, 09:00 UTC
make assume_32_bit_indexing configurable (#32044)
Signed-off-by: Laith Sakka <lsakka@meta.com>
Laith Sakka committed
Jan 10, 2026 at 23:15 UTC
46eb30f5196edf5691376febb7923bfa620649da
2 files changed
+9
-1
vllm/compilation/decorators.py
+3
-1
@@ -522,7 +522,9 @@ def _support_torch_compile(
522
# assume_32bit_indexing is only available in torch 2.10.0.dev+
523
inductor_config_patches = {}
524
if is_torch_equal_or_newer("2.10.0.dev"):
525
- inductor_config_patches["assume_32bit_indexing"] = True
525
+ inductor_config_patches["assume_32bit_indexing"] = (
526
+ self.compilation_config.dynamic_shapes_config.assume_32_bit_indexing
527
+ )
528
529
with (
530
patch.object(
vllm/config/compilation.py
+6
@@ -278,7 +278,11 @@ class DynamicShapesConfig:
278
artifacts also.
279
When type is backed, aot_compile must be disabled for this mode to work.
280
until this change picked up https://github.com/pytorch/pytorch/pull/169239.
281
+ """
282
283
+ assume_32_bit_indexing: bool = True
284
+ """
285
+ whether all tensor sizes can use 32 bit indexing.
286
"""
287
288
def compute_hash(self) -> str:
@@ -640,6 +644,7 @@ class CompilationConfig:
644
"compilation_time",
645
"static_forward_context",
646
"pass_config", # handled separately below
647
+ "dynamic_shapes_config", # handled separately below
648
}
649
650
from vllm.config.utils import get_hash_factors, hash_factors
@@ -647,6 +652,7 @@ class CompilationConfig:
652
factors = get_hash_factors(self, ignored_factors)
653
654
factors["pass_config"] = self.pass_config.compute_hash()
655
+ factors["dynamic_shapes_config"] = self.dynamic_shapes_config.compute_hash()
656
return hash_factors(factors)
657
658
def __repr__(self) -> str: