@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1085b64425

[Benchmark] Retire stale benchmarks and consolidate RMSNorm (#57083)

Signed-off-by: Misha Goin <mgoin64@gmail.com>

Misha Goin committed Sep 16, 2026 at 13:21 UTC 1085b64425a9e6f5ca52876ad32e55fda5665f4e

7 files changed +3 -780

benchmarks/cutlass_benchmarks/utils.py deleted

-39

@@ -1,39 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -# Cutlass bench utils

5 -

6 -import torch

7 -

8 -

9 -def to_fp8(tensor: torch.Tensor) -> torch.Tensor:

10 - finfo = torch.finfo(torch.float8_e4m3fn)

11 - return torch.round(tensor.clamp(min=finfo.min, max=finfo.max)).to(

12 - dtype=torch.float8_e4m3fn

13 - )

14 -

15 -

16 -def to_int8(tensor: torch.Tensor) -> torch.Tensor:

17 - return torch.round(tensor.clamp(min=-128, max=127)).to(dtype=torch.int8)

18 -

19 -

20 -def to_bf16(tensor: torch.Tensor) -> torch.Tensor:

21 - return tensor.to(dtype=torch.bfloat16)

22 -

23 -

24 -def to_fp16(tensor: torch.Tensor) -> torch.Tensor:

25 - return tensor.to(dtype=torch.float16)

26 -

27 -

28 -def make_rand_tensors(

29 - dtype: torch.dtype, m: int, n: int, k: int

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

31 - a = torch.randn((m, k), device="cuda") * 5

32 - b = torch.randn((n, k), device="cuda").t() * 5

33 -

34 - if dtype == torch.int8:

35 - return to_int8(a), to_int8(b)

36 - if dtype == torch.float8_e4m3fn:

37 - return to_fp8(a), to_fp8(b)

38 -

39 - raise ValueError("unsupported dtype")

benchmarks/cutlass_benchmarks/w8a8_benchmarks.py deleted

-372

@@ -1,372 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -import argparse

5 -import copy

6 -import itertools

7 -import pickle as pkl

8 -import time

9 -from collections.abc import Callable, Iterable

10 -

11 -import torch

12 -import torch.utils.benchmark as TBenchmark

13 -from torch.utils.benchmark import Measurement as TMeasurement

14 -from utils import make_rand_tensors

15 -from weight_shapes import WEIGHT_SHAPES

16 -

17 -from vllm import _custom_ops as ops

18 -from vllm.model_executor.layers.quantization.utils.fp8_utils import (

19 - w8a8_triton_block_scaled_mm,

20 -)

21 -from vllm.utils.argparse_utils import FlexibleArgumentParser

22 -from vllm.utils.math_utils import cdiv

23 -

24 -DEFAULT_MODELS = list(WEIGHT_SHAPES.keys())

25 -DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512]

26 -DEFAULT_TP_SIZES = [1]

27 -

28 -

29 -# bench

30 -def bench_fn(

31 - label: str, sub_label: str, description: str, fn: Callable, *args, **kwargs

32 -) -> TMeasurement:

33 - min_run_time = 1

34 -

35 - globals = {

36 - "args": args,

37 - "kwargs": kwargs,

38 - "fn": fn,

39 - }

40 - return TBenchmark.Timer(

41 - stmt="fn(*args, **kwargs)",

42 - globals=globals,

43 - label=label,

44 - sub_label=sub_label,

45 - description=description,

46 - ).blocked_autorange(min_run_time=min_run_time)

47 -

48 -

49 -def bench_int8(

50 - dtype: torch.dtype,

51 - m: int,

52 - k: int,

53 - n: int,

54 - label: str,

55 - sub_label: str,

56 - bench_kernels: list[str] | None = None,

57 -) -> Iterable[TMeasurement]:

58 - """Benchmark INT8-based kernels."""

59 - assert dtype == torch.int8

60 - a, b = make_rand_tensors(torch.int8, m, n, k)

61 - scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32)

62 - scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32)

63 - bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16)

64 - azp = torch.zeros((m,), device="cuda", dtype=torch.int32)

65 - azp_adj = torch.zeros((n,), device="cuda", dtype=torch.int32)

