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 7, 22:45 UTC
[Bugfix] Fix kernel benchmark (#33752)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Feb 23, 2026 at 13:18 UTC
7291d1b288558d48508e1a17c37b0aa170332264
14 files changed
+26
benchmarks/fused_kernels/layernorm_rms_benchmarks.py
+2
@@ -13,6 +13,7 @@ from torch.utils.benchmark import Measurement as TMeasurement
13
from tqdm import tqdm
14
15
import vllm._custom_ops as ops
16
+from vllm.benchmarks.lib.utils import default_vllm_config
17
from vllm.model_executor.layers.layernorm import RMSNorm
18
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
19
per_token_group_quant_fp8,
@@ -291,6 +292,7 @@ def print_timers(timers: Iterable[TMeasurement]):
292
compare.print()
293
294
295
+@default_vllm_config()
296
def main():
297
torch.set_default_device("cuda")
298
bench_params = get_bench_params()
benchmarks/kernels/benchmark_activation.py
+2
@@ -7,6 +7,7 @@ import itertools
7
import torch
8
9
import vllm.model_executor.layers.activation # noqa F401
10
+from vllm.benchmarks.lib.utils import default_vllm_config
11
from vllm.model_executor.custom_op import op_registry
12
from vllm.triton_utils import triton
13
from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -18,6 +19,7 @@ intermediate_size = [3072, 9728, 12288]
19
configs = list(itertools.product(batch_size_range, seq_len_range, intermediate_size))
20
21
22
+@default_vllm_config()
23
def benchmark_activation(
24
batch_size: int,
25
seq_len: int,
benchmarks/kernels/benchmark_block_fp8_gemm.py
renamed
+2
@@ -8,6 +8,7 @@ os.environ["VLLM_USE_DEEP_GEMM"] = "0"
8
9
import torch
10
11
+from vllm.benchmarks.lib.utils import default_vllm_config
12
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
13
W8A8BlockFp8LinearOp,
14
)
@@ -40,6 +41,7 @@ DEEPSEEK_V3_SHAPES = [
41
]
42
43
44
+@default_vllm_config()
45
def build_w8a8_block_fp8_runner(M, N, K, block_size, device, use_cutlass):
46
"""Build runner function for w8a8 block fp8 matmul."""
47
factor_for_scale = 1e-2
benchmarks/kernels/benchmark_fp8_gemm.py
renamed
benchmarks/kernels/benchmark_int8_gemm.py
renamed
benchmarks/kernels/benchmark_layernorm.py
+2
@@ -5,12 +5,14 @@ import time
5
6
import torch
7
8
+from vllm.benchmarks.lib.utils import default_vllm_config
9
from vllm.model_executor.layers.layernorm import RMSNorm
10
from vllm.utils.argparse_utils import FlexibleArgumentParser
11
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
12
13
14
@torch.inference_mode()
15
+@default_vllm_config()
16
def main(
17
num_tokens: int,
18
hidden_size: int,
benchmarks/kernels/benchmark_mrope.py
+2
@@ -36,6 +36,7 @@ from typing import Any
36
import numpy as np
37
import torch
38
39
+from vllm.benchmarks.lib.utils import default_vllm_config
40
from vllm.model_executor.layers.rotary_embedding import get_rope
41
from vllm.transformers_utils.config import get_config
42
from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -78,6 +79,7 @@ def calculate_stats(times: list[float]) -> dict[str, float]:
79
}
80
81
82
+@default_vllm_config()
83
def benchmark_mrope(
84
model_name: str,
85
num_tokens: int,
benchmarks/kernels/benchmark_mxfp4_qutlass.py
renamed
benchmarks/kernels/benchmark_nvfp4_gemm.py
renamed
benchmarks/kernels/benchmark_nvfp4_quant.py
renamed
benchmarks/kernels/benchmark_nvfp4_qutlass.py
renamed
benchmarks/kernels/benchmark_per_token_quant_fp8.py
renamed
+2
@@ -7,6 +7,7 @@ from unittest.mock import patch
7
import pandas as pd
8
import torch
9
10
+from vllm.benchmarks.lib.utils import default_vllm_config
11
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
12
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
13
from vllm.triton_utils import triton
@@ -84,6 +85,7 @@ def calculate_diff(
85
configs = []
86
87
88
+@default_vllm_config()
89
def benchmark_quantization(
90
batch_size,
91
hidden_size,
benchmarks/kernels/benchmark_rope.py
+2
@@ -5,6 +5,7 @@ import itertools
5
6
import torch
7
8
+from vllm.benchmarks.lib.utils import default_vllm_config
9
from vllm.model_executor.layers.rotary_embedding import get_rope
10
from vllm.triton_utils import triton
11
from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -29,6 +30,7 @@ def get_benchmark(head_size, rotary_dim, is_neox_style, device):
30
args={},
31
)
32
)
33
+ @default_vllm_config()
34
def benchmark(batch_size, seq_len, num_heads, provider):
35
dtype = torch.bfloat16
36
max_position = 8192
vllm/benchmarks/lib/utils.py
+12
@@ -5,6 +5,7 @@ import argparse
5
import json
6
import math
7
import os
8
+from contextlib import contextmanager
9
from typing import Any
10
11
@@ -117,3 +118,14 @@ def write_to_json(filename: str, records: list) -> None:
118
cls=InfEncoder,
119
default=lambda o: f"<{type(o).__name__} is not JSON serializable>",
120
)
121
+
122
+
123
+@contextmanager
124
+def default_vllm_config():
125
+ """Set a default VllmConfig for cases that directly test CustomOps or pathways
126
+ that use get_current_vllm_config() outside of a full engine context.
127
+ """
128
+ from vllm.config import VllmConfig, set_current_vllm_config
129
+
130
+ with set_current_vllm_config(VllmConfig()):
131
+ yield