dc6908ac6a · tduncan/Nvidia-Optimization-Stack-VLLM

tests/kernels/attention/test_use_trtllm_attention.py

+8 -8

@@ -55,37 +55,37 @@ def _clear_supports_cache():

55 # supports_trtllm_attention

56

57

58 -@patch("vllm.utils.flashinfer.vllm_is_batch_invariant", return_value=True)

59 -def test_supports_batch_invariant_disables(_mock):

58 +@patch("vllm.envs.VLLM_BATCH_INVARIANT", True)

59 +def test_supports_batch_invariant_disables():

60 assert supports_trtllm_attention() is False

61

62

63 -@patch("vllm.utils.flashinfer.vllm_is_batch_invariant", return_value=False)

63 +@patch("vllm.envs.VLLM_BATCH_INVARIANT", False)

64 @patch(

65 "vllm.utils.flashinfer.current_platform.is_device_capability_family",

66 return_value=True,

67 )

68 @patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=True)

69 -def test_supports_sm100_with_artifactory(_art, _cap, _bi):

69 +def test_supports_sm100_with_artifactory(_art, _cap):

70 assert supports_trtllm_attention() is True

71

72

73 -@patch("vllm.utils.flashinfer.vllm_is_batch_invariant", return_value=False)

73 +@patch("vllm.envs.VLLM_BATCH_INVARIANT", False)

74 @patch(

75 "vllm.utils.flashinfer.current_platform.is_device_capability_family",

76 return_value=False,

77 )

78 -def test_supports_non_sm100_platform(_cap, _bi):

78 +def test_supports_non_sm100_platform(_cap):

79 assert supports_trtllm_attention() is False

80

81

82 -@patch("vllm.utils.flashinfer.vllm_is_batch_invariant", return_value=False)

82 +@patch("vllm.envs.VLLM_BATCH_INVARIANT", False)

83 @patch(

84 "vllm.utils.flashinfer.current_platform.is_device_capability_family",

85 return_value=True,

86 )

87 @patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=False)

88 -def test_supports_sm100_without_artifactory(_art, _cap, _bi):

88 +def test_supports_sm100_without_artifactory(_art, _cap):

89 assert supports_trtllm_attention() is False

90

91

tests/kernels/moe/test_grouped_topk.py

+2 -2

@@ -8,7 +8,7 @@ Run `pytest tests/kernels/moe/test_grouped_topk.py`.

8 import pytest

9 import torch

10

11 -import vllm.model_executor.layers.batch_invariant as batch_invariant

11 +import vllm.envs as envs