66 -

67 - bench_fns = {

68 - "pytorch_bf16_bf16_bf16_matmul-no-scales": lambda: torch.mm(

69 - a.to(dtype=torch.bfloat16), b.to(dtype=torch.bfloat16)

70 - ),

71 - "pytorch_fp16_fp16_fp16_matmul-no-scales": lambda: torch.mm(

72 - a.to(dtype=torch.float16), b.to(dtype=torch.float16)

73 - ),

74 - "cutlass_i8_i8_bf16_scaled_mm": lambda: ops.cutlass_scaled_mm(

75 - a, b, scale_a, scale_b, torch.bfloat16

76 - ),

77 - "cutlass_i8_i8_bf16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm(

78 - a, b, scale_a, scale_b, torch.bfloat16, bias

79 - ),

80 - "cutlass_i8_i8_bf16_scaled_mm_azp": lambda: ops.cutlass_scaled_mm_azp(

81 - a, b, scale_a, scale_b, torch.bfloat16, azp_adj

82 - ),

83 - "cutlass_i8_i8_bf16_scaled_mm_azp_bias": lambda: ops.cutlass_scaled_mm_azp(

84 - a, b, scale_a, scale_b, torch.bfloat16, azp_adj, None, bias

85 - ),

86 - "cutlass_i8_i8_bf16_scaled_mm_azp_pt": lambda: ops.cutlass_scaled_mm_azp(

87 - a, b, scale_a, scale_b, torch.bfloat16, azp_adj, azp

88 - ),

89 - "cutlass_i8_i8_bf16_scaled_mm_azp_pt_bias": lambda: ops.cutlass_scaled_mm_azp(

90 - a, b, scale_a, scale_b, torch.bfloat16, azp_adj, azp, bias

91 - ),

92 - }

93 -

94 - timers = []

95 - for name, fn in bench_fns.items():

96 - # If bench_kernels is None, run all. Otherwise, run only exact matches.

97 - if bench_kernels is None or name in bench_kernels:

98 - print(f"Running {name}")

99 - timers.append(bench_fn(label, sub_label, name, fn))

100 -

101 - return timers

102 -

103 -

104 -def bench_fp8(

105 - dtype: torch.dtype,

106 - m: int,

107 - k: int,

108 - n: int,

109 - label: str,

110 - sub_label: str,

111 - bench_kernels: list[str] | None = None,

112 -) -> Iterable[TMeasurement]:

113 - """Benchmark FP8-based kernels."""

114 - assert dtype == torch.float8_e4m3fn

115 - a, b = make_rand_tensors(torch.float8_e4m3fn, m, n, k)

116 - a_cont = a.contiguous()

117 - scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32)

118 - scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32)

119 -

120 - block_scale_a = torch.rand((m, cdiv(k, 128)), device="cuda", dtype=torch.float32)

121 - block_scale_b = torch.rand(

122 - cdiv(k, 128), cdiv(n, 128), device="cuda", dtype=torch.float32

123 - )

124 - block_scale_a_M_major = block_scale_a.t().contiguous().t()

125 - block_scale_b_K_major = block_scale_b.t().contiguous().t()

126 - bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16)

127 -

128 - print(m, k, n)

129 -

