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 3, 22:45 UTC
[Refactor] Cleanup batch invariant dead code (#41993)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
May 11, 2026 at 10:48 UTC
4b64fc2cbf2a10c0027f7c80681af0ede0c3799a
2 files changed
+4
-35
csrc/core/batch_invariant.hpp
-1
@@ -1,7 +1,6 @@
1
#pragma once
2
#include <cstdlib>
3
#include <string>
4
-#include <cctype>
4
5
namespace vllm {
6
vllm/model_executor/layers/batch_invariant.py
+4
-34
@@ -8,15 +8,12 @@ from typing import Any
8
import torch
9
10
import vllm.envs as envs
11
-from vllm.logger import init_logger
11
from vllm.platforms import current_platform
12
from vllm.triton_utils import tl, triton
13
from vllm.utils.mem_utils import get_max_shared_memory_bytes
14
from vllm.utils.platform_utils import num_compute_units
15
from vllm.utils.torch_utils import is_torch_equal_or_newer
16
18
-logger = init_logger(__name__)
19
-
17
18
def _matmul_launch_metadata(
19
grid: Callable[..., Any], kernel: Any, args: dict[str, Any]
@@ -24,15 +21,8 @@ def _matmul_launch_metadata(
21
ret = {}
22
m, n, k = args["M"], args["N"], args["K"]
23
ret["name"] = f"{kernel.name} [M={m}, N={n}, K={k}]"
27
- if "tiles_per_update" in args:
28
- ret["name"] = (
29
- f"{kernel.name} [M={m}, N={n}, K={k}, "
30
- f"tiles_per_update={args['tiles_per_update']:02}]"
31
- )
32
- if "c_ptr" in args:
33
- bytes_per_elem = args["c_ptr"].element_size()
34
- else:
35
- bytes_per_elem = 1 if args["FP8_OUTPUT"] else 2
24
+
25
+ bytes_per_elem = args["c_ptr"].element_size()
26
ret[f"flops{bytes_per_elem * 8}"] = 2.0 * m * n * k
27
ret["bytes"] = bytes_per_elem * (m * k + n * k + m * n)
28
return ret
@@ -191,7 +181,6 @@ def matmul_persistent(
181
"num_warps": 8,
182
},
183
}
194
- # print(a.device, b.device, c.device)
184
matmul_kernel_persistent[grid](
185
a,
186
b,
@@ -420,7 +409,7 @@ def log_softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor:
409
input: Input tensor
410
dim: Dimension along which to compute log_softmax
411
(only -1 or last dim supported)
423
- >> Stashed changes
412
+
413
Returns:
414
Tensor with log_softmax applied along the specified dimension
415
"""
@@ -910,18 +899,11 @@ def linear_batch_invariant(input, weight, bias=None):
899
900
_batch_invariant_MODE = False
901
_batch_invariant_LIB = None
913
-_original_torch_bmm = None
914
-_original_fp16_reduction_precision = None
915
-_original_bf16_reduction_precision = None
916
-_original_cublas_workspace_cfg = None
917
-_original_cublaslt_workspace_size = None
902
_fp16_block_size_n = 256
903
904
905
def enable_batch_invariant_mode():
922
- global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm
923
- global _original_fp16_reduction_precision, _original_bf16_reduction_precision
924
- global _original_cublas_workspace_cfg, _original_cublaslt_workspace_size
906
+ global _batch_invariant_MODE, _batch_invariant_LIB
907
global _fp16_block_size_n
908
909
if _batch_invariant_MODE:
@@ -941,10 +923,6 @@ def enable_batch_invariant_mode():
923
# Hopper (SM90) and Blackwell (SM100): the only source of batch
924
# variance is split-k, which we disable via the cuBLAS workspace
925
# config.
944
- _original_cublas_workspace_cfg = os.environ.get("CUBLAS_WORKSPACE_CONFIG", None)
945
- _original_cublaslt_workspace_size = os.environ.get(
946
- "CUBLASLT_WORKSPACE_SIZE", None
947
- )
926
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
927
os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1"
928
@@ -966,16 +944,8 @@ def enable_batch_invariant_mode():
944
_batch_invariant_LIB.impl(
945
"aten::bmm", bmm_batch_invariant, "CUDA", allow_override=True
946
)
969
- _original_torch_bmm = torch.bmm
947
torch.bmm = bmm_batch_invariant
948
972
- _original_bf16_reduction_precision = (
973
- torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction
974
- )
975
- _original_fp16_reduction_precision = (
976
- torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction
977
- )
978
-
949
reduced_precision_val = (
950
(False, False) if is_torch_equal_or_newer("2.10.0") else False
951
)