@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e7acb20076

Fork 0

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, 00:30 UTC

[Feature] Batch invariant torch.compile (#27660)

Signed-off-by: PaulZhang12 <paulzhan@fb.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

Paul Zhang committed Oct 30, 2025 at 16:11 UTC e7acb200766a0f8f006f9b6fd961dfdceabd7269

4 files changed +82 -9

vllm/config/model.py

-7

@@ -20,9 +20,6 @@ from vllm.config.pooler import PoolerConfig

20 from vllm.config.scheduler import RunnerType

21 from vllm.config.utils import assert_hashable, config, getattr_iter

22 from vllm.logger import init_logger

23 -from vllm.model_executor.layers.batch_invariant import (

24 - vllm_is_batch_invariant,

25 -)

23 from vllm.platforms import current_platform

24 from vllm.transformers_utils.config import (

25 ConfigFormat,

@@ -436,10 +433,6 @@ class ModelConfig:

433 skip_mm_profiling: bool | None,

434 video_pruning_rate: float | None,

435 ) -> None:

439 - # Enable batch invariance settings if requested

440 - if vllm_is_batch_invariant():

441 - self.enforce_eager = True

442 -

436 # Set the default seed to 0 in V1.

437 # NOTE(woosuk): In V1, we use separate processes for workers (unless

438 # VLLM_ENABLE_V1_MULTIPROCESSING=0), so setting a seed here

vllm/envs.py

+7 -1

@@ -251,6 +251,9 @@ def disable_compile_cache() -> bool:

251

252

253 def use_aot_compile() -> bool:

254 + from vllm.model_executor.layers.batch_invariant import (

255 + vllm_is_batch_invariant,

256 + )

257 from vllm.utils.torch_utils import is_torch_equal_or_newer

258

259 default_value = (

@@ -259,7 +262,10 @@ def use_aot_compile() -> bool:

262 else "0"

263 )

264

262 - return os.environ.get("VLLM_USE_AOT_COMPILE", default_value) == "1"

265 + return (

266 + not vllm_is_batch_invariant()

267 + and os.environ.get("VLLM_USE_AOT_COMPILE", default_value) == "1"

268 + )

269

270

271 def env_with_choices(

vllm/model_executor/layers/batch_invariant.py

+71

@@ -11,6 +11,7 @@ import torch

11 import vllm.envs as envs

12 from vllm.logger import init_logger

13 from vllm.triton_utils import tl, triton

14 +from vllm.utils.torch_utils import is_torch_equal_or_newer

16 logger = init_logger(__name__)

@@ -716,6 +717,10 @@ def linear_batch_invariant(input, weight, bias=None):

717 _batch_invariant_MODE = False

718 _batch_invariant_LIB = None

719 _original_torch_bmm = None

720 +_original_fp16_reduction_precision = None

721 +_original_bf16_reduction_precision = None

722 +_original_cublas_workspace_cfg = None

723 +_original_cublaslt_workspace_size = None

724

725

726 def is_batch_invariant_mode_enabled():

@@ -724,6 +729,8 @@ def is_batch_invariant_mode_enabled():

729

730 def enable_batch_invariant_mode():

731 global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm

732 + global _original_fp16_reduction_precision, _original_bf16_reduction_precision

733 + global _original_cublas_workspace_cfg, _original_cublaslt_workspace_size

734 if _batch_invariant_MODE:

735 return

736

@@ -745,14 +752,75 @@ def enable_batch_invariant_mode():

752 _original_torch_bmm = torch.bmm

753 torch.bmm = bmm_batch_invariant

754

755 + _original_bf16_reduction_precision = (

756 + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction

757 + )

758 + _original_fp16_reduction_precision = (

759 + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction

760 + )

761 +

762 + reduced_precision_val = (

763 + (False, False) if is_torch_equal_or_newer("2.10.0.dev") else False

764 + )

765 + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = (

766 + reduced_precision_val

767 + )

768 + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = (

769 + reduced_precision_val

770 + )

771 + torch.backends.cuda.preferred_blas_library(backend="cublaslt")

772 +

773 + if not is_torch_equal_or_newer("2.10.0.dev"):

774 + _original_cublas_workspace_cfg = os.environ.get("CUBLAS_WORKSPACE_CONFIG", None)

775 + _original_cublaslt_workspace_size = os.environ.get(

776 + "CUBLASLT_WORKSPACE_SIZE", None

777 + )

778 + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"

779 + os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1"

780 +

781

782 def disable_batch_invariant_mode():

783 global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm

784 + global _original_fp16_reduction_precision, _original_bf16_reduction_precision

785 + global _original_cublas_workspace_cfg, _original_cublaslt_workspace_size

786 + if not _batch_invariant_MODE:

787 + return

788 +

789 if _batch_invariant_LIB is not None:

790 _batch_invariant_LIB._destroy()

791 if _original_torch_bmm is not None:

792 torch.bmm = _original_torch_bmm

793 _original_torch_bmm = None

794 +

795 + if _original_bf16_reduction_precision is not None:

796 + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = (

797 + _original_bf16_reduction_precision

798 + )

799 + _original_bf16_reduction_precision = None

800 + if _original_fp16_reduction_precision is not None:

801 + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = (

802 + _original_fp16_reduction_precision

803 + )

804 + _original_fp16_reduction_precision = None

805 +

806 + torch.backends.cuda.preferred_blas_library(backend="default")

807 +

808 + if not is_torch_equal_or_newer("2.10.0.dev"):

809 + # Set cublas env vars to previous results. If previous results are None,

810 + # that means the env vars were not set, so we should remove them.

811 + if _original_cublas_workspace_cfg:

812 + os.environ["CUBLAS_WORKSPACE_CONFIG"] = _original_cublas_workspace_cfg

813 + elif "CUBLAS_WORKSPACE_CONFIG" in os.environ:

814 + del os.environ["CUBLAS_WORKSPACE_CONFIG"]

815 +

816 + if _original_cublaslt_workspace_size:

817 + os.environ["CUBLASLT_WORKSPACE_SIZE"] = _original_cublaslt_workspace_size

818 + elif "CUBLASLT_WORKSPACE_SIZE" in os.environ:

819 + del os.environ["CUBLASLT_WORKSPACE_SIZE"]

820 +

821 + _original_cublas_workspace_cfg = None

822 + _original_cublaslt_workspace_size = None

823 +

824 _batch_invariant_MODE = False

825 _batch_invariant_LIB = None

826

@@ -831,6 +899,9 @@ def override_envs_for_invariance():

899 os.environ["NCCL_NTHREADS"] = "1"

900 os.environ["NCCL_SOCKET_NTHREADS"] = "1"

901

902 + # torch.compile settings

903 + os.environ["VLLM_USE_AOT_COMPILE"] = "0"

904 +

905

906 def init_batch_invariance():

907 # this will hit all the csrc overrides as well

vllm/model_executor/layers/quantization/fp8.py

+4 -1

@@ -363,6 +363,7 @@ class Fp8LinearMethod(LinearMethodBase):

363 self.use_marlin = False

364

365 self.use_aiter_and_is_supported = check_aiter_fp8_linear_support()

366 + self.use_deep_gemm = is_deep_gemm_supported()

367

368 self.weight_block_size = self.quant_config.weight_block_size

369 self.block_quant = self.weight_block_size is not None

@@ -545,8 +546,10 @@ class Fp8LinearMethod(LinearMethodBase):

546 # if batch invariant mode is enabled, prefer DeepGEMM FP8 path

547 # we will use BF16 dequant when DeepGEMM is not supported.

548 if vllm_is_batch_invariant():

549 + # Call is_deep_gemm_supported() ahead of time for torch.compile

550 + # dynamo has trouble tracing through

551 if self.block_quant and should_use_deepgemm_for_fp8_linear(

549 - torch.bfloat16, layer.weight, None

552 + torch.bfloat16, layer.weight, self.use_deep_gemm

553 ):

554 # use group quant consistent with block size across K

555 assert self.act_q_group_shape is not None