130 - bench_fns = {

131 - "pytorch_bf16_bf16_bf16_matmul-no-scales": lambda: torch.mm(

132 - a.to(dtype=torch.bfloat16), b.to(dtype=torch.bfloat16)

133 - ),

134 - "pytorch_fp16_fp16_fp16_matmul-no-scales": lambda: torch.mm(

135 - a.to(dtype=torch.float16), b.to(dtype=torch.float16)

136 - ),

137 - "pytorch_fp8_fp8_fp16_scaled_mm": lambda: torch._scaled_mm(

138 - a, b, scale_a, scale_b, out_dtype=torch.float16

139 - ),

140 - "pytorch_fp8_fp8_fp16_scaled_mm_fast_accum": lambda: torch._scaled_mm(

141 - a, b, scale_a, scale_b, out_dtype=torch.float16, use_fast_accum=True

142 - ),

143 - "pytorch_fp8_fp8_bf16_scaled_mm": lambda: torch._scaled_mm(

144 - a, b, scale_a, scale_b, out_dtype=torch.bfloat16

145 - ),

146 - "pytorch_fp8_fp8_bf16_scaled_mm_fast_accum": lambda: torch._scaled_mm(

147 - a, b, scale_a, scale_b, out_dtype=torch.bfloat16, use_fast_accum=True

148 - ),

149 - "cutlass_fp8_fp8_bf16_scaled_mm": lambda: ops.cutlass_scaled_mm(

150 - a, b, scale_a, scale_b, torch.bfloat16

151 - ),

152 - "cutlass_fp8_fp8_fp16_scaled_mm": lambda: ops.cutlass_scaled_mm(

153 - a, b, scale_a, scale_b, torch.float16

154 - ),

155 - "cutlass_fp8_fp8_bf16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm(

156 - a, b, scale_a, scale_b, torch.bfloat16, bias

157 - ),

158 - "cutlass_fp8_fp8_fp16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm(

159 - a, b, scale_a, scale_b, torch.float16, bias.to(dtype=torch.float16)

160 - ),

161 - "triton_fp8_fp8_fp16_scaled_mm_blockwise": lambda: w8a8_triton_block_scaled_mm(

162 - a_cont, b.t(), block_scale_a, block_scale_b.t(), (128, 128)

163 - ),

164 - "cutlass_fp8_fp8_fp16_scaled_mm_blockwise": lambda: ops.cutlass_scaled_mm(

165 - a, b, block_scale_a_M_major, block_scale_b_K_major, torch.float16

166 - ),

167 - }

168 -

169 - timers = []

170 - for name, fn in bench_fns.items():

171 - # If bench_kernels is None, run all. Otherwise, run only exact matches.

172 - if bench_kernels is None or name in bench_kernels:

173 - print(f"Running {name}")

174 - timers.append(bench_fn(label, sub_label, name, fn))

175 -

176 - return timers

177 -

178 -

179 -def bench(

180 - dtype: torch.dtype,

181 - m: int,

182 - k: int,

183 - n: int,

184 - label: str,

185 - sub_label: str,

186 - bench_kernels: list[str] | None = None,

187 -) -> Iterable[TMeasurement]:

188 - if dtype == torch.int8:

189 - return bench_int8(dtype, m, k, n, label, sub_label, bench_kernels)

190 - if dtype == torch.float8_e4m3fn:

191 - return bench_fp8(dtype, m, k, n, label, sub_label, bench_kernels)

192 - raise ValueError("unsupported type")

193 -

194 -

195 -# runner

196 -def print_timers(timers: Iterable[TMeasurement]):

197 - compare = TBenchmark.Compare(timers)

198 - compare.print()

199 -

200 -

201 -def run(

202 - dtype: torch.dtype,

203 - MKNs: Iterable[tuple[int, int, int]],

204 - bench_kernels: list[str] | None = None,

205 -) -> Iterable[TMeasurement]:

206 - results = []

207 - for m, k, n in MKNs:

208 - timers = bench(

209 - dtype,

210 - m,

211 - k,

212 - n,

213 - f"scaled-{dtype}-gemm",

214 - f"MKN=({m}x{k}x{n})",

215 - bench_kernels=bench_kernels,

216 - )

217 - print_timers(timers)

218 - results.extend(timers)

219 - return results

220 -

221 -

222 -def make_output(

223 - data: Iterable[TMeasurement],

224 - MKNs: Iterable[tuple[int, int, int]],

225 - base_description: str,

226 - timestamp=None,

227 -):

228 - print(f"== All Results {base_description} ====")

229 - print_timers(data)

230 -

231 - # pickle all the results

232 - timestamp = int(time.time()) if timestamp is None else timestamp

233 - with open(f"{base_description}-{timestamp}.pkl", "wb") as f:

234 - pkl.dump(data, f)

235 -

236 -

237 -def run_square_bench(args):

238 - dim_sizes = list(range(args.dim_start, args.dim_end + 1, args.dim_increment))

239 - MKNs = list(zip(dim_sizes, dim_sizes, dim_sizes))

