[Refactor] Refactor batch invariance folder (#53619)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Aug 24, 2026 at 17:31 UTC
f59bb0bd8c5ad8a0144ae9abb123df9b2e9425a6
11 files changed
+12
-10
.github/CODEOWNERS
+1
-1
@@ -11,7 +11,7 @@
11
/vllm/model_executor/layers/mamba @tdoublep @tomeras91
12
/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
13
/vllm/model_executor/model_loader @22quinn
14
-/vllm/model_executor/layers/batch_invariant.py @yewentao256
14
+/vllm/model_executor/determinism/ @yewentao256
15
/vllm/model_executor/custom_op.py @shen-shanshan
16
/vllm/ir @ProExpertProg
17
/vllm/kernels/ @ProExpertProg @tjtanaa
examples/rl/rlhf_async_new_apis.py
+1
-1
@@ -123,7 +123,7 @@ class TrainModel:
123
"""Ray actor that wraps the training model on a dedicated GPU."""
124
125
def __init__(self, model_name: str):
126
- from vllm.model_executor.layers.batch_invariant import (
126
+ from vllm.model_executor.determinism.batch_invariant import (
127
init_batch_invariance,
128
)
129
tests/v1/determinism/test_matmul_batch_invariant.py
+2
-2
@@ -11,8 +11,8 @@ import pytest
11
import torch
12
from utils import skip_unsupported
13
14
-from vllm.model_executor.layers.batch_invariant import matmul_batch_invariant
15
-from vllm.model_executor.layers.batch_invariant_configs import (
14
+from vllm.model_executor.determinism.batch_invariant import matmul_batch_invariant
15
+from vllm.model_executor.determinism.batch_invariant_configs import (
16
_BATCH_INVARIANT_MATMUL_TUNED_CONFIGS,
17
_get_tuned_matmul_arch_family,
18
)
tests/v1/determinism/test_rms_norm_batch_invariant.py
+1
-1
@@ -6,7 +6,7 @@ import pytest
6
import torch
7
from utils import skip_if_not_cuda, skip_unsupported
8
9
-from vllm.model_executor.layers.batch_invariant import (
9
+from vllm.model_executor.determinism.batch_invariant import (
10
rms_norm_batch_invariant,
11
)
12
from vllm.platforms import current_platform
vllm/model_executor/determinism/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/model_executor/determinism/batch_invariant.py
renamed
+1
-1
@@ -8,7 +8,7 @@ from typing import Any
8
import torch
9
10
import vllm.envs as envs
11
-from vllm.model_executor.layers.batch_invariant_configs import (
11
+from vllm.model_executor.determinism.batch_invariant_configs import (
12
_get_matmul_config,
13
resolve_tuned_matmul_configs,
14
)
vllm/model_executor/determinism/batch_invariant_configs.py
renamed
vllm/model_executor/layers/layernorm.py
+1
-1
@@ -11,7 +11,7 @@ import vllm.kernels # noqa: F401
11
from vllm import envs, ir
12
from vllm.logger import init_logger
13
from vllm.model_executor.custom_op import CustomOp
14
-from vllm.model_executor.layers.batch_invariant import rms_norm_batch_invariant
14
+from vllm.model_executor.determinism.batch_invariant import rms_norm_batch_invariant
15
16
logger = init_logger(__name__)
17
vllm/model_executor/layers/linear.py
+1
-1
@@ -21,7 +21,7 @@ from vllm.distributed import (
21
)
22
from vllm.logger import init_logger
23
from vllm.model_executor.custom_op import PluggableLayer
24
-from vllm.model_executor.layers.batch_invariant import (
24
+from vllm.model_executor.determinism.batch_invariant import (
25
linear_batch_invariant,
26
)
27
from vllm.model_executor.layers.quantization.base_config import (
vllm/model_executor/layers/vocab_parallel_embedding.py
+1
-1
@@ -16,7 +16,7 @@ from vllm.distributed import (
16
tensor_model_parallel_all_reduce,
17
)
18
from vllm.model_executor.custom_op import PluggableLayer
19
-from vllm.model_executor.layers.batch_invariant import (
19
+from vllm.model_executor.determinism.batch_invariant import (
20
linear_batch_invariant,
21
)
22
from vllm.model_executor.layers.quantization.base_config import (
vllm/v1/worker/gpu_worker.py
+1
-1
@@ -1416,7 +1416,7 @@ def init_worker_distributed_environment(
1416
) -> None:
1417
"""Initialize the distributed environment."""
1418
parallel_config = vllm_config.parallel_config
1419
- from vllm.model_executor.layers.batch_invariant import init_batch_invariance
1419
+ from vllm.model_executor.determinism.batch_invariant import init_batch_invariance
1420
1421
init_batch_invariance()
1422
override_envs_for_eplb(