12 from vllm.config import (

13 CompilationConfig,

14 VllmConfig,

@@ -69,7 +69,7 @@ def test_grouped_topk(

69

70 with set_current_vllm_config(vllm_config), monkeypatch.context() as m:

71 m.setenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", "0")

72 - m.setattr(batch_invariant, "VLLM_BATCH_INVARIANT", True)

72 + m.setattr(envs, "VLLM_BATCH_INVARIANT", True)

73 grouped_topk = GroupedTopk(

74 topk=topk,

75 renormalize=renormalize,

tests/v1/determinism/conftest.py

+2 -2

@@ -2,11 +2,11 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import pytest

4

5 -import vllm.model_executor.layers.batch_invariant as batch_invariant

5 +import vllm.envs as envs

6

7

8 @pytest.fixture(autouse=True)

9 def enable_batch_invariant_mode(monkeypatch: pytest.MonkeyPatch):

10 """Automatically enable batch invariant kernel overrides for all tests."""

11 - monkeypatch.setattr(batch_invariant, "VLLM_BATCH_INVARIANT", True)

11 + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True)

12 monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")

tests/v1/determinism/test_batch_invariance.py

+6 -10

@@ -15,7 +15,7 @@ from utils import (

15 skip_unsupported,

16 )

17

18 -import vllm.model_executor.layers.batch_invariant as batch_invariant

18 +import vllm.envs as envs

19 from vllm import LLM, SamplingParams

20

21 IS_DEVICE_CAPABILITY_BELOW_90 = is_device_capability_below_90()

@@ -173,11 +173,9 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(

173

174 # For batch invariance, disable custom all-reduce to ensure deterministic

175 # all-reduce operations (custom all-reduce may not be deterministic)

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

177 - vllm_is_batch_invariant,

178 - )

176 + import vllm.envs as envs

177

180 - disable_custom_ar = vllm_is_batch_invariant()

178 + disable_custom_ar = envs.VLLM_BATCH_INVARIANT

179

180 if disable_custom_ar:

181 print(f"\n{'=' * 80}")

@@ -454,7 +452,7 @@ def test_logprobs_without_batch_invariance_should_fail(

452 """

453 # CRITICAL: Disable batch invariance for this test

454 monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0")

457 - monkeypatch.setattr(batch_invariant, "VLLM_BATCH_INVARIANT", False)

455 + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", False)

456 seed = int(os.getenv("VLLM_TEST_SEED", "12345"))

457 random.seed(seed)

458 tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))

@@ -674,11 +672,9 @@ def test_decode_logprobs_match_prefill_logprobs(

672 random.seed(seed)

673 tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))

674

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

678 - vllm_is_batch_invariant,

679 - )

675 + import vllm.envs as envs

676

681 - disable_custom_ar = vllm_is_batch_invariant()

677 + disable_custom_ar = envs.VLLM_BATCH_INVARIANT

678

679 if disable_custom_ar:

680 print(f"\n{'=' * 80}")

vllm/config/parallel.py

+1 -4

@@ -14,9 +14,6 @@ from typing_extensions import Self

14 import vllm.envs as envs

15 from vllm.config.utils import config

16 from vllm.logger import init_logger

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

18 - vllm_is_batch_invariant,

19 -)

17 from vllm.platforms import current_platform

18 from vllm.utils.network_utils import get_open_ports_list

19 from vllm.utils.torch_utils import cuda_device_count_stateless

@@ -786,7 +783,7 @@ class ParallelConfig:

783 from vllm.v1.executor import Executor

784

785 # Enable batch invariance settings if requested

789 - if vllm_is_batch_invariant():

786 + if envs.VLLM_BATCH_INVARIANT:

787 self.disable_custom_all_reduce = True

788

789 if (

vllm/config/vllm.py

+1 -3

@@ -1112,11 +1112,9 @@ class VllmConfig: # type: ignore[misc]

1112 "when cudagraph_mode piecewise cudagraphs is used, "

1113 f"cudagraph_mode={self.compilation_config.cudagraph_mode}"

1114 )

1115 - from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant

1116 -

1115 if (

1116 self.model_config

1119 - and vllm_is_batch_invariant()

1117 + and envs.VLLM_BATCH_INVARIANT

1118 and not self.model_config.disable_cascade_attn

1119 ):

1120 self.model_config.disable_cascade_attn = True

vllm/distributed/device_communicators/all_reduce_utils.py

+1 -4

@@ -19,9 +19,6 @@ import torch.multiprocessing as mp

19 import vllm.envs as envs

20 from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary

21 from vllm.logger import init_logger

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

23 - vllm_is_batch_invariant,

24 -)

22 from vllm.utils.system_utils import update_environment_variables

23 from vllm.utils.torch_utils import cuda_device_count_stateless

24

@@ -115,7 +112,7 @@ def should_nccl_symm_mem_allreduce(world_size: int, input_tensor: torch.Tensor)

112 is_symmetric_memory_enabled,

113 )

114

118 - if vllm_is_batch_invariant():

115 + if envs.VLLM_BATCH_INVARIANT:

116 return False

117

118 if not is_symmetric_memory_enabled():

vllm/distributed/device_communicators/symm_mem.py

+2 -4

@@ -5,13 +5,11 @@ import torch

5 import torch.distributed as dist

6 from torch.distributed import ProcessGroup

7

8 +import vllm.envs as envs

9 from vllm.distributed.device_communicators.all_reduce_utils import (

10 SYMM_MEM_ALL_REDUCE_MAX_SIZES,

11 )

12 from vllm.logger import init_logger

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

13 - vllm_is_batch_invariant,

14 -)

13 from vllm.platforms import current_platform

14

15 try:

@@ -112,7 +110,7 @@ class SymmMemCommunicator:

110 return

111 self.force_multimem = force_multimem

112 self.disabled = False

115 - if vllm_is_batch_invariant():

113 + if envs.VLLM_BATCH_INVARIANT:

114 self.disabled = True

115

116 def should_use_symm_mem(self, inp: torch.Tensor):

vllm/envs.py

+5 -4

@@ -74,6 +74,7 @@ if TYPE_CHECKING:

74 VLLM_TARGET_DEVICE: str = "cuda"

75 VLLM_MAIN_CUDA_VERSION: str = "12.9"

76 VLLM_FLOAT32_MATMUL_PRECISION: Literal["highest", "high", "medium"] = "highest"

77 + VLLM_BATCH_INVARIANT: bool = False

78 MAX_JOBS: str | None = None

79 NVCC_THREADS: str | None = None

80 VLLM_USE_PRECOMPILED: bool = False

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

281

282

283 def use_aot_compile() -> bool:

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

284 - vllm_is_batch_invariant,

285 - )

284 from vllm.utils.torch_utils import is_torch_equal_or_newer

285

286 default_value = (

@@ -292,7 +290,7 @@ def use_aot_compile() -> bool:

290 )

291

292 return (

295 - not vllm_is_batch_invariant()

293 + not bool(int(os.getenv("VLLM_BATCH_INVARIANT", "0")))

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

295 )

296

@@ -498,6 +496,9 @@ environment_variables: dict[str, Callable[[], Any]] = {

496 ["highest", "high", "medium"],

497 case_sensitive=False,

498 ),

499 + # Enable batch-invariant mode: deterministic results regardless of

500 + # batch composition. Requires NVIDIA GPU with compute capability >= 9.0.

501 + "VLLM_BATCH_INVARIANT": lambda: bool(int(os.getenv("VLLM_BATCH_INVARIANT", "0"))),

502 # Maximum number of compilation jobs to run in parallel.

503 # By default this is the number of CPUs

504 "MAX_JOBS": lambda: os.getenv("MAX_JOBS", None),

vllm/lora/ops/triton_ops/utils.py

+1 -2

@@ -11,12 +11,11 @@ import torch

11

12 from vllm import envs

13 from vllm.logger import init_logger

14 -from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant

14 from vllm.platforms import current_platform

15 from vllm.utils.math_utils import next_power_of_2

16

17 logger = init_logger(__name__)

19 -is_batch_invariant = vllm_is_batch_invariant()

18 +is_batch_invariant = envs.VLLM_BATCH_INVARIANT

19

20 _LORA_A_PTR_DICT: dict[tuple[int, ...], tuple[torch.tensor, ...]] = {}

21 _LORA_B_PTR_DICT: dict[tuple[int, ...], tuple[torch.tensor, ...]] = {}

vllm/model_executor/kernels/linear/scaled_mm/marlin.py

+1 -2

@@ -6,7 +6,6 @@ from collections.abc import Sequence

6 import torch

7

8 import vllm.envs as envs

9 -from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant

9 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

10 process_fp8_weight_block_strategy,

11 )

@@ -42,7 +41,7 @@ class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):

41 # Check if platform supports FP8 Marlin

42 if not is_fp8_marlin_supported():

43 return False, "FP8 Marlin requires compute capability 7.5 or higher"

45 - if vllm_is_batch_invariant():

44 + if envs.VLLM_BATCH_INVARIANT:

45 return False, "FP8 Marlin not supported for batch invariant execution."

46 if (

47 compute_capability is not None

vllm/model_executor/layers/attention/attention.py

+1 -2

@@ -15,7 +15,6 @@ from vllm.model_executor.layers.attention.kv_transfer_utils import (

15 maybe_transfer_kv_layer,

16 )

17 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

18 -from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant

18 from vllm.model_executor.layers.linear import (

19 UnquantizedLinearMethod,

20 )

@@ -296,7 +295,7 @@ class Attention(nn.Module, AttentionLayerBase):

295 if (

296 cache_config is not None

297 and cache_config.enable_prefix_caching

299 - and vllm_is_batch_invariant()

298 + and envs.VLLM_BATCH_INVARIANT

299 and (

300 self.attn_backend.get_name() == "FLASHINFER"

301 or self.attn_backend.get_name() == "TRITON_MLA"

vllm/model_executor/layers/attention/mla_attention.py

+2 -3

@@ -227,7 +227,6 @@ from vllm.model_executor.layers.attention.kv_transfer_utils import (

227 maybe_transfer_kv_layer,

228 )

229 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

230 -from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant

230 from vllm.model_executor.layers.linear import (

231 ColumnParallelLinear,

232 )

@@ -372,7 +371,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

371 if (

372 cache_config is not None

373 and cache_config.enable_prefix_caching

375 - and vllm_is_batch_invariant()

374 + and envs.VLLM_BATCH_INVARIANT

375 and (

376 self.attn_backend.get_name() == "TRITON_MLA"

377 or self.attn_backend.get_name() == "FLASHINFER"

@@ -2188,7 +2187,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2187 # ROCm leverages the upstream flash_attn, which takes a parameter

2188 # called "return_attn_probs" instead of return_softmax_lse

2189 kwargs["return_attn_probs"] = return_softmax_lse

2191 - if vllm_is_batch_invariant():

2190 + if envs.VLLM_BATCH_INVARIANT:

2191 kwargs["num_splits"] = 1

2192

2193 attn_out = self.flash_attn_varlen_func(

vllm/model_executor/layers/batch_invariant.py

+2 -16

@@ -6,6 +6,7 @@ from typing import Any

6

7 import torch

8

9 +import vllm.envs as envs

10 from vllm.logger import init_logger

11 from vllm.platforms import current_platform

12 from vllm.triton_utils import tl, triton

@@ -986,21 +987,6 @@ def enable_batch_invariant_mode():

987 torch.backends.cuda.preferred_blas_library(backend="cublaslt")

988

989

989 -def _read_vllm_batch_invariant() -> bool:

990 - val = os.getenv("VLLM_BATCH_INVARIANT", "0")

991 - try:

992 - return int(val) != 0

993 - except ValueError:

994 - return False

995 -

996 -

997 -VLLM_BATCH_INVARIANT: bool = _read_vllm_batch_invariant()

998 -

999 -

1000 -def vllm_is_batch_invariant() -> bool:

1001 - return VLLM_BATCH_INVARIANT

1002 -

1003 -

990 def override_envs_for_invariance(

991 attention_backend: AttentionBackendEnum | None,

992 ):

@@ -1059,7 +1045,7 @@ def init_batch_invariance(

1045 attention_backend: AttentionBackendEnum | None,

1046 ):

1047 # this will hit all the csrc overrides as well

1062 - if vllm_is_batch_invariant():

1048 + if envs.VLLM_BATCH_INVARIANT:

1049 override_envs_for_invariance(attention_backend)

1050 enable_batch_invariant_mode()

1051

vllm/model_executor/layers/fused_moe/fused_moe.py

+2 -5

@@ -14,9 +14,6 @@ import vllm.envs as envs

14 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

15 from vllm import _custom_ops as ops

16 from vllm.logger import init_logger

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

18 - vllm_is_batch_invariant,

19 -)

17 from vllm.model_executor.layers.fused_moe.activation import (

18 MoEActivation,

19 apply_moe_activation,

@@ -1051,7 +1048,7 @@ def get_moe_configs(

1048 """

1049

1050 # Avoid optimizing for the batch invariant case. Use default config

1054 - if vllm_is_batch_invariant():

1051 + if envs.VLLM_BATCH_INVARIANT:

1052 return None

1053

1054 # First look up if an optimized configuration is available in the configs

@@ -1232,7 +1229,7 @@ def get_default_config(

1229 dtype: str | None,

1230 block_shape: list[int] | None = None,

1231 ) -> dict[str, int]:

1235 - if vllm_is_batch_invariant():

1232 + if envs.VLLM_BATCH_INVARIANT:

1233 return {

1234 "BLOCK_SIZE_M": 64,

1235 "BLOCK_SIZE_N": 64,

vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py

+2 -4

@@ -6,11 +6,9 @@ from collections.abc import Callable

6 import torch

7

8 import vllm._custom_ops as ops

9 +import vllm.envs as envs

10 from vllm._aiter_ops import rocm_aiter_ops

11 from vllm.distributed.eplb.eplb_state import EplbLayerState

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

12 - vllm_is_batch_invariant,

13 -)

12 from vllm.model_executor.layers.fused_moe.config import (

13 RoutingMethodType,

14 get_routing_method_type,

@@ -160,7 +158,7 @@ def fused_topk_bias(

158 ) + e_score_correction_bias.unsqueeze(0)

159

160 # For batch invariance, use sorted=True to ensure deterministic expert selection

163 - use_sorted = vllm_is_batch_invariant()

161 + use_sorted = envs.VLLM_BATCH_INVARIANT

162 topk_indices = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=use_sorted)[1]

163 topk_weights = scores.gather(1, topk_indices)

164 if renormalize:

vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py

+1 -4

@@ -10,9 +10,6 @@ from vllm import envs as envs

10 from vllm._aiter_ops import rocm_aiter_ops

11 from vllm.distributed.eplb.eplb_state import EplbLayerState

12 from vllm.model_executor.custom_op import CustomOp

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

14 - vllm_is_batch_invariant,

15 -)

13 from vllm.model_executor.layers.fused_moe.config import (

14 RoutingMethodType,

15 get_routing_method_type,

@@ -135,7 +132,7 @@ def grouped_topk(

132 ) # [n, n_group]

133

134 # For batch invariance, use sorted=True to ensure deterministic expert selection

138 - use_sorted = vllm_is_batch_invariant()

135 + use_sorted = envs.VLLM_BATCH_INVARIANT

136 group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=use_sorted)[

137 1

138 ] # [n, top_k_group]

vllm/model_executor/layers/layernorm.py

+4 -5

@@ -12,7 +12,6 @@ from vllm.logger import init_logger

12 from vllm.model_executor.custom_op import CustomOp

13 from vllm.model_executor.layers.batch_invariant import (

14 rms_norm_batch_invariant,

15 - vllm_is_batch_invariant,

15 )

16 from vllm.platforms import current_platform

17

@@ -57,7 +56,7 @@ def rms_norm(

56 ) -> torch.Tensor:

57 from vllm import _custom_ops as ops

58

60 - if vllm_is_batch_invariant():

59 + if envs.VLLM_BATCH_INVARIANT:

60 return rms_norm_batch_invariant(x, weight, variance_epsilon)

61 out = torch.empty_like(x)

62 ops.rms_norm(

@@ -77,7 +76,7 @@ def fused_add_rms_norm(

76 ) -> tuple[torch.Tensor, torch.Tensor]:

77 from vllm import _custom_ops as ops

78

80 - if vllm_is_batch_invariant():

79 + if envs.VLLM_BATCH_INVARIANT:

80 return rms_norm_batch_invariant(

81 x + residual, weight, variance_epsilon

82 ), x + residual

@@ -300,7 +299,7 @@ class RMSNorm(CustomOp):

299 and x.is_cuda

300 and x.dim() >= 2

301 and self.has_weight

303 - and not vllm_is_batch_invariant()

302 + and not envs.VLLM_BATCH_INVARIANT

303 and self.weight.data.dtype == x.dtype

304 and self.weight.data.is_contiguous()

305 ):

@@ -328,7 +327,7 @@ class RMSNorm(CustomOp):

327 and x.dtype == residual.dtype

328 and x.dim() >= 2

329 and self.has_weight

331 - and not vllm_is_batch_invariant()

330 + and not envs.VLLM_BATCH_INVARIANT

331 and self.weight.data.dtype == x.dtype

332 and self.weight.data.is_contiguous()

333 ):

vllm/model_executor/layers/linear.py

+2 -2

@@ -7,6 +7,7 @@ from abc import abstractmethod

7 import torch

8 from torch.nn.parameter import Parameter, UninitializedParameter

9

10 +import vllm.envs as envs

11 from vllm.distributed import (

12 divide,

13 get_tensor_model_parallel_rank,

@@ -19,7 +20,6 @@ from vllm.logger import init_logger

20 from vllm.model_executor.custom_op import PluggableLayer

21 from vllm.model_executor.layers.batch_invariant import (

22 linear_batch_invariant,

22 - vllm_is_batch_invariant,

23 )

24 from vllm.model_executor.layers.quantization.base_config import (

25 QuantizationConfig,

@@ -223,7 +223,7 @@ class UnquantizedLinearMethod(LinearMethodBase):

223 x: torch.Tensor,

224 bias: torch.Tensor | None = None,

225 ) -> torch.Tensor:

226 - if vllm_is_batch_invariant() and current_platform.is_cuda_alike():

226 + if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike():

227 return linear_batch_invariant(x, layer.weight, bias)

228 return dispatch_unquantized_gemm()(layer, x, layer.weight, bias)

229

vllm/model_executor/layers/quantization/fp8.py

+2 -4

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

7 from torch.nn import Module

8 from torch.utils._python_dispatch import TorchDispatchMode

9

10 +import vllm.envs as envs

11 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

12 from vllm import _custom_ops as ops

13 from vllm._aiter_ops import rocm_aiter_ops

@@ -17,9 +18,6 @@ from vllm.model_executor.kernels.linear import (

18 )

19 from vllm.model_executor.kernels.linear.scaled_mm import MarlinFP8ScaledMMLinearKernel

20 from vllm.model_executor.layers.attention import Attention

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

21 - vllm_is_batch_invariant,

22 -)

21 from vllm.model_executor.layers.fused_moe import (

22 FusedMoE,

23 FusedMoEMethodBase,

@@ -441,7 +439,7 @@ class Fp8LinearMethod(LinearMethodBase):

439 ) -> torch.Tensor:

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

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

444 - if vllm_is_batch_invariant():

442 + if envs.VLLM_BATCH_INVARIANT:

443 if self.block_quant:

444 assert self.weight_block_size is not None

445 return self.w8a8_block_fp8_linear.apply(

vllm/model_executor/layers/quantization/utils/fp8_utils.py

+1 -3

@@ -305,9 +305,7 @@ def _flashinfer_fp8_blockscale_gemm_impl(

305 )

306 return output

307

308 - from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant

309 -

310 - if vllm_is_batch_invariant():

308 + if envs.VLLM_BATCH_INVARIANT:

309 return run_deepgemm(input, weight, weight_scale)

310

311 condition = input.shape[0] < 32

vllm/utils/flashinfer.py

+1 -4

@@ -19,9 +19,6 @@ import torch

19

20 import vllm.envs as envs

21 from vllm.logger import init_logger

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

23 - vllm_is_batch_invariant,

24 -)

22 from vllm.platforms import current_platform

23

24 logger = init_logger(__name__)

@@ -289,7 +286,7 @@ def supports_trtllm_attention() -> bool:

286 NVIDIA artifactory is accessible, and batch-invariant mode is not enabled.

287 """

288 # Batch-invariant mode disables TRTLLM attention

292 - if vllm_is_batch_invariant():

289 + if envs.VLLM_BATCH_INVARIANT:

290 return False

291

292 # Requires SM100 and NVIDIA artifactory to be accessible to download cubins

vllm/v1/attention/backends/fa_utils.py

+2 -2

@@ -3,8 +3,8 @@

3

4 from typing import Any

5

6 +import vllm.envs as envs

7 from vllm.logger import init_logger

7 -from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant

8 from vllm.platforms import current_platform

9

10 logger = init_logger(__name__)

@@ -114,7 +114,7 @@ def get_flash_attn_version(

114

115 # FA4 currently uses batch-shape-dependent scheduling

116 # heuristics on SM100+, which breaks batch invariance.

117 - if vllm_is_batch_invariant() and fa_version == 4:

117 + if envs.VLLM_BATCH_INVARIANT and fa_version == 4:

118 logger.warning_once(

119 "Cannot use FA version 4 with batch invariance, "

120 "defaulting to FA version 2.",

vllm/v1/attention/backends/flash_attn.py

+5 -7

@@ -33,6 +33,7 @@ if is_flash_attn_varlen_func_available():

33 get_scheduler_metadata,

34 reshape_and_cache_flash,

35 )

36 +import vllm.envs as envs

37 from vllm.config import (

38 VllmConfig,

39 get_current_vllm_config,

@@ -42,9 +43,6 @@ from vllm.config import (

43 from vllm.config.cache import CacheDType

44 from vllm.distributed.parallel_state import get_dcp_group

45 from vllm.logger import init_logger

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

46 - vllm_is_batch_invariant,

47 -)

46 from vllm.platforms.interface import DeviceCapability

47 from vllm.utils.math_utils import cdiv, round_up

48 from vllm.v1.attention.backend import (

@@ -402,7 +400,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

400 # we only set num_splits when using cuda graphs.

401 max_num_splits = self.max_num_splits

402

405 - if vllm_is_batch_invariant():

403 + if envs.VLLM_BATCH_INVARIANT:

404 max_num_splits = 1

405

406 def schedule(

@@ -601,7 +599,7 @@ class FlashAttentionImpl(AttentionImpl):

599 scope="local",

600 )

601 # Cache the batch invariant result for use in forward passes

604 - self.batch_invariant_enabled = vllm_is_batch_invariant()

602 + self.batch_invariant_enabled = envs.VLLM_BATCH_INVARIANT

603

604 if is_quantized_kv_cache(self.kv_cache_dtype) and not flash_attn_supports_fp8():

605 raise NotImplementedError(

@@ -1124,7 +1122,7 @@ def cascade_attention(

1122 # s_aux is incorporated into prefix_lse inside the GPU kernel,

1123 # enabling its effect during the final attention merge.

1124 s_aux=s_aux,

1127 - num_splits=1 if vllm_is_batch_invariant() else max_num_splits,

1125 + num_splits=1 if envs.VLLM_BATCH_INVARIANT else max_num_splits,

1126 )

1127

1128 descale_shape = (cu_query_lens.shape[0] - 1, key_cache.shape[-2])

@@ -1149,7 +1147,7 @@ def cascade_attention(

1147 q_descale=q_descale.expand(descale_shape) if q_descale is not None else None,

1148 k_descale=k_descale.expand(descale_shape) if k_descale is not None else None,

1149 v_descale=v_descale.expand(descale_shape) if v_descale is not None else None,

1152 - num_splits=1 if vllm_is_batch_invariant() else max_num_splits,

1150 + num_splits=1 if envs.VLLM_BATCH_INVARIANT else max_num_splits,

1151 )

1152

1153 # Merge prefix and suffix outputs, and store the result in output.

vllm/v1/attention/backends/flashinfer.py

+2 -5

@@ -28,9 +28,6 @@ from vllm.config import (

28 from vllm.config.cache import CacheDType

29 from vllm.distributed.parallel_state import get_dcp_group

30 from vllm.logger import init_logger

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

32 - vllm_is_batch_invariant,

33 -)

31 from vllm.model_executor.layers.quantization.utils.quant_utils import (

32 QuantKey,

33 kFp8StaticTensorSym,

@@ -544,7 +541,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

541 ) = None # Wrapper for prefill/append

542 self._decode_wrapper = None # Wrapper for decode (general shape)

543

547 - if vllm_is_batch_invariant():

544 + if envs.VLLM_BATCH_INVARIANT:

545 self.decode_fixed_split_size = 2048

546 self.prefill_fixed_split_size = 4096

547 self.disable_split_kv = True

@@ -719,7 +716,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

716 def _get_workspace_buffer(self):

717 if self._workspace_buffer is None:

718 buffer_size = envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE

722 - if vllm_is_batch_invariant():

719 + if envs.VLLM_BATCH_INVARIANT:

720 buffer_size = FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT

721 self._workspace_buffer = torch.zeros(

722 buffer_size, dtype=torch.uint8, device=self.device

vllm/v1/attention/backends/flex_attention.py

+2 -4

@@ -20,12 +20,10 @@ from torch.nn.attention.flex_attention import (

20 or_masks,

21 )

22

23 +import vllm.envs as envs

24 from vllm.config import VllmConfig

25 from vllm.config.cache import CacheDType

26 from vllm.logger import init_logger

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

27 - vllm_is_batch_invariant,

28 -)

27 from vllm.platforms import current_platform

28 from vllm.utils.math_utils import cdiv

29 from vllm.utils.torch_utils import is_torch_equal_or_newer

@@ -995,7 +993,7 @@ def get_kernel_options(

993 return block_size

994 return candidate

995

998 - if vllm_is_batch_invariant():

996 + if envs.VLLM_BATCH_INVARIANT:

997 kernel_options["BLOCK_M"] = 16

998 kernel_options["BLOCK_N"] = 16

999 kernel_options["IS_DIVISIBLE"] = False

vllm/v1/attention/backends/mla/flashattn_mla.py

+3 -5

@@ -6,6 +6,7 @@ from typing import ClassVar

6

7 import torch

8

9 +import vllm.envs as envs

10 from vllm.config import VllmConfig

11 from vllm.config.cache import CacheDType

12 from vllm.logger import init_logger

@@ -17,9 +18,6 @@ from vllm.model_executor.layers.attention.mla_attention import (

18 MLACommonMetadataBuilder,

19 QueryLenSupport,

20 )

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

21 - vllm_is_batch_invariant,

22 -)

21 from vllm.platforms.interface import DeviceCapability

22 from vllm.utils.math_utils import round_up

23 from vllm.v1.attention.backend import (

@@ -152,7 +150,7 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]

150 vllm_config.attention_config.flash_attn_max_num_splits_for_cuda_graph

151 )

152

155 - if vllm_is_batch_invariant():

153 + if envs.VLLM_BATCH_INVARIANT:

154 self.max_num_splits = 1

155

156 def _schedule_decode(

@@ -209,7 +207,7 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]

207 # we only set num_splits when using cuda graphs.

208 max_num_splits = self.max_num_splits

209

212 - if vllm_is_batch_invariant():

210 + if envs.VLLM_BATCH_INVARIANT:

211 max_num_splits = 1

212

213 scheduler_metadata = self._schedule_decode(

vllm/v1/attention/backends/mla/flashmla.py

+2 -4

@@ -6,6 +6,7 @@ from typing import ClassVar

6

7 import torch

8

9 +import vllm.envs as envs

10 from vllm.config import VllmConfig

11 from vllm.config.cache import CacheDType

12 from vllm.logger import init_logger

@@ -17,9 +18,6 @@ from vllm.model_executor.layers.attention.mla_attention import (

18 MLACommonMetadataBuilder,

19 QueryLenSupport,

20 )

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

21 - vllm_is_batch_invariant,

22 -)

21 from vllm.platforms.interface import DeviceCapability

22 from vllm.utils.platform_utils import num_compute_units

23 from vllm.v1.attention.backend import (

@@ -256,7 +254,7 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

254 q = reshape_query_for_spec_decode(q, num_decodes)

255

256 scheduler_metadata = attn_metadata.decode.scheduler_metadata

259 - if vllm_is_batch_invariant() and not self.kv_cache_dtype.startswith("fp8"):

257 + if envs.VLLM_BATCH_INVARIANT and not self.kv_cache_dtype.startswith("fp8"):

258 device = q.device

259 dtype = torch.int32

260

vllm/v1/attention/backends/mla/triton_mla.py

+2 -4

@@ -5,6 +5,7 @@ from typing import ClassVar

5

6 import torch

7

8 +import vllm.envs as envs

9 from vllm.config.cache import CacheDType

10 from vllm.logger import init_logger

11 from vllm.model_executor.layers.attention.mla_attention import (

@@ -12,9 +13,6 @@ from vllm.model_executor.layers.attention.mla_attention import (

13 MLACommonImpl,

14 MLACommonMetadata,

15 )

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

16 - vllm_is_batch_invariant,

17 -)

16 from vllm.platforms.interface import DeviceCapability

17 from vllm.v1.attention.backend import (

18 AttentionLayer,

@@ -151,7 +149,7 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

149 lse = torch.zeros(B, q_num_heads, dtype=q.dtype, device=q.device)

150

151 # For batch invariance, use only 1 split to ensure deterministic reduction

154 - num_kv_splits = 1 if vllm_is_batch_invariant() else 4

152 + num_kv_splits = 1 if envs.VLLM_BATCH_INVARIANT else 4

153

154 # TODO(lucas) Allocate ahead of time

155 attn_logits = torch.empty(

vllm/v1/attention/ops/triton_unified_attention.py

+2 -2

@@ -9,13 +9,13 @@

9

10 import torch

11

12 +import vllm.envs as envs

13 from vllm.logger import init_logger

13 -from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant

14 from vllm.platforms import current_platform

15 from vllm.triton_utils import tl, triton

16

17 logger = init_logger(__name__)

18 -is_batch_invariant = vllm_is_batch_invariant()

18 +is_batch_invariant = envs.VLLM_BATCH_INVARIANT

19 float8_info = torch.finfo(current_platform.fp8_dtype())

20

21

[Bugfix] Register VLLM_BATCH_INVARIANT in envs.py to fix spurious unknown env var warning (#35007)