b2f78cbad4 · tduncan/Nvidia-Optimization-Stack-VLLM

csrc/core/batch_invariant.hpp

+4 -4

@@ -5,11 +5,11 @@

5

6 namespace vllm {

7

8 -// vllm_kernel_override_batch_invariant(); returns true

9 -// if env VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT=1

10 -inline bool vllm_kernel_override_batch_invariant() {

8 +// vllm_is_batch_invariant(); returns true

9 +// if env VLLM_BATCH_INVARIANT=1

10 +inline bool vllm_is_batch_invariant() {

11 static bool cached = []() {

12 - std::string env_key = "VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT";

12 + std::string env_key = "VLLM_BATCH_INVARIANT";

13 const char* val = std::getenv(env_key.c_str());

14 return (val && std::atoi(val) != 0) ? 1 : 0;

15 }();

csrc/layernorm_kernels.cu

+2 -2

@@ -426,7 +426,7 @@ void fused_add_rms_norm(torch::Tensor& input, // [..., hidden_size]

426 wt_ptr % req_alignment_bytes == 0;

427 bool offsets_are_multiple_of_vector_width =

428 hidden_size % vector_width == 0 && input_stride % vector_width == 0;

429 - bool batch_invariant_launch = vllm::vllm_kernel_override_batch_invariant();

429 + bool batch_invariant_launch = vllm::vllm_is_batch_invariant();

430 if (ptrs_are_aligned && offsets_are_multiple_of_vector_width &&

431 !batch_invariant_launch) {

432 LAUNCH_FUSED_ADD_RMS_NORM(8);

@@ -474,7 +474,7 @@ void poly_norm(torch::Tensor& out, // [..., hidden_size]

474 auto inp_ptr = reinterpret_cast<std::uintptr_t>(input.data_ptr());

475 auto out_ptr = reinterpret_cast<std::uintptr_t>(out.data_ptr());

476 bool ptrs_are_aligned = inp_ptr % 16 == 0 && out_ptr % 16 == 0;

477 - bool batch_invariant_launch = vllm::vllm_kernel_override_batch_invariant();

477 + bool batch_invariant_launch = vllm::vllm_is_batch_invariant();

478 if (ptrs_are_aligned && hidden_size % 8 == 0 && !batch_invariant_launch) {

479 LAUNCH_FUSED_POLY_NORM(8);

480 } else {

csrc/layernorm_quant_kernels.cu

+1 -1

@@ -254,7 +254,7 @@ void fused_add_rms_norm_static_fp8_quant(

254 auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());

255 bool ptrs_are_aligned =

256 inp_ptr % 16 == 0 && res_ptr % 16 == 0 && wt_ptr % 16 == 0;

257 - bool batch_invariant_launch = vllm::vllm_kernel_override_batch_invariant();

257 + bool batch_invariant_launch = vllm::vllm_is_batch_invariant();

258 if (ptrs_are_aligned && hidden_size % 8 == 0 && input_stride % 8 == 0 &&

259 !batch_invariant_launch) {

260 LAUNCH_FUSED_ADD_RMS_NORM(8);

tests/v1/e2e/test_async_sched_and_preempt.py

+1 -1

@@ -39,7 +39,7 @@ def test_preempt_and_async_scheduling_e2e(monkeypatch: pytest.MonkeyPatch):

39

40 with monkeypatch.context() as m:

41 m.setenv("VLLM_ATTENTION_BACKEND", "FLEX_ATTENTION")

42 - # m.setenv("VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT", "1")

42 + # m.setenv("VLLM_BATCH_INVARIANT", "1")

43

44 outputs: list[tuple[str, list]] = []

45 for test_preemption in [False, True]:

tests/v1/generation/test_batch_invariance.py

+12 -12

@@ -19,14 +19,14 @@ hopper_only = pytest.mark.skipif(

19 @pytest.fixture(autouse=True)

20 def enable_batch_invariant_mode():

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

22 - old_value = os.environ.get("VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT")

23 - os.environ["VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT"] = "1"

22 + old_value = os.environ.get("VLLM_BATCH_INVARIANT")

23 + os.environ["VLLM_BATCH_INVARIANT"] = "1"

24 yield

25 # Restore original value after test

26 if old_value is None:

27 - os.environ.pop("VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT", None)

27 + os.environ.pop("VLLM_BATCH_INVARIANT", None)

28 else:

29 - os.environ["VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT"] = old_value

29 + os.environ["VLLM_BATCH_INVARIANT"] = old_value

30

31

32 def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str:

@@ -231,10 +231,10 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(backend):

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

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

233 from vllm.model_executor.layers.batch_invariant import (

234 - vllm_kernel_override_batch_invariant,

234 + vllm_is_batch_invariant,

235 )

236

237 - disable_custom_ar = vllm_kernel_override_batch_invariant()

237 + disable_custom_ar = vllm_is_batch_invariant()

238

239 if disable_custom_ar:

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

@@ -494,8 +494,8 @@ def test_logprobs_WITHOUT_batch_invariance_should_FAIL(backend):

494 os.environ["VLLM_ATTENTION_BACKEND"] = backend

495

496 # CRITICAL: Disable batch invariance for this test

497 - old_value = os.environ.get("VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT")

498 - os.environ["VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT"] = "0"

497 + old_value = os.environ.get("VLLM_BATCH_INVARIANT")

498 + os.environ["VLLM_BATCH_INVARIANT"] = "0"

499

500 try:

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

@@ -687,9 +687,9 @@ def test_logprobs_WITHOUT_batch_invariance_should_FAIL(backend):

687 finally:

688 # Restore original value

689 if old_value is None:

690 - os.environ.pop("VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT", None)

690 + os.environ.pop("VLLM_BATCH_INVARIANT", None)

691 else:

692 - os.environ["VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT"] = old_value

692 + os.environ["VLLM_BATCH_INVARIANT"] = old_value

693

694

695 @hopper_only

@@ -718,10 +718,10 @@ def test_decode_logprobs_match_prefill_logprobs(backend):

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

719

720 from vllm.model_executor.layers.batch_invariant import (

721 - vllm_kernel_override_batch_invariant,

721 + vllm_is_batch_invariant,

722 )

723

724 - disable_custom_ar = vllm_kernel_override_batch_invariant()

724 + disable_custom_ar = vllm_is_batch_invariant()

725

726 if disable_custom_ar:

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

vllm/config/model.py

+2 -2

@@ -21,7 +21,7 @@ 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_kernel_override_batch_invariant,

24 + vllm_is_batch_invariant,

25 )

26 from vllm.platforms import current_platform

27 from vllm.transformers_utils.config import (

@@ -423,7 +423,7 @@ class ModelConfig:

423 video_pruning_rate: float | None,

424 ) -> None:

425 # Enable batch invariance settings if requested

426 - if vllm_kernel_override_batch_invariant():

426 + if vllm_is_batch_invariant():

427 self.enforce_eager = True

428

429 # Set the default seed to 0 in V1.

vllm/config/parallel.py

+2 -2

@@ -15,7 +15,7 @@ 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_kernel_override_batch_invariant,

18 + vllm_is_batch_invariant,

19 )

20 from vllm.platforms import current_platform

21 from vllm.utils import cuda_device_count_stateless, get_open_ports_list

@@ -565,7 +565,7 @@ class ParallelConfig:

565 from vllm.executor.executor_base import ExecutorBase

566

567 # Enable batch invariance settings if requested

568 - if vllm_kernel_override_batch_invariant():

568 + if vllm_is_batch_invariant():

569 self.disable_custom_all_reduce = True

570

571 if (

vllm/distributed/device_communicators/all_reduce_utils.py

+2 -2

@@ -20,7 +20,7 @@ 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_kernel_override_batch_invariant,

23 + vllm_is_batch_invariant,

24 )

25 from vllm.utils import cuda_device_count_stateless, update_environment_variables

26

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

74 is_symmetric_memory_enabled,

75 )

76

77 - if vllm_kernel_override_batch_invariant():

77 + if vllm_is_batch_invariant():

78 return False

79

80 if not is_symmetric_memory_enabled():

vllm/distributed/device_communicators/symm_mem.py

+2 -2

@@ -10,7 +10,7 @@ from vllm.distributed.device_communicators.all_reduce_utils import (

10 )

11 from vllm.logger import init_logger

12 from vllm.model_executor.layers.batch_invariant import (

13 - vllm_kernel_override_batch_invariant,

13 + vllm_is_batch_invariant,

14 )

15 from vllm.platforms import current_platform

16

@@ -103,7 +103,7 @@ class SymmMemCommunicator:

103 return

104 self.force_multimem = force_multimem

105 self.disabled = False

106 - if vllm_kernel_override_batch_invariant():

106 + if vllm_is_batch_invariant():

107 self.disabled = True

108

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

vllm/model_executor/layers/batch_invariant.py

+3 -3

@@ -741,8 +741,8 @@ def get_batch_invariant_attention_block_size() -> AttentionBlockSize:

741 return AttentionBlockSize(block_m=16, block_n=16)

742

743

744 -def vllm_kernel_override_batch_invariant():

745 - env_key = "VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT"

744 +def vllm_is_batch_invariant():

745 + env_key = "VLLM_BATCH_INVARIANT"

746 is_overridden = False

747 val = os.getenv(env_key, "0")

748 try:

@@ -797,7 +797,7 @@ def override_envs_for_invariance():

797

798 def init_batch_invariance():

799 # this will hit all the csrc overrides as well

800 - if vllm_kernel_override_batch_invariant():

800 + if vllm_is_batch_invariant():

801 override_envs_for_invariance()

802 enable_batch_invariant_mode()

803

vllm/model_executor/layers/fused_moe/fused_moe.py

+5 -5

@@ -16,7 +16,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk

16 from vllm import _custom_ops as ops

17 from vllm.logger import init_logger

18 from vllm.model_executor.layers.batch_invariant import (

19 - vllm_kernel_override_batch_invariant,

19 + vllm_is_batch_invariant,

20 )

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

22 FUSED_MOE_UNQUANTIZED_CONFIG,

@@ -841,7 +841,7 @@ def get_moe_configs(

841 """

842

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

844 - if vllm_kernel_override_batch_invariant():

844 + if vllm_is_batch_invariant():

845 return None

846

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

@@ -976,7 +976,7 @@ def get_default_config(

976 dtype: str | None,

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

978 ) -> dict[str, int]:

979 - if vllm_kernel_override_batch_invariant():

979 + if vllm_is_batch_invariant():

980 config = {

981 "BLOCK_SIZE_M": 64,

982 "BLOCK_SIZE_N": 64,

@@ -1136,7 +1136,7 @@ def fused_topk_bias(

1136 ) + e_score_correction_bias.unsqueeze(0)

1137

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

1139 - use_sorted = vllm_kernel_override_batch_invariant()

1139 + use_sorted = vllm_is_batch_invariant()

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

1141 topk_weights = scores.gather(1, topk_indices)

1142 if renormalize:

@@ -1200,7 +1200,7 @@ def grouped_topk(

1200 ) # [n, n_group]

1201

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

1203 - use_sorted = vllm_kernel_override_batch_invariant()

1203 + use_sorted = vllm_is_batch_invariant()

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

1205 1

1206 ] # [n, top_k_group]

vllm/model_executor/layers/layernorm.py

+3 -3

@@ -10,7 +10,7 @@ import vllm.envs as envs

10 from vllm.model_executor.custom_op import CustomOp

11 from vllm.model_executor.layers.batch_invariant import (

12 rms_norm_batch_invariant,

13 - vllm_kernel_override_batch_invariant,

13 + vllm_is_batch_invariant,

14 )

15 from vllm.platforms import current_platform

16 from vllm.utils import direct_register_custom_op

@@ -25,7 +25,7 @@ def rms_norm(

25 ) -> torch.Tensor:

26 from vllm import _custom_ops as ops

27

28 - if vllm_kernel_override_batch_invariant():

28 + if vllm_is_batch_invariant():

29 return rms_norm_batch_invariant(x, weight, variance_epsilon)

30 out = torch.empty_like(x)

31 ops.rms_norm(

@@ -45,7 +45,7 @@ def fused_add_rms_norm(

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

46 from vllm import _custom_ops as ops

47

48 - if vllm_kernel_override_batch_invariant():

48 + if vllm_is_batch_invariant():

49 return rms_norm_batch_invariant(

50 x + residual, weight, variance_epsilon

51 ), x + residual

vllm/model_executor/layers/quantization/fp8.py

+3 -3

@@ -15,7 +15,7 @@ from vllm import _custom_ops as ops

15 from vllm.distributed import get_tensor_model_parallel_world_size

16 from vllm.logger import init_logger

17 from vllm.model_executor.layers.batch_invariant import (

18 - vllm_kernel_override_batch_invariant,

18 + vllm_is_batch_invariant,

19 )

20 from vllm.model_executor.layers.fused_moe import (

21 FusedMoE,

@@ -356,7 +356,7 @@ class Fp8LinearMethod(LinearMethodBase):

356 # Disable marlin for rocm

357 if current_platform.is_rocm():

358 self.use_marlin = False

359 - if vllm_kernel_override_batch_invariant():

359 + if vllm_is_batch_invariant():

360 self.use_marlin = False

361

362 self.use_aiter_and_is_supported = check_aiter_fp8_linear_support()

@@ -540,7 +540,7 @@ class Fp8LinearMethod(LinearMethodBase):

540 bias: torch.Tensor | None = None,

541 ) -> torch.Tensor:

542 # If batch invariant mode is enabled, dequantize and use BF16 compute

543 - if vllm_kernel_override_batch_invariant():

543 + if vllm_is_batch_invariant():

544 # Dequantize FP8 weights to BF16

545 weight_fp8 = layer.weight.to(torch.bfloat16)

546 weight_scale = layer.weight_scale.to(torch.bfloat16)

vllm/v1/attention/backends/flash_attn.py

+5 -5

@@ -35,7 +35,7 @@ from vllm.config import VllmConfig, get_layers_from_vllm_config

35 from vllm.distributed.parallel_state import get_dcp_group

36 from vllm.logger import init_logger

37 from vllm.model_executor.layers.batch_invariant import (

38 - vllm_kernel_override_batch_invariant,

38 + vllm_is_batch_invariant,

39 )

40 from vllm.utils import cdiv

41 from vllm.v1.attention.backends.utils import (

@@ -308,7 +308,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

308 # we only set num_splits when using cuda graphs.

309 max_num_splits = self.max_num_splits

310

311 - if vllm_kernel_override_batch_invariant():

311 + if vllm_is_batch_invariant():

312 max_num_splits = 1

313

314 def schedule(

@@ -484,7 +484,7 @@ class FlashAttentionImpl(AttentionImpl):

484 self.attn_type = attn_type

485 self.vllm_flash_attn_version = get_flash_attn_version()

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

487 - self.batch_invariant_enabled = vllm_kernel_override_batch_invariant()

487 + self.batch_invariant_enabled = vllm_is_batch_invariant()

488

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

490 raise NotImplementedError(

@@ -963,7 +963,7 @@ def cascade_attention(

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

964 # enabling its effect during the final attention merge.

965 s_aux=s_aux,

966 - num_splits=1 if vllm_kernel_override_batch_invariant() else 0,

966 + num_splits=1 if vllm_is_batch_invariant() else 0,

967 )

968

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

@@ -988,7 +988,7 @@ def cascade_attention(

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

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

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

991 - num_splits=1 if vllm_kernel_override_batch_invariant() else 0,

991 + num_splits=1 if vllm_is_batch_invariant() else 0,

992 )

993

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

vllm/v1/attention/backends/flashinfer.py

+3 -3

@@ -25,7 +25,7 @@ from vllm.attention.backends.abstract import (

25 from vllm.config import CUDAGraphMode, VllmConfig

26 from vllm.logger import init_logger

27 from vllm.model_executor.layers.batch_invariant import (

28 - vllm_kernel_override_batch_invariant,

28 + vllm_is_batch_invariant,

29 )

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

31 QuantKey,

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

291 self._prefill_wrapper = None # Wrapper for prefill/append

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

293

294 - if vllm_kernel_override_batch_invariant():

294 + if vllm_is_batch_invariant():

295 self.decode_fixed_split_size = 2048

296 self.prefill_fixed_split_size = 4096

297 self.disable_split_kv = True

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

404 def _get_workspace_buffer(self):

405 if self._workspace_buffer is None:

406 buffer_size = FLASHINFER_WORKSPACE_BUFFER_SIZE

407 - if vllm_kernel_override_batch_invariant():

407 + if vllm_is_batch_invariant():

408 buffer_size = FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT

409 self._workspace_buffer = torch.zeros(

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

vllm/v1/attention/backends/flex_attention.py

+2 -2

@@ -26,7 +26,7 @@ from vllm.attention.backends.abstract import (

26 from vllm.config import VllmConfig

27 from vllm.logger import init_logger

28 from vllm.model_executor.layers.batch_invariant import (

29 - vllm_kernel_override_batch_invariant,

29 + vllm_is_batch_invariant,

30 )

31 from vllm.utils import cdiv, is_torch_equal_or_newer

32 from vllm.v1.attention.backends.utils import (

@@ -863,7 +863,7 @@ def get_kernel_options(

863 kernel_options: dict[str, int | bool] = {

864 "FORCE_USE_FLEX_ATTENTION": True,

865 }

866 - if vllm_kernel_override_batch_invariant():

866 + if vllm_is_batch_invariant():

867 kernel_options["BLOCK_M"] = 16

868 kernel_options["BLOCK_N"] = 16

869 kernel_options["IS_DIVISIBLE"] = False

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

+2 -2

@@ -212,7 +212,7 @@ from vllm.config import VllmConfig, get_current_vllm_config

212 from vllm.distributed.parallel_state import get_dcp_group, is_global_first_rank

213 from vllm.logger import init_logger

214 from vllm.model_executor.layers.batch_invariant import (

215 - vllm_kernel_override_batch_invariant,

215 + vllm_is_batch_invariant,

216 )

217 from vllm.model_executor.layers.linear import (

218 ColumnParallelLinear,

@@ -1283,7 +1283,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

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

1284 # called "return_attn_probs" instead of return_softmax_lse

1285 kwargs["return_attn_probs"] = return_softmax_lse

1286 - if vllm_kernel_override_batch_invariant():

1286 + if vllm_is_batch_invariant():

1287 kwargs["num_splits"] = 1

1288

1289 attn_out = self.flash_attn_varlen_func(

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

+3 -3

@@ -19,7 +19,7 @@ from vllm.attention.utils.fa_utils import (

19 from vllm.config import VllmConfig

20 from vllm.logger import init_logger

21 from vllm.model_executor.layers.batch_invariant import (

22 - vllm_kernel_override_batch_invariant,

22 + vllm_is_batch_invariant,

23 )

24 from vllm.v1.attention.backends.mla.common import (

25 MLACommonBackend,

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

110 # pre-allocated during capture.

111 self.max_num_splits = envs.VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH

112

113 - if vllm_kernel_override_batch_invariant():

113 + if vllm_is_batch_invariant():

114 self.max_num_splits = 1

115

116 def _schedule_decode(

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

181 # we only set num_splits when using cuda graphs.

182 max_num_splits = self.max_num_splits

183

184 - if vllm_kernel_override_batch_invariant():

184 + if vllm_is_batch_invariant():

185 max_num_splits = 1

186

187 metadata = FlashAttnMLADecodeMetadata(

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

+2 -2

@@ -15,7 +15,7 @@ from vllm.attention.ops.flashmla import (

15 from vllm.config import VllmConfig

16 from vllm.logger import init_logger

17 from vllm.model_executor.layers.batch_invariant import (

18 - vllm_kernel_override_batch_invariant,

18 + vllm_is_batch_invariant,

19 )

20 from vllm.v1.attention.backends.mla.common import (

21 MLACommonBackend,

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

234

235 tile_scheduler_metadata = attn_metadata.decode.tile_scheduler_metadata

236 num_splits = attn_metadata.decode.num_splits

237 - if vllm_kernel_override_batch_invariant():

237 + if vllm_is_batch_invariant():

238 device = q.device

239 dtype = torch.int32

240

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

+2 -2

@@ -14,7 +14,7 @@ from vllm.attention.ops.triton_decode_attention import decode_attention_fwd

14 from vllm.attention.ops.triton_flash_attention import triton_attention

15 from vllm.logger import init_logger

16 from vllm.model_executor.layers.batch_invariant import (

17 - vllm_kernel_override_batch_invariant,

17 + vllm_is_batch_invariant,

18 )

19 from vllm.platforms import current_platform

20 from vllm.triton_utils import HAS_TRITON

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

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

164

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

166 - num_kv_splits = 1 if vllm_kernel_override_batch_invariant() else 4

166 + num_kv_splits = 1 if vllm_is_batch_invariant() else 4

167

168 # TODO(lucas) Allocate ahead of time

169 attn_logits = torch.empty(

[small][batch invariance] Rename the env and internal flags to simplify usage (#26855)