240 - data = run(args.dtype, MKNs, bench_kernels=args.kernels)

241 - make_output(data, MKNs, f"square_bench-{args.dtype}")

242 -

243 -

244 -def run_range_bench(args):

245 - dim_sizes = list(range(args.dim_start, args.dim_end, args.dim_increment))

246 - n = len(dim_sizes)

247 - Ms = [args.m_constant] * n if args.m_constant is not None else dim_sizes

248 - Ks = [args.k_constant] * n if args.k_constant is not None else dim_sizes

249 - Ns = [args.n_constant] * n if args.n_constant is not None else dim_sizes

250 - MKNs = list(zip(Ms, Ks, Ns))

251 - data = run(args.dtype, MKNs, bench_kernels=args.kernels)

252 - make_output(data, MKNs, f"range_bench-{args.dtype}")

253 -

254 -

255 -def run_model_bench(args):

256 - print("Benchmarking models:")

257 - for i, model in enumerate(args.models):

258 - print(f"[{i}] {model}")

259 -

260 - def model_shapes(model_name: str, tp_size: int) -> list[tuple[int, int]]:

261 - KNs = []

262 - for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model_name]):

263 - KN[tp_split_dim] = KN[tp_split_dim] // tp_size

264 - KNs.append(KN)

265 - return KNs

266 -

267 - model_bench_data = []

268 - models_tps = list(itertools.product(args.models, args.tp_sizes))

269 - for model, tp_size in models_tps:

270 - Ms = args.batch_sizes

271 - KNs = model_shapes(model, tp_size)

272 - MKNs = []

273 - for m in Ms:

274 - for k, n in KNs:

275 - MKNs.append((m, k, n))

276 -

277 - data = run(args.dtype, MKNs, bench_kernels=args.kernels)

278 - model_bench_data.append(data)

279 -

280 - # Print all results

281 - for data, model_tp in zip(model_bench_data, models_tps):

282 - model, tp_size = model_tp

283 - print(f"== Results {args.dtype} {model}-TP{tp_size} ====")

284 - print_timers(data)

285 -

286 - timestamp = int(time.time())

287 -

288 - all_data = []

289 - for d in model_bench_data:

290 - all_data.extend(d)

291 - # pickle all data

292 - with open(f"model_bench-{args.dtype}-{timestamp}.pkl", "wb") as f:

293 - pkl.dump(all_data, f)

294 -

295 -

296 -if __name__ == "__main__":

297 -

298 - def to_torch_dtype(dt):

299 - if dt == "int8":

300 - return torch.int8

301 - if dt == "fp8":

302 - return torch.float8_e4m3fn

303 - raise ValueError("unsupported dtype")

304 -

305 - parser = FlexibleArgumentParser(

306 - description="""

307 -Benchmark Cutlass GEMM.

308 -

309 - To run square GEMMs:

310 - python3 ./benchmarks/cutlass_benchmarks/w8a8_benchmarks.py --dtype fp8 square_bench --dim-start 128 --dim-end 512 --dim-increment 64

311 -

312 - To run constant N and K and sweep M:

313 - python3 ./benchmarks/cutlass_benchmarks/w8a8_benchmarks.py --dtype fp8 range_bench --dim-start 128 --dim-end 512 --dim-increment 64 --n-constant 16384 --k-constant 16384

314 -

315 - To run dimensions from a model:

316 - python3 ./benchmarks/cutlass_benchmarks/w8a8_benchmarks.py --dtype fp8 model_bench --models meta-llama/Llama-2-7b-hf --batch-sizes 16 --tp-sizes 1

317 -

318 - Output:

319 - - a .pkl file, that is a list of raw torch.benchmark.utils.Measurements for the pytorch and cutlass implementations for the various GEMMs.

320 - """, # noqa: E501

321 - formatter_class=argparse.RawTextHelpFormatter,

322 - )

323 -

324 - parser.add_argument(

325 - "--dtype",

326 - type=to_torch_dtype,

327 - required=True,

328 - help="Available options are ['int8', 'fp8']",

329 - )

330 - parser.add_argument(

331 - "--kernels",

332 - nargs="+",

333 - type=str,

334 - default=None,

335 - help="Exact names of the kernels to benchmark. If not set, runs all kernels.",

336 - )

337 -

338 - subparsers = parser.add_subparsers(dest="cmd")

339 -

340 - square_parser = subparsers.add_parser("square_bench")

341 - square_parser.add_argument("--dim-start", type=int, required=True)

342 - square_parser.add_argument("--dim-end", type=int, required=True)

343 - square_parser.add_argument("--dim-increment", type=int, required=True)

344 - square_parser.set_defaults(func=run_square_bench)

345 -

346 - range_parser = subparsers.add_parser("range_bench")

347 - range_parser.add_argument("--dim-start", type=int, required=True)

348 - range_parser.add_argument("--dim-end", type=int, required=True)

349 - range_parser.add_argument("--dim-increment", type=int, required=True)

350 - range_parser.add_argument("--m-constant", type=int, default=None)

351 - range_parser.add_argument("--n-constant", type=int, default=None)

352 - range_parser.add_argument("--k-constant", type=int, default=None)

353 - range_parser.set_defaults(func=run_range_bench)

354 -

355 - model_parser = subparsers.add_parser("model_bench")

356 - model_parser.add_argument(

357 - "--models",

358 - nargs="+",

359 - type=str,

360 - default=DEFAULT_MODELS,

361 - choices=WEIGHT_SHAPES.keys(),

362 - )

363 - model_parser.add_argument(

364 - "--tp-sizes", nargs="+", type=int, default=DEFAULT_TP_SIZES

365 - )

366 - model_parser.add_argument(

367 - "--batch-sizes", nargs="+", type=int, default=DEFAULT_BATCH_SIZES

368 - )

369 - model_parser.set_defaults(func=run_model_bench)

370 -

371 - args = parser.parse_args()

372 - args.func(args)

benchmarks/cutlass_benchmarks/weight_shapes.py deleted

-46

@@ -1,46 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -# Weight Shapes are in the format

5 -# ([K, N], TP_SPLIT_DIM)

6 -# Example:

7 -# A shape of ([14336, 4096], 0) indicates the following GEMM shape,

8 -# - TP1 : K = 14336, N = 4096

9 -# - TP2 : K = 7168, N = 4096

10 -# A shape of ([4096, 6144], 1) indicates the following GEMM shape,

11 -# - TP1 : K = 4096, N = 6144

12 -# - TP4 : K = 4096, N = 1536

13 -

14 -# TP1 shapes

15 -WEIGHT_SHAPES = {

16 - "mistralai/Mistral-7B-v0.1": [

17 - ([4096, 6144], 1),

18 - ([4096, 4096], 0),

19 - ([4096, 28672], 1),

20 - ([14336, 4096], 0),

21 - ],

22 - "meta-llama/Llama-2-7b-hf": [

23 - ([4096, 12288], 1),

24 - ([4096, 4096], 0),

25 - ([4096, 22016], 1),

26 - ([11008, 4096], 0),

27 - ],

28 - "meta-llama/Llama-3-8b": [

29 - ([4096, 6144], 1),

30 - ([4096, 4096], 0),

31 - ([4096, 28672], 1),

32 - ([14336, 4096], 0),

33 - ],

34 - "meta-llama/Llama-2-13b-hf": [

35 - ([5120, 15360], 1),

36 - ([5120, 5120], 0),

37 - ([5120, 27648], 1),

38 - ([13824, 5120], 0),

39 - ],

40 - "meta-llama/Llama-2-70b-hf": [

41 - ([8192, 10240], 1),

42 - ([8192, 8192], 0),

43 - ([8192, 57344], 1),

44 - ([28672, 8192], 0),

45 - ],

46 -}

benchmarks/kernels/benchmark_rmsnorm.py deleted

-255

@@ -1,255 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -import itertools

5 -

6 -import torch

7 -from flashinfer.norm import fused_add_rmsnorm, rmsnorm

8 -from torch import nn

9 -

10 -from vllm import _custom_ops as vllm_ops

11 -from vllm.triton_utils import triton

12 -

13 -

14 -class HuggingFaceRMSNorm(nn.Module):

15 - def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:

16 - super().__init__()

17 - self.weight = nn.Parameter(torch.ones(hidden_size))

18 - self.variance_epsilon = eps

19 -

20 - def forward(

21 - self,

22 - x: torch.Tensor,

23 - residual: torch.Tensor | None = None,

24 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

25 - orig_dtype = x.dtype

26 - x = x.to(torch.float32)

27 - if residual is not None:

28 - x = x + residual.to(torch.float32)

29 - residual = x.to(orig_dtype)

30 -

31 - variance = x.pow(2).mean(dim=-1, keepdim=True)

32 - x = x * torch.rsqrt(variance + self.variance_epsilon)

33 - x = x.to(orig_dtype) * self.weight

34 - if residual is None:

35 - return x

36 - else:

37 - return x, residual

38 -

39 -

40 -def rmsnorm_naive(

41 - x: torch.Tensor,

42 - weight: torch.Tensor,

43 - residual: torch.Tensor | None = None,

44 - eps: float = 1e-6,

45 -):

46 - naive_norm = HuggingFaceRMSNorm(x.shape[-1], eps=eps)

47 - naive_norm.weight = nn.Parameter(weight)

48 - naive_norm = naive_norm.to(x.device)

49 -

50 - orig_shape = x.shape

51 - x = x.view(-1, x.shape[-1])

52 - if residual is not None:

53 - residual = residual.view(-1, residual.shape[-1])

54 -

55 - output = naive_norm(x, residual)

56 -

57 - if isinstance(output, tuple):

58 - output = (output[0].view(orig_shape), output[1].view(orig_shape))

59 - else:

60 - output = output.view(orig_shape)

61 - return output

62 -

63 -

64 -def rmsnorm_flashinfer(

65 - x: torch.Tensor,

66 - weight: torch.Tensor,

67 - residual: torch.Tensor | None = None,

68 - eps: float = 1e-6,

69 -):

70 - orig_shape = x.shape

71 - x = x.view(-1, x.shape[-1])

72 - if residual is not None:

73 - residual = residual.view(-1, residual.shape[-1])

74 -

75 - if residual is not None:

76 - fused_add_rmsnorm(x, residual, weight, eps)

77 - output = (x, residual)

78 - else:

79 - output = rmsnorm(x, weight, eps)

80 -

81 - if isinstance(output, tuple):

82 - output = (output[0].view(orig_shape), output[1].view(orig_shape))

83 - else:

84 - output = output.view(orig_shape)

85 - return output

86 -

87 -

88 -def rmsnorm_vllm(

89 - x: torch.Tensor,

90 - weight: torch.Tensor,

91 - residual: torch.Tensor | None = None,

92 - eps: float = 1e-6,

93 -):

94 - orig_shape = x.shape

95 - x = x.view(-1, x.shape[-1])

96 - if residual is not None:

97 - residual = residual.view(-1, residual.shape[-1])

98 -

99 - if residual is not None:

100 - vllm_ops.fused_add_rms_norm(x, residual, weight, eps)

101 - output = (x, residual)

102 - else:

103 - out = torch.empty_like(x)

104 - vllm_ops.rms_norm(out, x, weight, eps)

105 - output = out

106 -

107 - if isinstance(output, tuple):

108 - output = (output[0].view(orig_shape), output[1].view(orig_shape))

109 - else:

110 - output = output.view(orig_shape)

111 - return output

112 -

113 -

114 -def calculate_diff(batch_size, seq_len, hidden_size, use_residual=True):

115 - dtype = torch.bfloat16

116 - x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda")

117 - weight = torch.ones(hidden_size, dtype=dtype, device="cuda")

118 - residual = torch.randn_like(x) if use_residual else None

119 -

120 - output_naive = rmsnorm_naive(

121 - x.clone(), weight, residual.clone() if residual is not None else None

122 - )

123 - output_flashinfer = rmsnorm_flashinfer(

124 - x.clone(), weight, residual.clone() if residual is not None else None

125 - )

126 - output_vllm = rmsnorm_vllm(

127 - x.clone(), weight, residual.clone() if residual is not None else None

128 - )

129 -

130 - if use_residual:

131 - output_naive = output_naive[0]

132 - output_flashinfer = output_flashinfer[0]

133 - output_vllm = output_vllm[0]

134 -

135 - print(f"Naive output={output_naive}")

136 - print(f"FlashInfer output={output_flashinfer}")

137 - print(f"vLLM output={output_vllm}")

138 -

139 - if torch.allclose(

140 - output_naive, output_flashinfer, atol=1e-2, rtol=1e-2

141 - ) and torch.allclose(output_naive, output_vllm, atol=1e-2, rtol=1e-2):

142 - print("✅ All implementations match")

143 - else:

144 - print("❌ Implementations differ")

145 -

146 -

147 -batch_size_range = [2**i for i in range(0, 7, 2)]

148 -seq_length_range = [2**i for i in range(6, 11, 1)]

149 -head_num_range = [32, 48]

150 -configs = list(itertools.product(head_num_range, batch_size_range, seq_length_range))

151 -

152 -

153 -def get_benchmark(use_residual):

154 - @triton.testing.perf_report(

155 - triton.testing.Benchmark(

156 - x_names=["head_num", "batch_size", "seq_len"],

157 - x_vals=[list(_) for _ in configs],

158 - line_arg="provider",

159 - line_vals=["huggingface", "flashinfer", "vllm"],

160 - line_names=["HuggingFace", "FlashInfer", "vLLM"],

161 - styles=[("blue", "-"), ("green", "-"), ("red", "-")],

162 - ylabel="us",

163 - plot_name=f"rmsnorm-perf-{'with' if use_residual else 'without'}-residual",

164 - args={},

165 - )

166 - )

167 - def benchmark(head_num, batch_size, seq_len, provider):

168 - dtype = torch.bfloat16

169 - hidden_size = head_num * 128 # assuming head_dim = 128

170 -

171 - x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda")

172 - weight = torch.ones(hidden_size, dtype=dtype, device="cuda")

173 - residual = torch.randn_like(x) if use_residual else None

174 -

175 - quantiles = [0.5, 0.2, 0.8]

176 -

177 - if provider == "huggingface":

178 - ms, min_ms, max_ms = triton.testing.do_bench(

179 - lambda: rmsnorm_naive(

180 - x.clone(),

181 - weight,

182 - residual.clone() if residual is not None else None,

183 - ),

184 - quantiles=quantiles,

185 - )

186 - elif provider == "flashinfer":

187 - ms, min_ms, max_ms = triton.testing.do_bench(

188 - lambda: rmsnorm_flashinfer(

189 - x.clone(),

190 - weight,

191 - residual.clone() if residual is not None else None,

192 - ),

193 - quantiles=quantiles,

194 - )

195 - else:

196 - ms, min_ms, max_ms = triton.testing.do_bench(

197 - lambda: rmsnorm_vllm(

198 - x.clone(),

199 - weight,

200 - residual.clone() if residual is not None else None,

201 - ),

202 - quantiles=quantiles,

203 - )

204 -

205 - return 1000 * ms, 1000 * max_ms, 1000 * min_ms

206 -

207 - return benchmark

208 -

209 -

210 -if __name__ == "__main__":

211 - import argparse

212 -

213 - parser = argparse.ArgumentParser()

214 - parser.add_argument(

215 - "--batch-size",

216 - type=int,

217 - default=4,

218 - help="Batch size",

219 - )

220 - parser.add_argument(

221 - "--seq-len",

222 - type=int,

223 - default=128,

224 - help="Sequence length",

225 - )

226 - parser.add_argument(

227 - "--hidden-size",

228 - type=int,

229 - default=4096,

230 - help="Hidden size (2nd dimension) of the sequence",

231 - )

232 - parser.add_argument(

233 - "--use-residual", action="store_true", help="Whether to use residual connection"

234 - )

235 - parser.add_argument(

236 - "--save-path",

237 - type=str,

238 - default="./configs/rmsnorm/",

239 - help="Path to save rmsnorm benchmark results",

240 - )

241 -

242 - args = parser.parse_args()

243 -

244 - # Run correctness test

245 - calculate_diff(

246 - batch_size=args.batch_size,

247 - seq_len=args.seq_len,

248 - hidden_size=args.hidden_size,

249 - use_residual=args.use_residual,

250 - )

251 -

252 - # Get the benchmark function with proper use_residual setting

253 - benchmark = get_benchmark(args.use_residual)

254 - # Run performance benchmark

255 - benchmark.run(print_data=True, save_path=args.save_path)

benchmarks/kernels/ir/shapes.py

+3 -2

@@ -18,10 +18,11 @@ COMMON_HIDDEN_SIZES = [

18 # Each entry maps an op name to a list of kwarg dicts that will be passed

19 # to that op's registered input generator via op.generate_inputs(**kwargs).

20 SHAPE_CONFIGS: dict[str, list[dict]] = {

21 - "rms_norm": [

21 + op_name: [

22 {"num_tokens": n, "hidden_size": d, "dtype": dtype}

23 for dtype in [torch.float16, torch.bfloat16, torch.float32]

24 for d in COMMON_HIDDEN_SIZES

25 for n in NUM_TOKENS

26 - ],

26 + ]

27 + for op_name in ("rms_norm", "fused_add_rms_norm")

28 }

benchmarks/overheads/benchmark_hashing.py deleted

-64

@@ -1,64 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -import cProfile

5 -import pstats

6 -

7 -from vllm import LLM, SamplingParams

8 -from vllm.utils.argparse_utils import FlexibleArgumentParser

9 -

10 -# A very long prompt, total number of tokens is about 15k.

11 -LONG_PROMPT = ["You are an expert in large language models, aren't you?"] * 1000

12 -LONG_PROMPT = " ".join(LONG_PROMPT)

13 -

14 -

15 -def main(args):

16 - llm = LLM(

17 - model=args.model,

18 - enforce_eager=True,

19 - enable_prefix_caching=True,

20 - tensor_parallel_size=args.tensor_parallel_size,

21 - )

22 -

23 - sampling_params = SamplingParams(temperature=0, max_tokens=args.output_len)

24 - profiler = cProfile.Profile()

25 -

26 - print("------warm up------")

27 - for i in range(3):

28 - output = llm.generate(LONG_PROMPT, sampling_params)

29 - print(output[0].outputs[0].text)

30 -

31 - print("------start generating------")

32 - for i in range(3):

33 - profiler.runctx(

34 - "llm.generate(LONG_PROMPT, sampling_params)", globals(), locals()

35 - )

36 -

37 - # analyze the runtime of hashing function

38 - stats = pstats.Stats(profiler)

39 - stats.sort_stats("cumulative")

40 - total_time = 0

41 - total_calls = 0

42 - for func in stats.stats:

43 - if "hash_of_block" in func[2]:

44 - total_time = stats.stats[func][3]

45 - total_calls = stats.stats[func][0]

46 - percentage = (total_time / stats.total_tt) * 100

47 - print(

48 - f"Hashing took {total_time:.2f} seconds,{percentage:.2f}% of the total runtime."

49 - )

50 -

51 -

52 -if __name__ == "__main__":

53 - parser = FlexibleArgumentParser(

54 - description="Benchmark the performance of hashing function in"

55 - "automatic prefix caching."

56 - )

57 - parser.add_argument("--model", type=str, default="lmsys/longchat-7b-16k")

58 - parser.add_argument("--tensor-parallel-size", "-tp", type=int, default=1)

59 - parser.add_argument("--output-len", type=int, default=10)

60 - parser.add_argument(

61 - "--enable-prefix-caching", action="store_true", help="enable prefix caching"

62 - )

63 - args = parser.parse_args()

64 - main(args)

tools/pre_commit/check_forbidden_imports.py

-2

@@ -60,8 +60,6 @@ CHECK_IMPORTS = {

60 "benchmarks/kernels/benchmark_lora.py",

61 "benchmarks/kernels/benchmark_machete.py",

62 "benchmarks/fused_kernels/layernorm_rms_benchmarks.py",

63 - "benchmarks/cutlass_benchmarks/w8a8_benchmarks.py",

64 - "benchmarks/cutlass_benchmarks/sparse_benchmarks.py",

63 # cloudpickle

64 "vllm/v1/executor/multiproc_executor.py",

65 "vllm/v1/executor/ray_executor.py",