@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 009d9e7590

Fork 0

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 15, 11:00 UTC

Convert `benchmarks` to `ruff format` (#18068)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed May 13, 2025 at 14:43 UTC 009d9e759070422b3f3bbc1cd7d77b5882f905ad

41 files changed +4003 -2961

.buildkite/pyproject.toml

-4

@@ -1,13 +1,9 @@

1 # This local pyproject file is part of the migration from yapf to ruff format.

2 # It uses the same core rules as the main pyproject.toml file, but with the

3 # following differences:

4 -# - isort profile is set to black

4 # - ruff line length is overridden to 88

5 # - deprecated typing ignores (UP006, UP035) have been removed

8 -[tool.isort]

9 -profile = "black"

10 -

7 [tool.ruff]

8 line-length = 88

9 exclude = [

.pre-commit-config.yaml

+1 -3

@@ -17,7 +17,7 @@ repos:

17 - id: ruff

18 args: [--output-format, github, --fix]

19 - id: ruff-format

20 - files: ^(.buildkite).*

20 + files: ^(.buildkite|benchmarks)/.*

21 - repo: https://github.com/codespell-project/codespell

22 rev: v2.4.1

23 hooks:

@@ -28,8 +28,6 @@ repos:

28 rev: 6.0.1

29 hooks:

30 - id: isort

31 - # necessary during the transition from yapf to ruff format

32 - args: [--resolve-all-configs, --config-root, .]

31 - repo: https://github.com/pre-commit/mirrors-clang-format

32 rev: v20.1.3

33 hooks:

benchmarks/backend_request_func.py

+94 -95

@@ -12,8 +12,7 @@ from typing import Optional, Union

12 import aiohttp

13 import huggingface_hub.constants

14 from tqdm.asyncio import tqdm

15 -from transformers import (AutoTokenizer, PreTrainedTokenizer,

16 - PreTrainedTokenizerFast)

15 +from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast

17 # NOTE(simon): do not import vLLM here so the benchmark script

18 # can run without vLLM installed.

@@ -43,8 +42,7 @@ class RequestFuncOutput:

42 latency: float = 0.0

43 output_tokens: int = 0

44 ttft: float = 0.0 # Time to first token

46 - itl: list[float] = field(

47 - default_factory=list) # list of inter-token latencies

45 + itl: list[float] = field(default_factory=list) # list of inter-token latencies

46 tpot: float = 0.0 # avg next-token latencies

47 prompt_len: int = 0

48 error: str = ""

@@ -57,8 +55,9 @@ async def async_request_tgi(

55 api_url = request_func_input.api_url

56 assert api_url.endswith("generate_stream")

60 - async with aiohttp.ClientSession(trust_env=True,

61 - timeout=AIOHTTP_TIMEOUT) as session:

58 + async with aiohttp.ClientSession(

59 + trust_env=True, timeout=AIOHTTP_TIMEOUT

60 + ) as session:

61 params = {

62 "max_new_tokens": request_func_input.output_len,

63 "do_sample": True,

@@ -105,8 +104,7 @@ async def async_request_tgi(

104

105 # Decoding phase

106 else:

108 - output.itl.append(timestamp -

109 - most_recent_timestamp)

107 + output.itl.append(timestamp - most_recent_timestamp)

108

109 most_recent_timestamp = timestamp

110

@@ -133,8 +131,9 @@ async def async_request_trt_llm(

131 api_url = request_func_input.api_url

132 assert api_url.endswith("generate_stream")

133

136 - async with aiohttp.ClientSession(trust_env=True,

137 - timeout=AIOHTTP_TIMEOUT) as session:

134 + async with aiohttp.ClientSession(

135 + trust_env=True, timeout=AIOHTTP_TIMEOUT

136 + ) as session:

137 payload = {

138 "accumulate_tokens": True,

139 "text_input": request_func_input.prompt,

@@ -159,8 +158,7 @@ async def async_request_trt_llm(

158 if not chunk_bytes:

159 continue

160

162 - chunk = chunk_bytes.decode("utf-8").removeprefix(

163 - "data:")

161 + chunk = chunk_bytes.decode("utf-8").removeprefix("data:")

162

163 data = json.loads(chunk)

164 output.generated_text += data["text_output"]

@@ -172,8 +170,7 @@ async def async_request_trt_llm(

170

171 # Decoding phase

172 else:

175 - output.itl.append(timestamp -

176 - most_recent_timestamp)

173 + output.itl.append(timestamp - most_recent_timestamp)

174

175 most_recent_timestamp = timestamp

176

@@ -197,9 +194,9 @@ async def async_request_deepspeed_mii(

194 request_func_input: RequestFuncInput,

195 pbar: Optional[tqdm] = None,

196 ) -> RequestFuncOutput:

200 - async with aiohttp.ClientSession(trust_env=True,

201 - timeout=AIOHTTP_TIMEOUT) as session:

202 -

197 + async with aiohttp.ClientSession(

198 + trust_env=True, timeout=AIOHTTP_TIMEOUT

199 + ) as session:

200 payload = {

201 "model": request_func_input.model,

202 "prompt": request_func_input.prompt,

@@ -217,19 +214,21 @@ async def async_request_deepspeed_mii(

214

215 st = time.perf_counter()

216 try:

220 - async with session.post(url=request_func_input.api_url,

221 - json=payload) as response:

217 + async with session.post(

218 + url=request_func_input.api_url, json=payload

219 + ) as response:

220 if response.status == 200:

221 parsed_resp = await response.json()

222 output.latency = time.perf_counter() - st

223 if "choices" in parsed_resp:

226 - output.generated_text = parsed_resp["choices"][0][

227 - "text"]

224 + output.generated_text = parsed_resp["choices"][0]["text"]

225 elif "text" in parsed_resp:

226 output.generated_text = parsed_resp["text"][0]

227 else:

231 - output.error = ("Unexpected response format: "

232 - "neither 'choices' nor 'text' found")

228 + output.error = (

229 + "Unexpected response format: "

230 + "neither 'choices' nor 'text' found"

231 + )

232 output.success = False

233 output.success = True

234 else:

@@ -250,15 +249,17 @@ async def async_request_openai_completions(

249 pbar: Optional[tqdm] = None,

250 ) -> RequestFuncOutput:

251 api_url = request_func_input.api_url

253 - assert api_url.endswith(

254 - ("completions", "profile")

255 - ), "OpenAI Completions API URL must end with 'completions' or 'profile'."

252 + assert api_url.endswith(("completions", "profile")), (

253 + "OpenAI Completions API URL must end with 'completions' or 'profile'."

254 + )

255

257 - async with aiohttp.ClientSession(trust_env=True,

258 - timeout=AIOHTTP_TIMEOUT) as session:

256 + async with aiohttp.ClientSession(

257 + trust_env=True, timeout=AIOHTTP_TIMEOUT

258 + ) as session:

259 payload = {

260 - "model": request_func_input.model_name \

261 - if request_func_input.model_name else request_func_input.model,

260 + "model": request_func_input.model_name

261 + if request_func_input.model_name

262 + else request_func_input.model,

263 "prompt": request_func_input.prompt,

264 "temperature": 0.0,

265 "repetition_penalty": 1.0,

@@ -273,9 +274,7 @@ async def async_request_openai_completions(

274 payload["ignore_eos"] = request_func_input.ignore_eos

275 if request_func_input.extra_body:

276 payload.update(request_func_input.extra_body)

276 - headers = {

277 - "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"

278 - }

277 + headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}

278

279 output = RequestFuncOutput()

280 output.prompt_len = request_func_input.prompt_len

@@ -284,8 +283,9 @@ async def async_request_openai_completions(

283 st = time.perf_counter()

284 most_recent_timestamp = st

285 try:

287 - async with session.post(url=api_url, json=payload,

288 - headers=headers) as response:

286 + async with session.post(

287 + url=api_url, json=payload, headers=headers

288 + ) as response:

289 if response.status == 200:

290 first_chunk_received = False

291 async for chunk_bytes in response.content:

@@ -293,8 +293,7 @@ async def async_request_openai_completions(

293 if not chunk_bytes:

294 continue

295

296 - chunk = chunk_bytes.decode("utf-8").removeprefix(

297 - "data: ")

296 + chunk = chunk_bytes.decode("utf-8").removeprefix("data: ")

297 if chunk != "[DONE]":

298 data = json.loads(chunk)

299

@@ -314,21 +313,20 @@ async def async_request_openai_completions(

313

314 # Decoding phase

315 else:

317 - output.itl.append(timestamp -

318 - most_recent_timestamp)

316 + output.itl.append(timestamp - most_recent_timestamp)

317

318 most_recent_timestamp = timestamp

319 generated_text += text or ""

320 elif usage := data.get("usage"):

323 - output.output_tokens = usage.get(

324 - "completion_tokens")

321 + output.output_tokens = usage.get("completion_tokens")

322 if first_chunk_received:

323 output.success = True

324 else:

325 output.success = False

326 output.error = (

327 "Never received a valid chunk to calculate TTFT."

331 - "This response will be marked as failed!")

328 + "This response will be marked as failed!"

329 + )

330 output.generated_text = generated_text

331 output.latency = most_recent_timestamp - st

332 else:

@@ -349,23 +347,22 @@ async def async_request_openai_chat_completions(

347 pbar: Optional[tqdm] = None,

348 ) -> RequestFuncOutput:

349 api_url = request_func_input.api_url

352 - assert api_url.endswith(

353 - ("chat/completions", "profile")

354 - ), "OpenAI Chat Completions API URL must end with 'chat/completions'."

350 + assert api_url.endswith(("chat/completions", "profile")), (

351 + "OpenAI Chat Completions API URL must end with 'chat/completions'."

352 + )

353

356 - async with aiohttp.ClientSession(trust_env=True,

357 - timeout=AIOHTTP_TIMEOUT) as session:

354 + async with aiohttp.ClientSession(

355 + trust_env=True, timeout=AIOHTTP_TIMEOUT

356 + ) as session:

357 content = [{"type": "text", "text": request_func_input.prompt}]

358 if request_func_input.multi_modal_content:

359 content.append(request_func_input.multi_modal_content)

360 payload = {

362 - "model": request_func_input.model_name \

363 - if request_func_input.model_name else request_func_input.model,

361 + "model": request_func_input.model_name

362 + if request_func_input.model_name

363 + else request_func_input.model,

364 "messages": [

365 - {

366 - "role": "user",

367 - "content": content

368 - },

365 + {"role": "user", "content": content},

366 ],

367 "temperature": 0.0,

368 "max_completion_tokens": request_func_input.output_len,

@@ -391,16 +388,16 @@ async def async_request_openai_chat_completions(

388 st = time.perf_counter()

389 most_recent_timestamp = st

390 try:

394 - async with session.post(url=api_url, json=payload,

395 - headers=headers) as response:

391 + async with session.post(

392 + url=api_url, json=payload, headers=headers

393 + ) as response:

394 if response.status == 200:

395 async for chunk_bytes in response.content:

396 chunk_bytes = chunk_bytes.strip()

397 if not chunk_bytes:

398 continue

399

402 - chunk = chunk_bytes.decode("utf-8").removeprefix(

403 - "data: ")

400 + chunk = chunk_bytes.decode("utf-8").removeprefix("data: ")

401 if chunk != "[DONE]":

402 timestamp = time.perf_counter()

403 data = json.loads(chunk)

@@ -414,13 +411,11 @@ async def async_request_openai_chat_completions(

411

412 # Decoding phase

413 else:

417 - output.itl.append(timestamp -

418 - most_recent_timestamp)

414 + output.itl.append(timestamp - most_recent_timestamp)

415

416 generated_text += content or ""

417 elif usage := data.get("usage"):

422 - output.output_tokens = usage.get(

423 - "completion_tokens")

418 + output.output_tokens = usage.get("completion_tokens")

419

420 most_recent_timestamp = timestamp

421

@@ -446,25 +441,28 @@ async def async_request_openai_audio(

441 ) -> RequestFuncOutput:

442 # Lazy import without PlaceholderModule to avoid vllm dep.

443 import soundfile

444 +

445 api_url = request_func_input.api_url

450 - assert api_url.endswith(

451 - ("transcriptions", "translations"

452 - )), "OpenAI Chat Completions API URL must end with 'transcriptions' "

446 + assert api_url.endswith(("transcriptions", "translations")), (

447 + "OpenAI Chat Completions API URL must end with 'transcriptions' "

448 + )

449 "or `translations`."

450

455 - async with aiohttp.ClientSession(trust_env=True,

456 - timeout=AIOHTTP_TIMEOUT) as session:

451 + async with aiohttp.ClientSession(

452 + trust_env=True, timeout=AIOHTTP_TIMEOUT

453 + ) as session:

454 content = [{"type": "text", "text": request_func_input.prompt}]

455 payload = {

459 - "model": request_func_input.model_name \

460 - if request_func_input.model_name else request_func_input.model,

456 + "model": request_func_input.model_name

457 + if request_func_input.model_name

458 + else request_func_input.model,

459 "temperature": 0.0,

460 "max_completion_tokens": request_func_input.output_len,

461 "stream": True,

462 "language": "en",

463 # Flattened due to multipart/form-data

464 "stream_include_usage": True,

467 - "stream_continuous_usage_stats": True

465 + "stream_continuous_usage_stats": True,

466 }

467 if request_func_input.extra_body:

468 payload.update(request_func_input.extra_body)

@@ -479,9 +477,9 @@ async def async_request_openai_audio(

477 buffer.seek(0)

478 return buffer

479

482 - with to_bytes(*request_func_input.multi_modal_content['audio']) as f:

480 + with to_bytes(*request_func_input.multi_modal_content["audio"]) as f:

481 form = aiohttp.FormData()

484 - form.add_field('file', f, content_type='audio/wav')

482 + form.add_field("file", f, content_type="audio/wav")

483 for key, value in payload.items():

484 form.add_field(key, str(value))

485

@@ -493,24 +491,22 @@ async def async_request_openai_audio(

491 st = time.perf_counter()

492 most_recent_timestamp = st

493 try:

496 - async with session.post(url=api_url,

497 - data=form,

498 - headers=headers) as response:

494 + async with session.post(

495 + url=api_url, data=form, headers=headers

496 + ) as response:

497 if response.status == 200:

498 async for chunk_bytes in response.content:

499 chunk_bytes = chunk_bytes.strip()

500 if not chunk_bytes:

501 continue

502

505 - chunk = chunk_bytes.decode("utf-8").removeprefix(

506 - "data: ")

503 + chunk = chunk_bytes.decode("utf-8").removeprefix("data: ")

504 if chunk != "[DONE]":

505 timestamp = time.perf_counter()

506 data = json.loads(chunk)

507

508 if choices := data.get("choices"):

512 - content = choices[0]["delta"].get(

513 - "content")

509 + content = choices[0]["delta"].get("content")

510 # First token

511 if ttft == 0.0:

512 ttft = timestamp - st

@@ -519,12 +515,14 @@ async def async_request_openai_audio(

515 # Decoding phase

516 else:

517 output.itl.append(

522 - timestamp - most_recent_timestamp)

518 + timestamp - most_recent_timestamp

519 + )

520

521 generated_text += content or ""

522 elif usage := data.get("usage"):

523 output.output_tokens = usage.get(

527 - "completion_tokens")

524 + "completion_tokens"

525 + )

526

527 most_recent_timestamp = timestamp

528

@@ -545,7 +543,7 @@ async def async_request_openai_audio(

543

544

545 def get_model(pretrained_model_name_or_path: str) -> str:

548 - if os.getenv('VLLM_USE_MODELSCOPE', 'False').lower() == 'true':

546 + if os.getenv("VLLM_USE_MODELSCOPE", "False").lower() == "true":

547 from modelscope import snapshot_download

548

549 from vllm.model_executor.model_loader.weight_utils import get_lock

@@ -556,7 +554,8 @@ def get_model(pretrained_model_name_or_path: str) -> str:

554 model_path = snapshot_download(

555 model_id=pretrained_model_name_or_path,

556 local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,

559 - ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"])

557 + ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"],

558 + )

559

560 return model_path

561 return pretrained_model_name_or_path

@@ -569,23 +568,23 @@ def get_tokenizer(

568 **kwargs,

569 ) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]:

570 if pretrained_model_name_or_path is not None and not os.path.exists(

572 - pretrained_model_name_or_path):

573 - pretrained_model_name_or_path = get_model(

574 - pretrained_model_name_or_path)

571 + pretrained_model_name_or_path

572 + ):

573 + pretrained_model_name_or_path = get_model(pretrained_model_name_or_path)

574 if tokenizer_mode == "slow":

575 if kwargs.get("use_fast", False):

577 - raise ValueError(

578 - "Cannot use the fast tokenizer in slow tokenizer mode.")

576 + raise ValueError("Cannot use the fast tokenizer in slow tokenizer mode.")

577 kwargs["use_fast"] = False

578 if tokenizer_mode == "mistral":

579 try:

580 from vllm.transformers_utils.tokenizer import MistralTokenizer

581 except ImportError as e:

584 - raise ImportError("MistralTokenizer requires vllm package.\n"

585 - "Please install it with `pip install vllm` "

586 - "to use mistral tokenizer mode.") from e

587 - return MistralTokenizer.from_pretrained(

588 - str(pretrained_model_name_or_path))

582 + raise ImportError(

583 + "MistralTokenizer requires vllm package.\n"

584 + "Please install it with `pip install vllm` "

585 + "to use mistral tokenizer mode."

586 + ) from e

587 + return MistralTokenizer.from_pretrained(str(pretrained_model_name_or_path))

588 else:

589 return AutoTokenizer.from_pretrained(

590 pretrained_model_name_or_path,

@@ -608,7 +607,7 @@ ASYNC_REQUEST_FUNCS = {

607 }

608

609 OPENAI_COMPATIBLE_BACKENDS = [

611 - k for k, v in ASYNC_REQUEST_FUNCS.items()

612 - if v in (async_request_openai_completions,

613 - async_request_openai_chat_completions)

610 + k

611 + for k, v in ASYNC_REQUEST_FUNCS.items()

612 + if v in (async_request_openai_completions, async_request_openai_chat_completions)

613 ]

benchmarks/benchmark_dataset.py

+186 -169

@@ -82,14 +82,12 @@ class BenchmarkDataset(ABC):

82 self.dataset_path = dataset_path

83 # Set the random seed, ensuring that a None value is replaced with the

84 # default seed.

85 - self.random_seed = (random_seed

86 - if random_seed is not None else self.DEFAULT_SEED)

85 + self.random_seed = random_seed if random_seed is not None else self.DEFAULT_SEED

86 self.data = None

88 def apply_multimodal_chat_transformation(

90 - self,

91 - prompt: str,

92 - mm_content: Optional[MultiModalDataDict] = None) -> list[dict]:

89 + self, prompt: str, mm_content: Optional[MultiModalDataDict] = None

90 + ) -> list[dict]:

91 """

92 Transform a prompt and optional multimodal content into a chat format.

93 This method is used for chat models that expect a specific conversation

@@ -111,8 +109,7 @@ class BenchmarkDataset(ABC):

109 NotImplementedError: If a subclass does not implement this method.

110 """

111 # TODO (jenniferzhao): add support for downloading data

114 - raise NotImplementedError(

115 - "load_data must be implemented in subclasses.")

112 + raise NotImplementedError("load_data must be implemented in subclasses.")

113

114 def get_random_lora_request(

115 self,

@@ -158,8 +155,9 @@ class BenchmarkDataset(ABC):

155 return lora_request, lora_tokenizer_cache[lora_id] or tokenizer

156

157 @abstractmethod

161 - def sample(self, tokenizer: PreTrainedTokenizerBase,

162 - num_requests: int) -> list[SampleRequest]:

158 + def sample(

159 + self, tokenizer: PreTrainedTokenizerBase, num_requests: int

160 + ) -> list[SampleRequest]:

161 """

162 Abstract method to generate sample requests from the dataset.

163

@@ -177,8 +175,9 @@ class BenchmarkDataset(ABC):

175 """

176 raise NotImplementedError("sample must be implemented in subclasses.")

177

180 - def maybe_oversample_requests(self, requests: list[SampleRequest],

181 - num_requests: int) -> None:

178 + def maybe_oversample_requests(

179 + self, requests: list[SampleRequest], num_requests: int

180 + ) -> None:

181 """

182 Oversamples the list of requests if its size is less than the desired

183 number.

@@ -189,11 +188,9 @@ class BenchmarkDataset(ABC):

188 """

189 if len(requests) < num_requests:

190 random.seed(self.random_seed)

192 - additional = random.choices(requests,

193 - k=num_requests - len(requests))

191 + additional = random.choices(requests, k=num_requests - len(requests))

192 requests.extend(additional)

195 - logger.info("Oversampled requests to reach %d total samples.",

196 - num_requests)

193 + logger.info("Oversampled requests to reach %d total samples.", num_requests)

194

195

196 # -----------------------------------------------------------------------------

@@ -218,14 +215,14 @@ def is_valid_sequence(

215 """

216 # Check for invalid conditions

217 prompt_too_short = prompt_len < min_len

221 - output_too_short = (not skip_min_output_len_check) and (output_len

222 - < min_len)

218 + output_too_short = (not skip_min_output_len_check) and (output_len < min_len)

219 prompt_too_long = prompt_len > max_prompt_len

220 combined_too_long = (prompt_len + output_len) > max_total_len

221

222 # Return True if none of the invalid conditions are met

227 - return not (prompt_too_short or output_too_short or prompt_too_long

228 - or combined_too_long)

223 + return not (

224 + prompt_too_short or output_too_short or prompt_too_long or combined_too_long

225 + )

226

227

228 @cache

@@ -257,28 +254,28 @@ def process_image(image: Any) -> Mapping[str, Any]:

254 Raises:

255 ValueError: If the input is not a supported type.

256 """

260 - if isinstance(image, dict) and 'bytes' in image:

261 - image = Image.open(BytesIO(image['bytes']))

257 + if isinstance(image, dict) and "bytes" in image:

258 + image = Image.open(BytesIO(image["bytes"]))

259 if isinstance(image, Image.Image):

260 image = image.convert("RGB")

261 with io.BytesIO() as image_data:

262 image.save(image_data, format="JPEG")

266 - image_base64 = base64.b64encode(

267 - image_data.getvalue()).decode("utf-8")

263 + image_base64 = base64.b64encode(image_data.getvalue()).decode("utf-8")

264 return {

265 "type": "image_url",

270 - "image_url": {

271 - "url": f"data:image/jpeg;base64,{image_base64}"

272 - },

266 + "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"},

267 }

268

269 if isinstance(image, str):

276 - image_url = (image if image.startswith(

277 - ("http://", "file://")) else f"file://{image}")

270 + image_url = (

271 + image if image.startswith(("http://", "file://")) else f"file://{image}"

272 + )

273 return {"type": "image_url", "image_url": {"url": image_url}}

274

280 - raise ValueError(f"Invalid image input {image}. Must be a PIL.Image.Image"

281 - " or str or dictionary with raw image bytes.")

275 + raise ValueError(

276 + f"Invalid image input {image}. Must be a PIL.Image.Image"

277 + " or str or dictionary with raw image bytes."

278 + )

279

280

281 # -----------------------------------------------------------------------------

@@ -318,8 +315,11 @@ class RandomDataset(BenchmarkDataset):

315 num_special_tokens = tokenizer.num_special_tokens_to_add()

316 real_input_len = input_len - num_special_tokens

317

321 - prefix_token_ids = (np.random.randint(

322 - 0, vocab_size, size=prefix_len).tolist() if prefix_len > 0 else [])

318 + prefix_token_ids = (

319 + np.random.randint(0, vocab_size, size=prefix_len).tolist()

320 + if prefix_len > 0

321 + else []

322 + )

323

324 # New sampling logic: [X * (1 - b), X * (1 + b)]

325 input_low = int(real_input_len * (1 - range_ratio))

@@ -329,21 +329,17 @@ class RandomDataset(BenchmarkDataset):

329

330 # Add logging for debugging

331 logger.info("Sampling input_len from [%s, %s]", input_low, input_high)

332 - logger.info("Sampling output_len from [%s, %s]", output_low,

333 - output_high)

334 -

335 - input_lens = np.random.randint(input_low,

336 - input_high + 1,

337 - size=num_requests)

338 - output_lens = np.random.randint(output_low,

339 - output_high + 1,

340 - size=num_requests)

332 + logger.info("Sampling output_len from [%s, %s]", output_low, output_high)

333 +

334 + input_lens = np.random.randint(input_low, input_high + 1, size=num_requests)

335 + output_lens = np.random.randint(output_low, output_high + 1, size=num_requests)

336 offsets = np.random.randint(0, vocab_size, size=num_requests)

337

338 requests = []

339 for i in range(num_requests):

345 - inner_seq = ((offsets[i] + i + np.arange(input_lens[i])) %

346 - vocab_size).tolist()

340 + inner_seq = (

341 + (offsets[i] + i + np.arange(input_lens[i])) % vocab_size

342 + ).tolist()

343 token_sequence = prefix_token_ids + inner_seq

344 prompt = tokenizer.decode(token_sequence)

345 # After decoding the prompt we have to encode and decode it again.

@@ -354,8 +350,9 @@ class RandomDataset(BenchmarkDataset):

350 # [1650, 939, 486] -> ['Ġcall', 'sh', 'ere']

351 # To avoid uncontrolled change of the prompt length,

352 # the encoded sequence is truncated before being decode again.

357 - re_encoded_sequence = tokenizer.encode(

358 - prompt, add_special_tokens=False)[:input_lens[i]]

353 + re_encoded_sequence = tokenizer.encode(prompt, add_special_tokens=False)[

354 + : input_lens[i]

355 + ]

356 prompt = tokenizer.decode(re_encoded_sequence)

357 total_input_len = prefix_len + int(input_lens[i])

358 requests.append(

@@ -363,7 +360,8 @@ class RandomDataset(BenchmarkDataset):

360 prompt=prompt,

361 prompt_len=total_input_len,

362 expected_output_len=int(output_lens[i]),

366 - ))

363 + )

364 + )

365 return requests

366

367

@@ -390,7 +388,8 @@ class ShareGPTDataset(BenchmarkDataset):

388 self.data = json.load(f)

389 # Filter entries with at least two conversation turns.

390 self.data = [

393 - entry for entry in self.data

391 + entry

392 + for entry in self.data

393 if "conversations" in entry and len(entry["conversations"]) >= 2

394 ]

395 random.seed(self.random_seed)

@@ -416,27 +415,28 @@ class ShareGPTDataset(BenchmarkDataset):

415 )

416

417 lora_request, tokenizer = self.get_random_lora_request(

419 - tokenizer=tokenizer, max_loras=max_loras, lora_path=lora_path)

418 + tokenizer=tokenizer, max_loras=max_loras, lora_path=lora_path

419 + )

420 prompt_ids = tokenizer(prompt).input_ids

421 completion_ids = tokenizer(completion).input_ids

422 prompt_len = len(prompt_ids)

423 - new_output_len = (len(completion_ids)

424 - if output_len is None else output_len)

425 - if not is_valid_sequence(prompt_len,

426 - new_output_len,

427 - skip_min_output_len_check=output_len

428 - is not None):

423 + new_output_len = len(completion_ids) if output_len is None else output_len

424 + if not is_valid_sequence(

425 + prompt_len,

426 + new_output_len,

427 + skip_min_output_len_check=output_len is not None,

428 + ):

429 continue

430 if enable_multimodal_chat:

431 - prompt = self.apply_multimodal_chat_transformation(

432 - prompt, None)

431 + prompt = self.apply_multimodal_chat_transformation(prompt, None)

432 samples.append(

433 SampleRequest(

434 prompt=prompt,

435 prompt_len=prompt_len,

436 expected_output_len=new_output_len,

437 lora_request=lora_request,

439 - ))

438 + )

439 + )

440 self.maybe_oversample_requests(samples, num_requests)

441 return samples

442

@@ -482,20 +482,20 @@ class SonnetDataset(BenchmarkDataset):

482 ) -> list:

483 # Calculate average token length for a poem line.

484 tokenized_lines = [tokenizer(line).input_ids for line in self.data]

485 - avg_len = sum(len(tokens)

486 - for tokens in tokenized_lines) / len(tokenized_lines)

485 + avg_len = sum(len(tokens) for tokens in tokenized_lines) / len(tokenized_lines)

486

487 # Build the base prompt.

488 base_prompt = "Pick as many lines as you can from these poem lines:\n"

489 base_msg = [{"role": "user", "content": base_prompt}]

491 - base_fmt = tokenizer.apply_chat_template(base_msg,

492 - add_generation_prompt=True,

493 - tokenize=False)

490 + base_fmt = tokenizer.apply_chat_template(

491 + base_msg, add_generation_prompt=True, tokenize=False

492 + )

493 base_offset = len(tokenizer(base_fmt).input_ids)

494 if input_len <= base_offset:

495 raise ValueError(

496 f"'input_len' must be higher than the base prompt length "

498 - f"({base_offset}).")

497 + f"({base_offset})."

498 + )

499

500 # Determine how many poem lines to use.

501 num_input_lines = round((input_len - base_offset) / avg_len)

@@ -504,21 +504,23 @@ class SonnetDataset(BenchmarkDataset):

504

505 samples = []

506 while len(samples) < num_requests:

507 - extra_lines = random.choices(self.data,

508 - k=num_input_lines - num_prefix_lines)

507 + extra_lines = random.choices(

508 + self.data, k=num_input_lines - num_prefix_lines

509 + )

510 prompt = f"{base_prompt}{''.join(prefix_lines + extra_lines)}"

511 msg = [{"role": "user", "content": prompt}]

512 prompt_formatted = tokenizer.apply_chat_template(

512 - msg, add_generation_prompt=True, tokenize=False)

513 + msg, add_generation_prompt=True, tokenize=False

514 + )

515 prompt_len = len(tokenizer(prompt_formatted).input_ids)

516 if prompt_len <= input_len:

517 samples.append(

518 SampleRequest(

517 - prompt=prompt_formatted

518 - if return_prompt_formatted else prompt,

519 + prompt=prompt_formatted if return_prompt_formatted else prompt,

520 prompt_len=prompt_len,

521 expected_output_len=output_len,

521 - ))

522 + )

523 + )

524 return samples

525

526

@@ -538,7 +540,9 @@ class BurstGPTDataset(BenchmarkDataset):

540 super().__init__(**kwargs)

541 self.load_data()

542

541 - def load_data(self, ):

543 + def load_data(

544 + self,

545 + ):

546 if self.dataset_path is None:

547 raise ValueError("dataset_path must be provided for loading data.")

548

@@ -552,8 +556,7 @@ class BurstGPTDataset(BenchmarkDataset):

556

557 def _sample_loaded_data(self, num_requests: int) -> list:

558 if num_requests <= len(self.data):

555 - data = self.data.sample(n=num_requests,

556 - random_state=self.random_seed)

559 + data = self.data.sample(n=num_requests, random_state=self.random_seed)

560 else:

561 data = self.data.sample(

562 n=num_requests,

@@ -577,7 +580,8 @@ class BurstGPTDataset(BenchmarkDataset):

580 input_len = int(data[i][2])

581 output_len = int(data[i][3])

582 lora_req, tokenizer = self.get_random_lora_request(

580 - tokenizer=tokenizer, max_loras=max_loras, lora_path=lora_path)

583 + tokenizer=tokenizer, max_loras=max_loras, lora_path=lora_path

584 + )

585 vocab_size = tokenizer.vocab_size

586 # Generate a synthetic prompt: a list of token IDs computed as (i +

587 # j) modulo vocab_size.

@@ -589,7 +593,8 @@ class BurstGPTDataset(BenchmarkDataset):

593 prompt_len=input_len,

594 expected_output_len=output_len,

595 lora_request=lora_req,

592 - ))

596 + )

597 + )

598 return samples

599

600

@@ -632,20 +637,23 @@ class HuggingFaceDataset(BenchmarkDataset):

637

638 class ConversationDataset(HuggingFaceDataset):

639 """Dataset for conversation data with multimodal support."""

640 +

641 SUPPORTED_DATASET_PATHS = {

636 - 'lmms-lab/LLaVA-OneVision-Data', 'Aeala/ShareGPT_Vicuna_unfiltered'

642 + "lmms-lab/LLaVA-OneVision-Data",

643 + "Aeala/ShareGPT_Vicuna_unfiltered",

644 }

645 IS_MULTIMODAL = True

646

640 - def sample(self,

641 - tokenizer: PreTrainedTokenizerBase,

642 - num_requests: int,

643 - output_len: Optional[int] = None,

644 - enable_multimodal_chat: bool = False,

645 - **kwargs) -> list:

647 + def sample(

648 + self,

649 + tokenizer: PreTrainedTokenizerBase,

650 + num_requests: int,

651 + output_len: Optional[int] = None,

652 + enable_multimodal_chat: bool = False,

653 + **kwargs,

654 + ) -> list:

655 # Filter examples with at least 2 conversations

647 - filtered_data = self.data.filter(

648 - lambda x: len(x["conversations"]) >= 2)

656 + filtered_data = self.data.filter(lambda x: len(x["conversations"]) >= 2)

657 sampled_requests = []

658 dynamic_output = output_len is None

659

@@ -661,24 +669,22 @@ class ConversationDataset(HuggingFaceDataset):

669 completion_len = len(completion_ids)

670 output_len = completion_len if dynamic_output else output_len

671 assert isinstance(output_len, int) and output_len > 0

664 - if dynamic_output and not is_valid_sequence(

665 - prompt_len, completion_len):

672 + if dynamic_output and not is_valid_sequence(prompt_len, completion_len):

673 continue

667 - mm_content = process_image(

668 - item["image"]) if "image" in item else None

674 + mm_content = process_image(item["image"]) if "image" in item else None

675 if enable_multimodal_chat:

676 # Note: when chat is enabled the request prompt_len is no longer

677 # accurate and we will be using request output to count the

678 # actual prompt len and output len

673 - prompt = self.apply_multimodal_chat_transformation(

674 - prompt, mm_content)

679 + prompt = self.apply_multimodal_chat_transformation(prompt, mm_content)

680 sampled_requests.append(

681 SampleRequest(

682 prompt=prompt,

683 prompt_len=prompt_len,

684 expected_output_len=output_len,

685 multi_modal_data=mm_content,

681 - ))

686 + )

687 + )

688 self.maybe_oversample_requests(sampled_requests, num_requests)

689 return sampled_requests

690

@@ -695,10 +701,8 @@ class VisionArenaDataset(HuggingFaceDataset):

701

702 DEFAULT_OUTPUT_LEN = 128

703 SUPPORTED_DATASET_PATHS = {

698 - "lmarena-ai/VisionArena-Chat":

699 - lambda x: x["conversation"][0][0]["content"],

700 - "lmarena-ai/vision-arena-bench-v0.1":

701 - lambda x: x["turns"][0][0]["content"]

704 + "lmarena-ai/VisionArena-Chat": lambda x: x["conversation"][0][0]["content"],

705 + "lmarena-ai/vision-arena-bench-v0.1": lambda x: x["turns"][0][0]["content"],

706 }

707 IS_MULTIMODAL = True

708

@@ -710,16 +714,14 @@ class VisionArenaDataset(HuggingFaceDataset):

714 enable_multimodal_chat: bool = False,

715 **kwargs,

716 ) -> list:

713 - output_len = (output_len

714 - if output_len is not None else self.DEFAULT_OUTPUT_LEN)

717 + output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN

718 sampled_requests = []

719 for item in self.data:

720 if len(sampled_requests) >= num_requests:

721 break

722 parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.dataset_path)

723 if parser_fn is None:

721 - raise ValueError(

722 - f"Unsupported dataset path: {self.dataset_path}")

724 + raise ValueError(f"Unsupported dataset path: {self.dataset_path}")

725 prompt = parser_fn(item)

726 mm_content = process_image(item["images"][0])

727 prompt_len = len(tokenizer(prompt).input_ids)

@@ -727,15 +729,15 @@ class VisionArenaDataset(HuggingFaceDataset):

729 # Note: when chat is enabled the request prompt_len is no longer

730 # accurate and we will be using request output to count the

731 # actual prompt len

730 - prompt = self.apply_multimodal_chat_transformation(

731 - prompt, mm_content)

732 + prompt = self.apply_multimodal_chat_transformation(prompt, mm_content)

733 sampled_requests.append(

734 SampleRequest(

735 prompt=prompt,

736 prompt_len=prompt_len,

737 expected_output_len=output_len,

738 multi_modal_data=mm_content,

738 - ))

739 + )

740 + )

741 self.maybe_oversample_requests(sampled_requests, num_requests)

742 return sampled_requests

743

@@ -760,14 +762,15 @@ class InstructCoderDataset(HuggingFaceDataset):

762 "likaixin/InstructCoder",

763 }

764

763 - def sample(self,

764 - tokenizer: PreTrainedTokenizerBase,

765 - num_requests: int,

766 - output_len: Optional[int] = None,

767 - enable_multimodal_chat: bool = False,

768 - **kwargs) -> list:

769 - output_len = (output_len

770 - if output_len is not None else self.DEFAULT_OUTPUT_LEN)

765 + def sample(

766 + self,

767 + tokenizer: PreTrainedTokenizerBase,

768 + num_requests: int,

769 + output_len: Optional[int] = None,

770 + enable_multimodal_chat: bool = False,

771 + **kwargs,

772 + ) -> list:

773 + output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN

774 sampled_requests = []

775 for item in self.data:

776 if len(sampled_requests) >= num_requests:

@@ -779,7 +782,8 @@ class InstructCoderDataset(HuggingFaceDataset):

782 prompt=prompt,

783 prompt_len=prompt_len,

784 expected_output_len=output_len,

782 - ))

785 + )

786 + )

787 self.maybe_oversample_requests(sampled_requests, num_requests)

788 return sampled_requests

789

@@ -794,38 +798,38 @@ class MTBenchDataset(HuggingFaceDataset):

798 MT-Bench Dataset.

799 https://huggingface.co/datasets/philschmid/mt-bench

800

797 - We create a single turn dataset for MT-Bench.

801 + We create a single turn dataset for MT-Bench.

802 This is similar to Spec decoding benchmark setup in vLLM

803 https://github.com/vllm-project/vllm/blob/9d98ab5ec/examples/offline_inference/eagle.py#L14-L18

800 - """ # noqa: E501

804 + """ # noqa: E501

805

806 DEFAULT_OUTPUT_LEN = 256 # avg len used in SD bench in vLLM

807 SUPPORTED_DATASET_PATHS = {

808 "philschmid/mt-bench",

809 }

810

807 - def sample(self,

808 - tokenizer: PreTrainedTokenizerBase,

809 - num_requests: int,

810 - output_len: Optional[int] = None,

811 - enable_multimodal_chat: bool = False,

812 - **kwargs) -> list:

813 - output_len = (output_len

814 - if output_len is not None else self.DEFAULT_OUTPUT_LEN)

811 + def sample(

812 + self,

813 + tokenizer: PreTrainedTokenizerBase,

814 + num_requests: int,

815 + output_len: Optional[int] = None,

816 + enable_multimodal_chat: bool = False,

817 + **kwargs,

818 + ) -> list:

819 + output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN

820 sampled_requests = []

821

822 for item in self.data:

823 if len(sampled_requests) >= num_requests:

824 break

820 - prompt = item['turns'][0]

825 + prompt = item["turns"][0]

826

827 # apply template

823 - prompt = tokenizer.apply_chat_template([{

824 - "role": "user",

825 - "content": prompt

826 - }],

827 - add_generation_prompt=True,

828 - tokenize=False)

828 + prompt = tokenizer.apply_chat_template(

829 + [{"role": "user", "content": prompt}],

830 + add_generation_prompt=True,

831 + tokenize=False,

832 + )

833

834 prompt_len = len(tokenizer(prompt).input_ids)

835 sampled_requests.append(

@@ -833,7 +837,8 @@ class MTBenchDataset(HuggingFaceDataset):

837 prompt=prompt,

838 prompt_len=prompt_len,

839 expected_output_len=output_len,

836 - ))

840 + )

841 + )

842 self.maybe_oversample_requests(sampled_requests, num_requests)

843 return sampled_requests

844

@@ -847,23 +852,27 @@ class AIMODataset(HuggingFaceDataset):

852 """

853 Dataset class for processing a AIMO dataset with reasoning questions.

854 """

855 +

856 SUPPORTED_DATASET_PATHS = {

851 - "AI-MO/aimo-validation-aime", "AI-MO/NuminaMath-1.5",

852 - "AI-MO/NuminaMath-CoT"

857 + "AI-MO/aimo-validation-aime",

858 + "AI-MO/NuminaMath-1.5",

859 + "AI-MO/NuminaMath-CoT",

860 }

861

855 - def sample(self,

856 - tokenizer: PreTrainedTokenizerBase,

857 - num_requests: int,

858 - output_len: Optional[int] = None,

859 - **kwargs) -> list:

862 + def sample(

863 + self,

864 + tokenizer: PreTrainedTokenizerBase,

865 + num_requests: int,

866 + output_len: Optional[int] = None,

867 + **kwargs,

868 + ) -> list:

869 sampled_requests = []

870 dynamic_output = output_len is None

871

872 for item in self.data:

873 if len(sampled_requests) >= num_requests:

874 break

866 - prompt, completion = item['problem'], item["solution"]

875 + prompt, completion = item["problem"], item["solution"]

876

877 prompt_ids = tokenizer(prompt).input_ids

878 completion_ids = tokenizer(completion).input_ids

@@ -871,10 +880,9 @@ class AIMODataset(HuggingFaceDataset):

880 completion_len = len(completion_ids)

881 output_len = completion_len if dynamic_output else output_len

882 assert isinstance(output_len, int) and output_len > 0

874 - if dynamic_output and not is_valid_sequence(prompt_len,

875 - completion_len,

876 - max_prompt_len=2048,

877 - max_total_len=32000):

883 + if dynamic_output and not is_valid_sequence(

884 + prompt_len, completion_len, max_prompt_len=2048, max_total_len=32000

885 + ):

886 continue

887 sampled_requests.append(

888 SampleRequest(

@@ -882,7 +890,8 @@ class AIMODataset(HuggingFaceDataset):

890 prompt_len=prompt_len,

891 expected_output_len=output_len,

892 multi_modal_data=None,

885 - ))

893 + )

894 + )

895 self.maybe_oversample_requests(sampled_requests, num_requests)

896 return sampled_requests

897

@@ -905,25 +914,25 @@ You are a code completion assistant and your task is to analyze user edits and t

914

915 ### Response:

916

908 -""" # noqa: E501

917 +""" # noqa: E501

918

919

920 def _format_zeta_prompt(

912 - sample: dict,

913 - original_start_marker: str = "<|editable_region_start|>") -> dict:

921 + sample: dict, original_start_marker: str = "<|editable_region_start|>"

922 +) -> dict:

923 """Format the zeta prompt for the Next Edit Prediction (NEP) dataset.

915 -

916 - This function formats examples from the NEP dataset

917 - into prompts and expected outputs. It could be

924 +

925 + This function formats examples from the NEP dataset

926 + into prompts and expected outputs. It could be

927 further extended to support more NEP datasets.

919 -

928 +

929 Args:

921 - sample: The dataset sample containing events,

930 + sample: The dataset sample containing events,

931 inputs, and outputs.

923 - original_start_marker: The marker indicating the

924 - start of the editable region. Defaults to

932 + original_start_marker: The marker indicating the

933 + start of the editable region. Defaults to

934 "<|editable_region_start|>".

926 -

935 +

936 Returns:

937 A dictionary with the formatted prompts and expected outputs.

938 """

@@ -953,10 +962,8 @@ class NextEditPredictionDataset(HuggingFaceDataset):

962 "zed-industries/zeta": _format_zeta_prompt,

963 }

964

956 - def sample(self, tokenizer: PreTrainedTokenizerBase, num_requests: int,

957 - **kwargs):

958 - formatting_prompt_func = self.MAPPING_PROMPT_FUNCS.get(

959 - self.dataset_path)

965 + def sample(self, tokenizer: PreTrainedTokenizerBase, num_requests: int, **kwargs):

966 + formatting_prompt_func = self.MAPPING_PROMPT_FUNCS.get(self.dataset_path)

967 if formatting_prompt_func is None:

968 raise ValueError(f"Unsupported dataset path: {self.dataset_path}")

969 samples = []

@@ -967,8 +974,10 @@ class NextEditPredictionDataset(HuggingFaceDataset):

974 prompt=sample["prompt"],

975 prompt_len=len(tokenizer(sample["prompt"]).input_ids),

976 expected_output_len=len(

970 - tokenizer(sample["expected_output"]).input_ids),

971 - ))

977 + tokenizer(sample["expected_output"]).input_ids

978 + ),

979 + )

980 + )

981 if len(samples) >= num_requests:

982 break

983 self.maybe_oversample_requests(samples, num_requests)

@@ -997,18 +1006,22 @@ class ASRDataset(HuggingFaceDataset):

1007 +----------------+----------------------------------------+--------------------------+-----------------------------+

1008

1000 - """ # noqa: E501

1009 + """ # noqa: E501

1010 +

1011 SUPPORTED_DATASET_PATHS = {

1002 - "openslr/librispeech_asr", "facebook/voxpopuli", "LIUM/tedlium",

1003 - "edinburghcstr/ami", "speechcolab/gigaspeech", "kensho/spgispeech"

1012 + "openslr/librispeech_asr",

1013 + "facebook/voxpopuli",

1014 + "LIUM/tedlium",

1015 + "edinburghcstr/ami",

1016 + "speechcolab/gigaspeech",

1017 + "kensho/spgispeech",

1018 }

1019

1020 DEFAULT_OUTPUT_LEN = 128

1021 IS_MULTIMODAL = True

1022

1023 # TODO Whisper-specific. Abstract interface when more models are supported.

1010 - TRANSCRIPTION_PREAMBLE = "<|startoftranscript|><|en|><|transcribe|>"\

1011 - "<|notimestamps|>"

1024 + TRANSCRIPTION_PREAMBLE = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"

1025 skip_long_audios: bool = True

1026

1027 def sample(

@@ -1019,8 +1032,8 @@ class ASRDataset(HuggingFaceDataset):

1032 **kwargs,

1033 ) -> list:

1034 import librosa

1022 - output_len = (output_len

1023 - if output_len is not None else self.DEFAULT_OUTPUT_LEN)

1035 +

1036 + output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN

1037 prompt = ASRDataset.TRANSCRIPTION_PREAMBLE

1038 prompt_len = len(tokenizer(prompt).input_ids)

1039 sampled_requests = []

@@ -1043,10 +1056,14 @@ class ASRDataset(HuggingFaceDataset):

1056 prompt_len=prompt_len,

1057 expected_output_len=output_len,

1058 multi_modal_data=mm_content,

1046 - ))

1059 + )

1060 + )

1061 if skipped:

1048 - logger.warning("%d samples discarded from dataset due to" \

1049 - " their length being greater than" \

1050 - " what Whisper supports.", skipped)

1062 + logger.warning(

1063 + "%d samples discarded from dataset due to"

1064 + " their length being greater than"

1065 + " what Whisper supports.",

1066 + skipped,

1067 + )

1068 self.maybe_oversample_requests(sampled_requests, num_requests)

1069 return sampled_requests

benchmarks/benchmark_latency.py

+41 -34

@@ -11,9 +11,9 @@ from typing import Any, Optional

12 import numpy as np

13 import torch

14 -from benchmark_utils import convert_to_pytorch_benchmark_format, write_to_json

14 from tqdm import tqdm

16 +from benchmark_utils import convert_to_pytorch_benchmark_format, write_to_json

17 from vllm import LLM, SamplingParams

18 from vllm.engine.arg_utils import EngineArgs

19 from vllm.inputs import PromptType

@@ -21,13 +21,14 @@ from vllm.sampling_params import BeamSearchParams

21 from vllm.utils import FlexibleArgumentParser

24 -def save_to_pytorch_benchmark_format(args: argparse.Namespace,

25 - results: dict[str, Any]) -> None:

24 +def save_to_pytorch_benchmark_format(

25 + args: argparse.Namespace, results: dict[str, Any]

26 +) -> None:

27 pt_records = convert_to_pytorch_benchmark_format(

28 args=args,

29 metrics={"latency": results["latencies"]},

29 - extra_info={k: results[k]

30 - for k in ["avg_latency", "percentiles"]})

30 + extra_info={k: results[k] for k in ["avg_latency", "percentiles"]},

31 + )

32 if pt_records:

33 pt_file = f"{os.path.splitext(args.output_json)[0]}.pytorch.json"

34 write_to_json(pt_file, pt_records)

@@ -42,9 +43,11 @@ def main(args: argparse.Namespace):

43 # the engine will automatically process the request in multiple batches.

44 llm = LLM(**dataclasses.asdict(engine_args))

45 assert llm.llm_engine.model_config.max_model_len >= (

45 - args.input_len +

46 - args.output_len), ("Please ensure that max_model_len is greater than"

47 - " the sum of input_len and output_len.")

46 + args.input_len + args.output_len

47 + ), (

48 + "Please ensure that max_model_len is greater than"

49 + " the sum of input_len and output_len."

50 + )

52 sampling_params = SamplingParams(

53 n=args.n,

@@ -55,18 +58,16 @@ def main(args: argparse.Namespace):

58 detokenize=not args.disable_detokenize,

59 )

60 print(sampling_params)

58 - dummy_prompt_token_ids = np.random.randint(10000,

59 - size=(args.batch_size,

60 - args.input_len))

61 - dummy_prompts: list[PromptType] = [{

62 - "prompt_token_ids": batch

63 - } for batch in dummy_prompt_token_ids.tolist()]

61 + dummy_prompt_token_ids = np.random.randint(

62 + 10000, size=(args.batch_size, args.input_len)

63 + )

64 + dummy_prompts: list[PromptType] = [

65 + {"prompt_token_ids": batch} for batch in dummy_prompt_token_ids.tolist()

66 + ]

68 def llm_generate():

69 if not args.use_beam_search:

67 - llm.generate(dummy_prompts,

68 - sampling_params=sampling_params,

69 - use_tqdm=False)

70 + llm.generate(dummy_prompts, sampling_params=sampling_params, use_tqdm=False)

71 else:

72 llm.beam_search(

73 dummy_prompts,

@@ -80,12 +81,13 @@ def main(args: argparse.Namespace):

81 def run_to_completion(profile_dir: Optional[str] = None):

82 if profile_dir:

83 with torch.profiler.profile(

83 - activities=[

84 - torch.profiler.ProfilerActivity.CPU,

85 - torch.profiler.ProfilerActivity.CUDA,

86 - ],

87 - on_trace_ready=torch.profiler.tensorboard_trace_handler(

88 - str(profile_dir)),

84 + activities=[

85 + torch.profiler.ProfilerActivity.CPU,

86 + torch.profiler.ProfilerActivity.CUDA,

87 + ],

88 + on_trace_ready=torch.profiler.tensorboard_trace_handler(

89 + str(profile_dir)

90 + ),

91 ) as p:

92 llm_generate()

93 print(p.key_averages().table(sort_by="self_cuda_time_total"))

@@ -103,8 +105,9 @@ def main(args: argparse.Namespace):

105 if args.profile:

106 profile_dir = args.profile_result_dir

107 if not profile_dir:

106 - profile_dir = (Path(".") / "vllm_benchmark_result" /

107 - f"latency_result_{time.time()}")

108 + profile_dir = (

109 + Path(".") / "vllm_benchmark_result" / f"latency_result_{time.time()}"

110 + )

111 print(f"Profiling (results will be saved to '{profile_dir}')...")

112 run_to_completion(profile_dir=profile_dir)

113 return

@@ -135,7 +138,8 @@ def main(args: argparse.Namespace):

138 if __name__ == "__main__":

139 parser = FlexibleArgumentParser(

140 description="Benchmark the latency of processing a single batch of "

138 - "requests till completion.")

141 + "requests till completion."

142 + )

143 parser.add_argument("--input-len", type=int, default=32)

144 parser.add_argument("--output-len", type=int, default=128)

145 parser.add_argument("--batch-size", type=int, default=8)

@@ -152,10 +156,9 @@ if __name__ == "__main__":

156 default=10,

157 help="Number of iterations to run for warmup.",

158 )

155 - parser.add_argument("--num-iters",

156 - type=int,

157 - default=30,

158 - help="Number of iterations to run.")

159 + parser.add_argument(

160 + "--num-iters", type=int, default=30, help="Number of iterations to run."

161 + )

162 parser.add_argument(

163 "--profile",

164 action="store_true",

@@ -165,8 +168,10 @@ if __name__ == "__main__":

168 "--profile-result-dir",

169 type=str,

170 default=None,

168 - help=("path to save the pytorch profiler output. Can be visualized "

169 - "with ui.perfetto.dev or Tensorboard."),

171 + help=(

172 + "path to save the pytorch profiler output. Can be visualized "

173 + "with ui.perfetto.dev or Tensorboard."

174 + ),

175 )

176 parser.add_argument(

177 "--output-json",

@@ -177,8 +182,10 @@ if __name__ == "__main__":

182 parser.add_argument(

183 "--disable-detokenize",

184 action="store_true",

180 - help=("Do not detokenize responses (i.e. do not include "

181 - "detokenization time in the latency measurement)"),

185 + help=(

186 + "Do not detokenize responses (i.e. do not include "

187 + "detokenization time in the latency measurement)"

188 + ),

189 )

190

191 parser = EngineArgs.add_cli_args(parser)

benchmarks/benchmark_long_document_qa_throughput.py

+51 -40

@@ -76,7 +76,7 @@ def repeat_prompts(prompts, repeat_count, mode: str):

76 - 'random': Shuffle the prompts randomly after repetition.

77 - 'tile': Repeat the entire prompt list in sequence.

78 Example: [1, 2, 3] -> [1, 2, 3, 1, 2, 3].

79 - - 'interleave': Repeat each prompt consecutively before moving to

79 + - 'interleave': Repeat each prompt consecutively before moving to

80 the next. Example: [1, 2, 3] -> [1, 1, 2, 2, 3, 3].

82 Returns:

@@ -86,20 +86,21 @@ def repeat_prompts(prompts, repeat_count, mode: str):

86 ValueError: If an invalid mode is provided.

87 """

88 print("Repeat mode: ", mode)

89 - if mode == 'random':

89 + if mode == "random":

90 repeated_prompts = prompts * repeat_count

91 random.shuffle(repeated_prompts)

92 return repeated_prompts

93 - elif mode == 'tile':

93 + elif mode == "tile":

94 return prompts * repeat_count

95 - elif mode == 'interleave':

95 + elif mode == "interleave":

96 repeated_prompts = []

97 for prompt in prompts:

98 repeated_prompts.extend([prompt] * repeat_count)

99 return repeated_prompts

100 else:

101 - raise ValueError(f"Invalid mode: {mode}, only support "

102 - "'random', 'tile', 'interleave'")

101 + raise ValueError(

102 + f"Invalid mode: {mode}, only support 'random', 'tile', 'interleave'"

103 + )

104

105

106 def main(args):

@@ -109,16 +110,16 @@ def main(args):

110 # we append the document id at the beginning to avoid any of the document

111 # being the prefix of other documents

112 prompts = [

112 - str(i) + ' '.join(['hi'] * args.document_length)

113 + str(i) + " ".join(["hi"] * args.document_length)

114 for i in range(args.num_documents)

115 ]

116

117 prompts = repeat_prompts(prompts, args.repeat_count, mode=args.repeat_mode)

118

119 warmup_prompts = [

119 - "This is warm up request " + str(i) + \

120 - ' '.join(['hi'] * args.document_length)

121 - for i in range(args.num_documents)]

120 + "This is warm up request " + str(i) + " ".join(["hi"] * args.document_length)

121 + for i in range(args.num_documents)

122 + ]

123

124 # Create the LLM engine

125 engine_args = EngineArgs.from_cli_args(args)

@@ -142,42 +143,52 @@ def main(args):

143

144 if __name__ == "__main__":

145 parser = FlexibleArgumentParser(

145 - description=

146 - 'Benchmark the performance with or without automatic prefix caching.')

146 + description="Benchmark the performance with or "

147 + "without automatic prefix caching."

148 + )

149

150 parser.add_argument(

149 - '--document-length',

151 + "--document-length",

152 type=int,

153 # Roughly the number of tokens for a system paper,

154 # excluding images

155 default=20000,

154 - help='Range of input lengths for sampling prompts,'

155 - 'specified as "min:max" (e.g., "128:256").')

156 -

157 - parser.add_argument('--num-documents',

158 - type=int,

159 - default=8,

160 - help='Range of input lengths for sampling prompts,'

161 - 'specified as "min:max" (e.g., "128:256").')

162 -

163 - parser.add_argument('--output-len', type=int, default=10)

164 -

165 - parser.add_argument('--repeat-count',

166 - type=int,

167 - default=2,

168 - help='Number of times to repeat each prompt')

169 -

170 - parser.add_argument("--repeat-mode",

171 - type=str,

172 - default='random',

173 - help='The mode to repeat prompts. The supported '

174 - 'modes are "random", "tile", and "interleave". '

175 - 'See repeat_prompts() in the source code for details.')

176 -

177 - parser.add_argument("--shuffle-seed",

178 - type=int,

179 - default=0,

180 - help='Random seed when the repeat mode is "random"')

156 + help="Range of input lengths for sampling prompts, "

157 + 'specified as "min:max" (e.g., "128:256").',

158 + )

159 +

160 + parser.add_argument(

161 + "--num-documents",

162 + type=int,

163 + default=8,

164 + help="Range of input lengths for sampling prompts, "

165 + 'specified as "min:max" (e.g., "128:256").',

166 + )

167 +

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

169 +

170 + parser.add_argument(

171 + "--repeat-count",

172 + type=int,

173 + default=2,

174 + help="Number of times to repeat each prompt",

175 + )

176 +

177 + parser.add_argument(

178 + "--repeat-mode",

179 + type=str,

180 + default="random",

181 + help="The mode to repeat prompts. The supported "

182 + 'modes are "random", "tile", and "interleave". '

183 + "See repeat_prompts() in the source code for details.",

184 + )

185 +

186 + parser.add_argument(

187 + "--shuffle-seed",

188 + type=int,

189 + default=0,

190 + help='Random seed when the repeat mode is "random"',

191 + )

192

193 parser = EngineArgs.add_cli_args(parser)

194 args = parser.parse_args()

benchmarks/benchmark_prefix_caching.py

+64 -52

@@ -63,8 +63,7 @@ class Request:

63 output_len: int

66 -def sample_tokens(tokenizer: PreTrainedTokenizerBase,

67 - length: int) -> list[int]:

66 +def sample_tokens(tokenizer: PreTrainedTokenizerBase, length: int) -> list[int]:

67 vocab = tokenizer.get_vocab()

68 all_special_ids = set(tokenizer.all_special_ids)

@@ -91,8 +90,10 @@ def sample_requests_from_dataset(

90 # Filter out the conversations with less than 2 turns.

91 dataset = [data for data in dataset if len(data["conversations"]) >= 2]

92 # Only keep the first two turns of each conversation.

94 - dataset = [(data["conversations"][0]["value"],

95 - data["conversations"][1]["value"]) for data in dataset]

93 + dataset = [

94 + (data["conversations"][0]["value"], data["conversations"][1]["value"])

95 + for data in dataset

96 + ]

98 # Shuffle the dataset.

99 random.shuffle(dataset)

@@ -113,8 +114,9 @@ def sample_requests_from_dataset(

114 completion = dataset[i][1]

115 completion_token_ids = tokenizer(completion).input_ids

116 prompt_len = len(prompt_token_ids)

116 - output_len = (len(completion_token_ids)

117 - if fixed_output_len is None else fixed_output_len)

117 + output_len = (

118 + len(completion_token_ids) if fixed_output_len is None else fixed_output_len

119 + )

120 if min_len <= prompt_len <= max_len:

121 filtered_requests.append(Request(prompt, prompt_len, output_len))

122

@@ -128,27 +130,27 @@ def sample_requests_from_random(

130 fixed_output_len: Optional[int],

131 prefix_len: int,

132 ) -> list[Request]:

131 -

133 requests = []

134 prefix_token_ids = sample_tokens(tokenizer, prefix_len)

135 min_len, max_len = input_length_range

136

137 for i in range(num_requests):

138 unique_part_token_ids = sample_tokens(

138 - tokenizer,

139 - random.randint(min_len - prefix_len, max_len - prefix_len))

139 + tokenizer, random.randint(min_len - prefix_len, max_len - prefix_len)

140 + )

141 prompt_token_ids = prefix_token_ids + unique_part_token_ids

142 prompt = tokenizer.decode(prompt_token_ids)

143 prompt_len = len(prompt_token_ids)

143 - assert (min_len <= prompt_len <= max_len

144 - ), f"prompt_len {prompt_len} out of range {min_len}:{max_len}"

144 + assert min_len <= prompt_len <= max_len, (

145 + f"prompt_len {prompt_len} out of range {min_len}:{max_len}"

146 + )

147 requests.append(Request(prompt, prompt_len, fixed_output_len))

148 return requests

149

150

149 -def repeat_and_sort_requests(requests: list[Request],

150 - repeat_count: int,

151 - sort: bool = False) -> list[str]:

151 +def repeat_and_sort_requests(

152 + requests: list[Request], repeat_count: int, sort: bool = False

153 +) -> list[str]:

154 repeated_requests = requests * repeat_count

155 if sort:

156 repeated_requests.sort(key=lambda x: x[1])

@@ -159,14 +161,14 @@ def repeat_and_sort_requests(requests: list[Request],

161

162 def main(args):

163 tokenizer = get_tokenizer(args.model, trust_remote_code=True)

162 - input_length_range = tuple(map(int, args.input_length_range.split(':')))

164 + input_length_range = tuple(map(int, args.input_length_range.split(":")))

165 random.seed(args.seed)

166 if args.dataset_path is not None:

167 if args.prefix_len > 0:

166 - raise ValueError("prefix-len is not supported when "

167 - "dataset-path is provided.")

168 - print(f"Start to sample {args.num_prompts} prompts "

169 - f"from {args.dataset_path}")

168 + raise ValueError(

169 + "prefix-len is not supported when dataset-path is provided."

170 + )

171 + print(f"Start to sample {args.num_prompts} prompts from {args.dataset_path}")

172 filtered_requests = sample_requests_from_dataset(

173 dataset_path=args.dataset_path,

174 num_requests=args.num_prompts,

@@ -196,14 +198,16 @@ def main(args):

198

199 llm = LLM(**dataclasses.asdict(engine_args))

200

199 - sampling_params = SamplingParams(temperature=0,

200 - max_tokens=args.output_len,

201 - detokenize=not args.disable_detokenize)

201 + sampling_params = SamplingParams(

202 + temperature=0,

203 + max_tokens=args.output_len,

204 + detokenize=not args.disable_detokenize,

205 + )

206

207 print("Testing filtered requests")

204 - prompts = repeat_and_sort_requests(filtered_requests,

205 - repeat_count=args.repeat_count,

206 - sort=args.sort)

208 + prompts = repeat_and_sort_requests(

209 + filtered_requests, repeat_count=args.repeat_count, sort=args.sort

210 + )

211

212 print("------start generating------")

213 test_prefix(

@@ -215,29 +219,35 @@ def main(args):

219

220 if __name__ == "__main__":

221 parser = FlexibleArgumentParser(

218 - description=

219 - 'Benchmark the performance with or without automatic prefix caching.')

220 - parser.add_argument("--dataset-path",

221 - type=str,

222 - default=None,

223 - help="Path to the dataset.")

224 - parser.add_argument('--output-len', type=int, default=10)

225 - parser.add_argument('--num-prompts',

226 - type=int,

227 - required=True,

228 - help="Number of the prompts sampled from dataset")

229 - parser.add_argument('--repeat-count',

230 - type=int,

231 - default=1,

232 - help='Number of times to repeat each prompt')

233 - parser.add_argument('--sort',

234 - action='store_true',

235 - help='Sort prompts by input length')

236 - parser.add_argument('--input-length-range',

237 - type=str,

238 - required=True,

239 - help='Range of input lengths for sampling prompts,'

240 - 'specified as "min:max" (e.g., "128:256").')

222 + description="Benchmark the performance with or without "

223 + "automatic prefix caching."

224 + )

225 + parser.add_argument(

226 + "--dataset-path", type=str, default=None, help="Path to the dataset."

227 + )

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

229 + parser.add_argument(

230 + "--num-prompts",

231 + type=int,

232 + required=True,

233 + help="Number of the prompts sampled from dataset",

234 + )

235 + parser.add_argument(

236 + "--repeat-count",

237 + type=int,

238 + default=1,

239 + help="Number of times to repeat each prompt",

240 + )

241 + parser.add_argument(

242 + "--sort", action="store_true", help="Sort prompts by input length"

243 + )

244 + parser.add_argument(

245 + "--input-length-range",

246 + type=str,

247 + required=True,

248 + help="Range of input lengths for sampling prompts,"

249 + 'specified as "min:max" (e.g., "128:256").',

250 + )

251 parser.add_argument(

252 "--prefix-len",

253 type=int,

@@ -248,10 +258,12 @@ if __name__ == "__main__":

258 "when dataset-path is not provided.",

259 )

260 parser.add_argument(

251 - '--disable-detokenize',

252 - action='store_true',

253 - help=("Do not detokenize responses (i.e. do not include "

254 - "detokenization time in the latency measurement)"),

261 + "--disable-detokenize",

262 + action="store_true",

263 + help=(

264 + "Do not detokenize responses (i.e. do not include "

265 + "detokenization time in the latency measurement)"

266 + ),

267 )

268

269 parser = EngineArgs.add_cli_args(parser)

benchmarks/benchmark_prioritization.py

+70 -52

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 """Benchmark offline prioritization."""

3 +

4 import argparse

5 import dataclasses

6 import json

@@ -13,7 +14,7 @@ from vllm.engine.arg_utils import EngineArgs

14 from vllm.utils import FlexibleArgumentParser

16 -#Select a equi-probable random priority

17 +# Select a equi-probable random priority

18 def get_random_flag():

19 return 0 if random.random() < 0.5 else 1

@@ -33,8 +34,10 @@ def sample_requests(

34 # Filter out the conversations with less than 2 turns.

35 dataset = [data for data in dataset if len(data["conversations"]) >= 2]

36 # Only keep the first two turns of each conversation.

36 - dataset = [(data["conversations"][0]["value"],

37 - data["conversations"][1]["value"]) for data in dataset]

37 + dataset = [

38 + (data["conversations"][0]["value"], data["conversations"][1]["value"])

39 + for data in dataset

40 + ]

42 # Shuffle the dataset.

43 random.shuffle(dataset)

@@ -51,8 +54,9 @@ def sample_requests(

54 completion = dataset[i][1]

55 completion_token_ids = tokenizer(completion).input_ids

56 prompt_len = len(prompt_token_ids)

54 - output_len = len(completion_token_ids

55 - ) if fixed_output_len is None else fixed_output_len

57 + output_len = (

58 + len(completion_token_ids) if fixed_output_len is None else fixed_output_len

59 + )

60 if prompt_len < 4 or output_len < 4:

61 # Prune too short sequences.

62 continue

@@ -74,13 +78,16 @@ def run_vllm(

78 disable_detokenize: bool = False,

79 ) -> float:

80 from vllm import LLM, SamplingParams

81 +

82 llm = LLM(**dataclasses.asdict(engine_args))

84 assert all(

85 llm.llm_engine.model_config.max_model_len >= (request[1] + request[2])

81 - for request in requests), (

82 - "Please ensure that max_model_len is greater than the sum of"

83 - " input_len and output_len for all requests.")

86 + for request in requests

87 + ), (

88 + "Please ensure that max_model_len is greater than the sum of"

89 + " input_len and output_len for all requests."

90 + )

92 # Add the requests to the engine.

93 prompts = []

@@ -97,7 +104,8 @@ def run_vllm(

104 ignore_eos=True,

105 max_tokens=output_len,

106 detokenize=not disable_detokenize,

100 - ))

107 + )

108 + )

109

110 start = time.perf_counter()

111 llm.generate(prompts, sampling_params, priority=priority, use_tqdm=True)

@@ -111,26 +119,33 @@ def main(args: argparse.Namespace):

119

120 # Sample the requests.

121 tokenizer = AutoTokenizer.from_pretrained(

114 - args.tokenizer, trust_remote_code=args.trust_remote_code)

122 + args.tokenizer, trust_remote_code=args.trust_remote_code

123 + )

124 if args.dataset is None:

125 # Synthesize a prompt with the given input length.

126 prompt = "hi" * (args.input_len - 1)

118 - requests = [(prompt, args.input_len, args.output_len,

119 - get_random_flag()) for _ in range(args.num_prompts)]

127 + requests = [

128 + (prompt, args.input_len, args.output_len, get_random_flag())

129 + for _ in range(args.num_prompts)

130 + ]

131 else:

121 - requests = sample_requests(args.dataset, args.num_prompts, tokenizer,

122 - args.output_len)

132 + requests = sample_requests(

133 + args.dataset, args.num_prompts, tokenizer, args.output_len

134 + )

135

136 if args.backend == "vllm":

125 - elapsed_time = run_vllm(requests, args.n,

126 - EngineArgs.from_cli_args(args),

127 - args.disable_detokenize)

137 + elapsed_time = run_vllm(

138 + requests, args.n, EngineArgs.from_cli_args(args), args.disable_detokenize

139 + )

140 else:

141 raise ValueError(f"Unknown backend: {args.backend}")

130 - total_num_tokens = sum(prompt_len + output_len

131 - for _, prompt_len, output_len, priority in requests)

132 - print(f"Throughput: {len(requests) / elapsed_time:.2f} requests/s, "

133 - f"{total_num_tokens / elapsed_time:.2f} tokens/s")

142 + total_num_tokens = sum(

143 + prompt_len + output_len for _, prompt_len, output_len, priority in requests

144 + )

145 + print(

146 + f"Throughput: {len(requests) / elapsed_time:.2f} requests/s, "

147 + f"{total_num_tokens / elapsed_time:.2f} tokens/s"

148 + )

149

150 # Output JSON results if specified

151 if args.output_json:

@@ -147,41 +162,44 @@ def main(args: argparse.Namespace):

162

163 if __name__ == "__main__":

164 parser = FlexibleArgumentParser(description="Benchmark the throughput.")

150 - parser.add_argument("--backend",

151 - type=str,

152 - choices=["vllm", "hf", "mii"],

153 - default="vllm")

154 - parser.add_argument("--dataset",

155 - type=str,

156 - default=None,

157 - help="Path to the dataset.")

158 - parser.add_argument("--input-len",

159 - type=int,

160 - default=None,

161 - help="Input prompt length for each request")

162 - parser.add_argument("--output-len",

163 - type=int,

164 - default=None,

165 - help="Output length for each request. Overrides the "

166 - "output length from the dataset.")

167 - parser.add_argument("--n",

168 - type=int,

169 - default=1,

170 - help="Number of generated sequences per prompt.")

171 - parser.add_argument("--num-prompts",

172 - type=int,

173 - default=200,

174 - help="Number of prompts to process.")

165 parser.add_argument(

176 - '--output-json',

166 + "--backend", type=str, choices=["vllm", "hf", "mii"], default="vllm"

167 + )

168 + parser.add_argument(

169 + "--dataset", type=str, default=None, help="Path to the dataset."

170 + )

171 + parser.add_argument(

172 + "--input-len",

173 + type=int,

174 + default=None,

175 + help="Input prompt length for each request",

176 + )

177 + parser.add_argument(

178 + "--output-len",

179 + type=int,

180 + default=None,

181 + help="Output length for each request. Overrides the "

182 + "output length from the dataset.",

183 + )

184 + parser.add_argument(

185 + "--n", type=int, default=1, help="Number of generated sequences per prompt."

186 + )

187 + parser.add_argument(

188 + "--num-prompts", type=int, default=200, help="Number of prompts to process."

189 + )

190 + parser.add_argument(

191 + "--output-json",

192 type=str,

193 default=None,

179 - help='Path to save the throughput results in JSON format.')

194 + help="Path to save the throughput results in JSON format.",

195 + )

196 parser.add_argument(

181 - '--disable-detokenize',

182 - action='store_true',

183 - help=("Do not detokenize responses (i.e. do not include "

184 - "detokenization time in the latency measurement)"),

197 + "--disable-detokenize",

198 + action="store_true",

199 + help=(

200 + "Do not detokenize responses (i.e. do not include "

201 + "detokenization time in the latency measurement)"

202 + ),

203 )

204

205 parser = EngineArgs.add_cli_args(parser)

benchmarks/benchmark_serving.py

+337 -248

@@ -20,6 +20,7 @@ On the client side, run:

20 --endpoint /generate_stream

21 to the end of the command above.

22 """

23 +

24 import argparse

25 import asyncio

26 import gc

@@ -34,12 +35,16 @@ from datetime import datetime

35 from typing import Any, Optional

37 import numpy as np

37 -from backend_request_func import (ASYNC_REQUEST_FUNCS,

38 - OPENAI_COMPATIBLE_BACKENDS, RequestFuncInput,

39 - RequestFuncOutput)

38 from tqdm.asyncio import tqdm

39 from transformers import PreTrainedTokenizerBase

41 +from backend_request_func import (

42 + ASYNC_REQUEST_FUNCS,

43 + OPENAI_COMPATIBLE_BACKENDS,

44 + RequestFuncInput,

45 + RequestFuncOutput,

46 +)

47 +

48 try:

49 from vllm.transformers_utils.tokenizer import get_tokenizer

50 except ImportError:

@@ -50,12 +55,21 @@ try:

55 except ImportError:

56 from argparse import ArgumentParser as FlexibleArgumentParser

53 -from benchmark_dataset import (AIMODataset, ASRDataset, BurstGPTDataset,

54 - ConversationDataset, HuggingFaceDataset,

55 - InstructCoderDataset, MTBenchDataset,

56 - NextEditPredictionDataset, RandomDataset,

57 - SampleRequest, ShareGPTDataset, SonnetDataset,

58 - VisionArenaDataset)

58 +from benchmark_dataset import (

59 + AIMODataset,

60 + ASRDataset,

61 + BurstGPTDataset,

62 + ConversationDataset,

63 + HuggingFaceDataset,

64 + InstructCoderDataset,

65 + MTBenchDataset,

66 + NextEditPredictionDataset,

67 + RandomDataset,

68 + SampleRequest,

69 + ShareGPTDataset,

70 + SonnetDataset,

71 + VisionArenaDataset,

72 +)

73 from benchmark_utils import convert_to_pytorch_benchmark_format, write_to_json

75 MILLISECONDS_TO_SECONDS_CONVERSION = 1000

@@ -118,7 +132,8 @@ async def get_request(

132

133 # Calculate scale parameter theta to maintain the desired request_rate.

134 assert burstiness > 0, (

121 - f"A positive burstiness factor is expected, but given {burstiness}.")

135 + f"A positive burstiness factor is expected, but given {burstiness}."

136 + )

137 theta = 1.0 / (request_rate * burstiness)

138

139 for request in input_requests:

@@ -164,8 +179,10 @@ def calculate_metrics(

179 # bundled together

180 # Note : this may inflate the output token count slightly

181 output_len = len(

167 - tokenizer(outputs[i].generated_text,

168 - add_special_tokens=False).input_ids)

182 + tokenizer(

183 + outputs[i].generated_text, add_special_tokens=False

184 + ).input_ids

185 + )

186 actual_output_lens.append(output_len)

187 total_input += input_requests[i].prompt_len

188 tpot = 0

@@ -188,16 +205,19 @@ def calculate_metrics(

205

206 if "ttft" in goodput_config_dict:

207 valid_metrics.append(ttfts)

191 - slo_values.append(goodput_config_dict["ttft"] /

192 - MILLISECONDS_TO_SECONDS_CONVERSION)

208 + slo_values.append(

209 + goodput_config_dict["ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION

210 + )

211 if "tpot" in goodput_config_dict:

212 valid_metrics.append(all_tpots)

195 - slo_values.append(goodput_config_dict["tpot"] /

196 - MILLISECONDS_TO_SECONDS_CONVERSION)

213 + slo_values.append(

214 + goodput_config_dict["tpot"] / MILLISECONDS_TO_SECONDS_CONVERSION

215 + )

216 if "e2el" in goodput_config_dict:

217 valid_metrics.append(e2els)

199 - slo_values.append(goodput_config_dict["e2el"] /

200 - MILLISECONDS_TO_SECONDS_CONVERSION)

218 + slo_values.append(

219 + goodput_config_dict["e2el"] / MILLISECONDS_TO_SECONDS_CONVERSION

220 + )

221

222 for req_metric in zip(*valid_metrics):

223 is_good_req = all([s >= r for s, r in zip(slo_values, req_metric)])

@@ -208,7 +228,8 @@ def calculate_metrics(

228 warnings.warn(

229 "All requests failed. This is likely due to a misconfiguration "

230 "on the benchmark arguments.",

211 - stacklevel=2)

231 + stacklevel=2,

232 + )

233 metrics = BenchmarkMetrics(

234 completed=completed,

235 total_input=total_input,

@@ -217,27 +238,31 @@ def calculate_metrics(

238 request_goodput=good_completed / dur_s,

239 output_throughput=sum(actual_output_lens) / dur_s,

240 total_token_throughput=(total_input + sum(actual_output_lens)) / dur_s,

220 - mean_ttft_ms=np.mean(ttfts or 0) *

221 - 1000, # ttfts is empty if streaming is not supported by backend

241 + mean_ttft_ms=np.mean(ttfts or 0)

242 + * 1000, # ttfts is empty if streaming is not supported by backend

243 std_ttft_ms=np.std(ttfts or 0) * 1000,

244 median_ttft_ms=np.median(ttfts or 0) * 1000,

224 - percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000)

225 - for p in selected_percentiles],

245 + percentiles_ttft_ms=[

246 + (p, np.percentile(ttfts or 0, p) * 1000) for p in selected_percentiles

247 + ],

248 mean_tpot_ms=np.mean(tpots or 0) * 1000,

249 std_tpot_ms=np.std(tpots or 0) * 1000,

250 median_tpot_ms=np.median(tpots or 0) * 1000,

229 - percentiles_tpot_ms=[(p, np.percentile(tpots or 0, p) * 1000)

230 - for p in selected_percentiles],

251 + percentiles_tpot_ms=[

252 + (p, np.percentile(tpots or 0, p) * 1000) for p in selected_percentiles

253 + ],

254 mean_itl_ms=np.mean(itls or 0) * 1000,

255 std_itl_ms=np.std(itls or 0) * 1000,

256 median_itl_ms=np.median(itls or 0) * 1000,

234 - percentiles_itl_ms=[(p, np.percentile(itls or 0, p) * 1000)

235 - for p in selected_percentiles],

257 + percentiles_itl_ms=[

258 + (p, np.percentile(itls or 0, p) * 1000) for p in selected_percentiles

259 + ],

260 mean_e2el_ms=np.mean(e2els or 0) * 1000,

261 std_e2el_ms=np.std(e2els or 0) * 1000,

262 median_e2el_ms=np.median(e2els or 0) * 1000,

239 - percentiles_e2el_ms=[(p, np.percentile(e2els or 0, p) * 1000)

240 - for p in selected_percentiles],

263 + percentiles_e2el_ms=[

264 + (p, np.percentile(e2els or 0, p) * 1000) for p in selected_percentiles

265 + ],

266 )

267

268 return metrics, actual_output_lens

@@ -270,10 +295,12 @@ async def benchmark(

295 raise ValueError(f"Unknown backend: {backend}")

296

297 print("Starting initial single prompt test run...")

273 - test_prompt, test_prompt_len, test_output_len, test_mm_content = \

274 - input_requests[0].prompt, input_requests[0].prompt_len, \

275 - input_requests[0].expected_output_len, \

276 - input_requests[0].multi_modal_data

298 + test_prompt, test_prompt_len, test_output_len, test_mm_content = (

299 + input_requests[0].prompt,

300 + input_requests[0].prompt_len,

301 + input_requests[0].expected_output_len,

302 + input_requests[0].multi_modal_data,

303 + )

304

305 assert test_mm_content is None or isinstance(test_mm_content, dict)

306 test_input = RequestFuncInput(

@@ -293,36 +320,36 @@ async def benchmark(

320 if not test_output.success:

321 raise ValueError(

322 "Initial test run failed - Please make sure benchmark arguments "

296 - f"are correctly specified. Error: {test_output.error}")

323 + f"are correctly specified. Error: {test_output.error}"

324 + )

325 else:

326 print("Initial test run completed. Starting main benchmark run...")

327

328 if lora_modules:

329 # For each input request, choose a LoRA module at random.

330 lora_modules = iter(

303 - [random.choice(lora_modules) \

304 - for _ in range(len(input_requests))])

331 + [random.choice(lora_modules) for _ in range(len(input_requests))]

332 + )

333

334 if profile:

335 print("Starting profiler...")

308 - profile_input = RequestFuncInput(model=model_id,

309 - model_name=model_name,

310 - prompt=test_prompt,

311 - api_url=base_url + "/start_profile",

312 - prompt_len=test_prompt_len,

313 - output_len=test_output_len,

314 - logprobs=logprobs,

315 - multi_modal_content=test_mm_content,

316 - ignore_eos=ignore_eos,

317 - extra_body=extra_body)

336 + profile_input = RequestFuncInput(

337 + model=model_id,

338 + model_name=model_name,

339 + prompt=test_prompt,

340 + api_url=base_url + "/start_profile",

341 + prompt_len=test_prompt_len,

342 + output_len=test_output_len,

343 + logprobs=logprobs,

344 + multi_modal_content=test_mm_content,

345 + ignore_eos=ignore_eos,

346 + extra_body=extra_body,

347 + )

348 profile_output = await request_func(request_func_input=profile_input)

349 if profile_output.success:

350 print("Profiler started")

351

322 - if burstiness == 1.0:

323 - distribution = "Poisson process"

324 - else:

325 - distribution = "Gamma distribution"

352 + distribution = "Poisson process" if burstiness == 1.0 else "Gamma distribution"

353

354 print(f"Traffic request rate: {request_rate}")

355 print(f"Burstiness factor: {burstiness} ({distribution})")

@@ -334,42 +361,45 @@ async def benchmark(

361 # and it will simplify the code in limited_request_func.

362 # semaphore = (asyncio.Semaphore(max_concurrency)

363 # if max_concurrency else contextlib.nullcontext())

337 - semaphore = (asyncio.Semaphore(max_concurrency)

338 - if max_concurrency else None)

364 + semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None

365

366 async def limited_request_func(request_func_input, pbar):

367 if semaphore is None:

342 - return await request_func(request_func_input=request_func_input,

343 - pbar=pbar)

368 + return await request_func(request_func_input=request_func_input, pbar=pbar)

369 async with semaphore:

345 - return await request_func(request_func_input=request_func_input,

346 - pbar=pbar)

370 + return await request_func(request_func_input=request_func_input, pbar=pbar)

371

372 benchmark_start_time = time.perf_counter()

373 tasks: list[asyncio.Task] = []

374 async for request in get_request(input_requests, request_rate, burstiness):

351 - prompt, prompt_len, output_len, mm_content = request.prompt, \

352 - request.prompt_len, request.expected_output_len, \

353 - request.multi_modal_data

375 + prompt, prompt_len, output_len, mm_content = (

376 + request.prompt,

377 + request.prompt_len,

378 + request.expected_output_len,

379 + request.multi_modal_data,

380 + )

381 req_model_id, req_model_name = model_id, model_name

382 if lora_modules:

383 req_lora_module = next(lora_modules)

384 req_model_id, req_model_name = req_lora_module, req_lora_module

385

359 - request_func_input = RequestFuncInput(model=req_model_id,

360 - model_name=req_model_name,

361 - prompt=prompt,

362 - api_url=api_url,

363 - prompt_len=prompt_len,

364 - output_len=output_len,

365 - logprobs=logprobs,

366 - multi_modal_content=mm_content,

367 - ignore_eos=ignore_eos,

368 - extra_body=extra_body)

386 + request_func_input = RequestFuncInput(

387 + model=req_model_id,

388 + model_name=req_model_name,

389 + prompt=prompt,

390 + api_url=api_url,

391 + prompt_len=prompt_len,

392 + output_len=output_len,

393 + logprobs=logprobs,

394 + multi_modal_content=mm_content,

395 + ignore_eos=ignore_eos,

396 + extra_body=extra_body,

397 + )

398 tasks.append(

399 asyncio.create_task(

371 - limited_request_func(request_func_input=request_func_input,

372 - pbar=pbar)))

400 + limited_request_func(request_func_input=request_func_input, pbar=pbar)

401 + )

402 + )

403 outputs: list[RequestFuncOutput] = await asyncio.gather(*tasks)

404

405 if profile:

@@ -401,22 +431,32 @@ async def benchmark(

431 goodput_config_dict=goodput_config_dict,

432 )

433

404 - print("{s:{c}^{n}}".format(s=' Serving Benchmark Result ', n=50, c='='))

434 + print("{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="="))

435 print("{:<40} {:<10}".format("Successful requests:", metrics.completed))

406 - print("{:<40} {:<10.2f}".format("Benchmark duration (s):",

407 - benchmark_duration))

436 + print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration))

437 print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input))

409 - print("{:<40} {:<10}".format("Total generated tokens:",

410 - metrics.total_output))

411 - print("{:<40} {:<10.2f}".format("Request throughput (req/s):",

412 - metrics.request_throughput))

438 + print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output))

439 + print(

440 + "{:<40} {:<10.2f}".format(

441 + "Request throughput (req/s):", metrics.request_throughput

442 + )

443 + )

444 if goodput_config_dict:

414 - print("{:<40} {:<10.2f}".format("Request goodput (req/s):",

415 - metrics.request_goodput))

416 - print("{:<40} {:<10.2f}".format("Output token throughput (tok/s):",

417 - metrics.output_throughput))

418 - print("{:<40} {:<10.2f}".format("Total Token throughput (tok/s):",

419 - metrics.total_token_throughput))

445 + print(

446 + "{:<40} {:<10.2f}".format(

447 + "Request goodput (req/s):", metrics.request_goodput

448 + )

449 + )

450 + print(

451 + "{:<40} {:<10.2f}".format(

452 + "Output token throughput (tok/s):", metrics.output_throughput

453 + )

454 + )

455 + print(

456 + "{:<40} {:<10.2f}".format(

457 + "Total Token throughput (tok/s):", metrics.total_token_throughput

458 + )

459 + )

460

461 result = {

462 "duration": benchmark_duration,

@@ -424,8 +464,7 @@ async def benchmark(

464 "total_input_tokens": metrics.total_input,

465 "total_output_tokens": metrics.total_output,

466 "request_throughput": metrics.request_throughput,

427 - "request_goodput:":

428 - metrics.request_goodput if goodput_config_dict else None,

467 + "request_goodput:": metrics.request_goodput if goodput_config_dict else None,

468 "output_throughput": metrics.output_throughput,

469 "total_token_throughput": metrics.total_token_throughput,

470 "input_lens": [output.prompt_len for output in outputs],

@@ -448,29 +487,35 @@ async def benchmark(

487 # metric.

488 if metric_attribute_name not in selected_percentile_metrics:

489 return

451 - print("{s:{c}^{n}}".format(s=metric_header, n=50, c='-'))

452 - print("{:<40} {:<10.2f}".format(

453 - f"Mean {metric_name} (ms):",

454 - getattr(metrics, f"mean_{metric_attribute_name}_ms")))

455 - print("{:<40} {:<10.2f}".format(

456 - f"Median {metric_name} (ms):",

457 - getattr(metrics, f"median_{metric_attribute_name}_ms")))

490 + print("{s:{c}^{n}}".format(s=metric_header, n=50, c="-"))

491 + print(

492 + "{:<40} {:<10.2f}".format(

493 + f"Mean {metric_name} (ms):",

494 + getattr(metrics, f"mean_{metric_attribute_name}_ms"),

495 + )

496 + )

497 + print(

498 + "{:<40} {:<10.2f}".format(

499 + f"Median {metric_name} (ms):",

500 + getattr(metrics, f"median_{metric_attribute_name}_ms"),

501 + )

502 + )

503 result[f"mean_{metric_attribute_name}_ms"] = getattr(

459 - metrics, f"mean_{metric_attribute_name}_ms")

504 + metrics, f"mean_{metric_attribute_name}_ms"

505 + )

506 result[f"median_{metric_attribute_name}_ms"] = getattr(

461 - metrics, f"median_{metric_attribute_name}_ms")

507 + metrics, f"median_{metric_attribute_name}_ms"

508 + )

509 result[f"std_{metric_attribute_name}_ms"] = getattr(

463 - metrics, f"std_{metric_attribute_name}_ms")

464 - for p, value in getattr(metrics,

465 - f"percentiles_{metric_attribute_name}_ms"):

510 + metrics, f"std_{metric_attribute_name}_ms"

511 + )

512 + for p, value in getattr(metrics, f"percentiles_{metric_attribute_name}_ms"):

513 p_word = str(int(p)) if int(p) == p else str(p)

467 - print("{:<40} {:<10.2f}".format(f"P{p_word} {metric_name} (ms):",

468 - value))

514 + print("{:<40} {:<10.2f}".format(f"P{p_word} {metric_name} (ms):", value))

515 result[f"p{p_word}_{metric_attribute_name}_ms"] = value

516

517 process_one_metric("ttft", "TTFT", "Time to First Token")

472 - process_one_metric("tpot", "TPOT",

473 - "Time per Output Token (excl. 1st token)")

518 + process_one_metric("tpot", "TPOT", "Time per Output Token (excl. 1st token)")

519 process_one_metric("itl", "ITL", "Inter-token Latency")

520 process_one_metric("e2el", "E2EL", "End-to-end Latency")

521

@@ -490,12 +535,14 @@ def check_goodput_args(args):

535 raise ValueError(

536 f"Invalid metric name found, {slo_name}: {slo_val}. "

537 "The service level objective name should be one of "

493 - f"{str(VALID_NAMES)}. ")

538 + f"{str(VALID_NAMES)}. "

539 + )

540 if slo_val < 0:

541 raise ValueError(

542 f"Invalid value found, {slo_name}: {slo_val}. "

543 "The service level objective value should be "

498 - "non-negative.")

544 + "non-negative."

545 + )

546 return goodput_config_dict

547

548

@@ -508,31 +555,42 @@ def parse_goodput(slo_pairs):

555 except ValueError as err:

556 raise argparse.ArgumentTypeError(

557 "Invalid format found for service level objectives. "

511 - "Specify service level objectives for goodput as \"KEY:VALUE\" "

558 + 'Specify service level objectives for goodput as "KEY:VALUE" '

559 "pairs, where the key is a metric name, and the value is a "

513 - "number in milliseconds.") from err

560 + "number in milliseconds."

561 + ) from err

562 return goodput_config_dict

563

564

517 -def save_to_pytorch_benchmark_format(args: argparse.Namespace,

518 - results: dict[str, Any],

519 - file_name: str) -> None:

565 +def save_to_pytorch_benchmark_format(

566 + args: argparse.Namespace, results: dict[str, Any], file_name: str

567 +) -> None:

568 metrics = [

521 - "median_ttft_ms", "mean_ttft_ms", "std_ttft_ms", "p99_ttft_ms",

522 - "mean_tpot_ms", "median_tpot_ms", "std_tpot_ms", "p99_tpot_ms",

523 - "median_itl_ms", "mean_itl_ms", "std_itl_ms", "p99_itl_ms"

569 + "median_ttft_ms",

570 + "mean_ttft_ms",

571 + "std_ttft_ms",

572 + "p99_ttft_ms",

573 + "mean_tpot_ms",

574 + "median_tpot_ms",

575 + "std_tpot_ms",

576 + "p99_tpot_ms",

577 + "median_itl_ms",

578 + "mean_itl_ms",

579 + "std_itl_ms",

580 + "p99_itl_ms",

581 ]

582 # These raw data might be useful, but they are rather big. They can be added

583 # later if needed

584 ignored_metrics = ["ttfts", "itls", "generated_texts", "errors"]

585 pt_records = convert_to_pytorch_benchmark_format(

586 args=args,

530 - metrics={k: [results[k]]

531 - for k in metrics},

587 + metrics={k: [results[k]] for k in metrics},

588 extra_info={

589 k: results[k]

534 - for k in results if k not in metrics and k not in ignored_metrics

535 - })

590 + for k in results

591 + if k not in metrics and k not in ignored_metrics

592 + },

593 + )

594 if pt_records:

595 # Don't use json suffix here as we don't want CI to pick it up

596 pt_file = f"{os.path.splitext(file_name)[0]}.pytorch.json"

@@ -557,34 +615,42 @@ def main(args: argparse.Namespace):

615 api_url = f"http://{args.host}:{args.port}{args.endpoint}"

616 base_url = f"http://{args.host}:{args.port}"

617

560 - tokenizer = get_tokenizer(tokenizer_id,

561 - tokenizer_mode=tokenizer_mode,

562 - trust_remote_code=args.trust_remote_code)

618 + tokenizer = get_tokenizer(

619 + tokenizer_id,

620 + tokenizer_mode=tokenizer_mode,

621 + trust_remote_code=args.trust_remote_code,

622 + )

623

624 if args.dataset_name is None:

625 raise ValueError(

626 "Please specify '--dataset-name' and the corresponding "

567 - "'--dataset-path' if required.")

627 + "'--dataset-path' if required."

628 + )

629

630 if args.dataset_name == "sonnet":

631 dataset = SonnetDataset(dataset_path=args.dataset_path)

632 # For the "sonnet" dataset, formatting depends on the backend.

633 if args.backend == "openai-chat":

573 - input_requests = dataset.sample(num_requests=args.num_prompts,

574 - input_len=args.sonnet_input_len,

575 - output_len=args.sonnet_output_len,

576 - prefix_len=args.sonnet_prefix_len,

577 - tokenizer=tokenizer,

578 - return_prompt_formatted=False)

634 + input_requests = dataset.sample(

635 + num_requests=args.num_prompts,

636 + input_len=args.sonnet_input_len,

637 + output_len=args.sonnet_output_len,

638 + prefix_len=args.sonnet_prefix_len,

639 + tokenizer=tokenizer,

640 + return_prompt_formatted=False,

641 + )

642 else:

643 assert tokenizer.chat_template or tokenizer.default_chat_template, (

581 - "Tokenizer/model must have chat template for sonnet dataset.")

582 - input_requests = dataset.sample(num_requests=args.num_prompts,

583 - input_len=args.sonnet_input_len,

584 - output_len=args.sonnet_output_len,

585 - prefix_len=args.sonnet_prefix_len,

586 - tokenizer=tokenizer,

587 - return_prompt_formatted=True)

644 + "Tokenizer/model must have chat template for sonnet dataset."

645 + )

646 + input_requests = dataset.sample(

647 + num_requests=args.num_prompts,

648 + input_len=args.sonnet_input_len,

649 + output_len=args.sonnet_output_len,

650 + prefix_len=args.sonnet_prefix_len,

651 + tokenizer=tokenizer,

652 + return_prompt_formatted=True,

653 + )

654

655 elif args.dataset_name == "hf":

656 # all following datasets are implemented from the

@@ -611,23 +677,30 @@ def main(args: argparse.Namespace):

677 dataset_class = ASRDataset

678 args.hf_split = "train"

679 else:

614 - supported_datasets = set([

615 - dataset_name for cls in HuggingFaceDataset.__subclasses__()

616 - for dataset_name in cls.SUPPORTED_DATASET_PATHS

617 - ])

680 + supported_datasets = set(

681 + [

682 + dataset_name

683 + for cls in HuggingFaceDataset.__subclasses__()

684 + for dataset_name in cls.SUPPORTED_DATASET_PATHS

685 + ]

686 + )

687 raise ValueError(

688 f"Unsupported dataset path: {args.dataset_path}. "

689 "Huggingface dataset only supports dataset_path"

690 f" from one of following: {supported_datasets}. "

691 "Please consider contributing if you would "

623 - "like to add support for additional dataset formats.")

692 + "like to add support for additional dataset formats."

693 + )

694

625 - if (dataset_class.IS_MULTIMODAL and backend not in \

626 - ["openai-chat", "openai-audio"]):

695 + if dataset_class.IS_MULTIMODAL and backend not in [

696 + "openai-chat",

697 + "openai-audio",

698 + ]:

699 # multi-modal benchmark is only available on OpenAI Chat backend.

700 raise ValueError(

629 - "Multi-modal content is only supported on 'openai-chat' and " \

630 - "'openai-audio' backend.")

701 + "Multi-modal content is only supported on 'openai-chat' and "

702 + "'openai-audio' backend."

703 + )

704 input_requests = dataset_class(

705 dataset_path=args.dataset_path,

706 dataset_subset=args.hf_subset,

@@ -642,26 +715,24 @@ def main(args: argparse.Namespace):

715 else:

716 # For datasets that follow a similar structure, use a mapping.

717 dataset_mapping = {

645 - "sharegpt":

646 - lambda: ShareGPTDataset(random_seed=args.seed,

647 - dataset_path=args.dataset_path).sample(

648 - tokenizer=tokenizer,

649 - num_requests=args.num_prompts,

650 - output_len=args.sharegpt_output_len,

651 - ),

652 - "burstgpt":

653 - lambda: BurstGPTDataset(random_seed=args.seed,

654 - dataset_path=args.dataset_path).

655 - sample(tokenizer=tokenizer, num_requests=args.num_prompts),

656 - "random":

657 - lambda: RandomDataset(dataset_path=args.dataset_path).sample(

718 + "sharegpt": lambda: ShareGPTDataset(

719 + random_seed=args.seed, dataset_path=args.dataset_path

720 + ).sample(

721 + tokenizer=tokenizer,

722 + num_requests=args.num_prompts,

723 + output_len=args.sharegpt_output_len,

724 + ),

725 + "burstgpt": lambda: BurstGPTDataset(

726 + random_seed=args.seed, dataset_path=args.dataset_path

727 + ).sample(tokenizer=tokenizer, num_requests=args.num_prompts),

728 + "random": lambda: RandomDataset(dataset_path=args.dataset_path).sample(

729 tokenizer=tokenizer,

730 num_requests=args.num_prompts,

731 prefix_len=args.random_prefix_len,

732 input_len=args.random_input_len,

733 output_len=args.random_output_len,

734 range_ratio=args.random_range_ratio,

664 - )

735 + ),

736 }

737

738 try:

@@ -677,15 +748,16 @@ def main(args: argparse.Namespace):

748 "top_p": args.top_p,

749 "top_k": args.top_k,

750 "min_p": args.min_p,

680 - "temperature": args.temperature

681 - }.items() if v is not None

751 + "temperature": args.temperature,

752 + }.items()

753 + if v is not None

754 }

755

756 # Sampling parameters are only supported by openai-compatible backend.

757 if sampling_params and args.backend not in OPENAI_COMPATIBLE_BACKENDS:

758 raise ValueError(

687 - "Sampling parameters are only supported by openai-compatible "

688 - "backends.")

759 + "Sampling parameters are only supported by openai-compatible backends."

760 + )

761

762 if "temperature" not in sampling_params:

763 sampling_params["temperature"] = 0.0 # Default to greedy decoding.

@@ -709,15 +781,14 @@ def main(args: argparse.Namespace):

781 disable_tqdm=args.disable_tqdm,

782 profile=args.profile,

783 selected_percentile_metrics=args.percentile_metrics.split(","),

712 - selected_percentiles=[

713 - float(p) for p in args.metric_percentiles.split(",")

714 - ],

784 + selected_percentiles=[float(p) for p in args.metric_percentiles.split(",")],

785 ignore_eos=args.ignore_eos,

786 goodput_config_dict=goodput_config_dict,

787 max_concurrency=args.max_concurrency,

788 lora_modules=args.lora_modules,

789 extra_body=sampling_params,

720 - ))

790 + )

791 + )

792

793 # Save config and results to json

794 if args.save_result or args.append_result:

@@ -742,8 +813,9 @@ def main(args: argparse.Namespace):

813 "Invalid metadata format. Please use KEY=VALUE format."

814 )

815 # Traffic

745 - result_json["request_rate"] = (args.request_rate if args.request_rate

746 - < float("inf") else "inf")

816 + result_json["request_rate"] = (

817 + args.request_rate if args.request_rate < float("inf") else "inf"

818 + )

819 result_json["burstiness"] = args.burstiness

820 result_json["max_concurrency"] = args.max_concurrency

821

@@ -753,24 +825,31 @@ def main(args: argparse.Namespace):

825 if not args.save_detailed:

826 # Remove fields with too many data points

827 for field in [

756 - "input_lens", "output_lens", "ttfts", "itls",

757 - "generated_texts", "errors"

828 + "input_lens",

829 + "output_lens",

830 + "ttfts",

831 + "itls",

832 + "generated_texts",

833 + "errors",

834 ]:

835 if field in result_json:

836 del result_json[field]

837

838 # Save to file

839 base_model_id = model_id.split("/")[-1]

764 - max_concurrency_str = (f"-concurrency{args.max_concurrency}"

765 - if args.max_concurrency is not None else "")

766 - file_name = f"{backend}-{args.request_rate}qps{max_concurrency_str}-{base_model_id}-{current_dt}.json" #noqa

840 + max_concurrency_str = (

841 + f"-concurrency{args.max_concurrency}"

842 + if args.max_concurrency is not None

843 + else ""

844 + )

845 + file_name = f"{backend}-{args.request_rate}qps{max_concurrency_str}-{base_model_id}-{current_dt}.json" # noqa

846 if args.result_filename:

847 file_name = args.result_filename

848 if args.result_dir:

849 file_name = os.path.join(args.result_dir, file_name)

771 - with open(file_name,

772 - mode="a+" if args.append_result else "w",

773 - encoding='utf-8') as outfile:

850 + with open(

851 + file_name, mode="a+" if args.append_result else "w", encoding="utf-8"

852 + ) as outfile:

853 # Append a newline.

854 if args.append_result and outfile.tell() != 0:

855 outfile.write("\n")

@@ -780,7 +859,8 @@ def main(args: argparse.Namespace):

859

860 if __name__ == "__main__":

861 parser = FlexibleArgumentParser(

783 - description="Benchmark the online serving throughput.")

862 + description="Benchmark the online serving throughput."

863 + )

864 parser.add_argument(

865 "--backend",

866 type=str,

@@ -809,11 +889,13 @@ if __name__ == "__main__":

889 choices=["sharegpt", "burstgpt", "sonnet", "random", "hf"],

890 help="Name of the dataset to benchmark on.",

891 )

812 - parser.add_argument("--dataset-path",

813 - type=str,

814 - default=None,

815 - help="Path to the sharegpt/sonnet dataset. "

816 - "Or the huggingface dataset ID if using HF dataset.")

892 + parser.add_argument(

893 + "--dataset-path",

894 + type=str,

895 + default=None,

896 + help="Path to the sharegpt/sonnet dataset. "

897 + "Or the huggingface dataset ID if using HF dataset.",

898 + )

899 parser.add_argument(

900 "--max-concurrency",

901 type=int,

@@ -825,7 +907,8 @@ if __name__ == "__main__":

907 "initiated, this argument will control how many are actually allowed "

908 "to execute at a time. This means that when used in combination, the "

909 "actual request rate may be lower than specified with --request-rate, "

828 - "if the server is not processing requests fast enough to keep up.")

910 + "if the server is not processing requests fast enough to keep up.",

911 + )

912

913 parser.add_argument(

914 "--model",

@@ -836,8 +919,7 @@ if __name__ == "__main__":

919 parser.add_argument(

920 "--tokenizer",

921 type=str,

839 - help=

840 - "Name or path of the tokenizer, if not using the default tokenizer.", # noqa: E501

922 + help="Name or path of the tokenizer, if not using the default tokenizer.", # noqa: E501

923 )

924 parser.add_argument("--use-beam-search", action="store_true")

925 parser.add_argument(

@@ -850,11 +932,13 @@ if __name__ == "__main__":

932 "--logprobs",

933 type=int,

934 default=None,

853 - help=("Number of logprobs-per-token to compute & return as part of "

854 - "the request. If unspecified, then either (1) if beam search "

855 - "is disabled, no logprobs are computed & a single dummy "

856 - "logprob is returned for each token; or (2) if beam search "

857 - "is enabled 1 logprob per token is computed"),

935 + help=(

936 + "Number of logprobs-per-token to compute & return as part of "

937 + "the request. If unspecified, then either (1) if beam search "

938 + "is disabled, no logprobs are computed & a single dummy "

939 + "logprob is returned for each token; or (2) if beam search "

940 + "is enabled 1 logprob per token is computed"

941 + ),

942 )

943 parser.add_argument(

944 "--request-rate",

@@ -938,35 +1022,38 @@ if __name__ == "__main__":

1022 "--ignore-eos",

1023 action="store_true",

1024 help="Set ignore_eos flag when sending the benchmark request."

941 - "Warning: ignore_eos is not supported in deepspeed_mii and tgi.")

1025 + "Warning: ignore_eos is not supported in deepspeed_mii and tgi.",

1026 + )

1027 parser.add_argument(

1028 "--percentile-metrics",

1029 type=str,

1030 default="ttft,tpot,itl",

1031 help="Comma-separated list of selected metrics to report percentils. "

1032 "This argument specifies the metrics to report percentiles. "

948 - "Allowed metric names are \"ttft\", \"tpot\", \"itl\", \"e2el\". "

949 - "Default value is \"ttft,tpot,itl\".")

1033 + 'Allowed metric names are "ttft", "tpot", "itl", "e2el". '

1034 + 'Default value is "ttft,tpot,itl".',

1035 + )

1036 parser.add_argument(

1037 "--metric-percentiles",

1038 type=str,

1039 default="99",

1040 help="Comma-separated list of percentiles for selected metrics. "

955 - "To report 25-th, 50-th, and 75-th percentiles, use \"25,50,75\". "

956 - "Default value is \"99\". "

957 - "Use \"--percentile-metrics\" to select metrics.",

1041 + 'To report 25-th, 50-th, and 75-th percentiles, use "25,50,75". '

1042 + 'Default value is "99". '

1043 + 'Use "--percentile-metrics" to select metrics.',

1044 )

1045 parser.add_argument(

1046 "--goodput",

1047 nargs="+",

1048 required=False,

963 - help="Specify service level objectives for goodput as \"KEY:VALUE\" "

1049 + help='Specify service level objectives for goodput as "KEY:VALUE" '

1050 "pairs, where the key is a metric name, and the value is in "

965 - "milliseconds. Multiple \"KEY:VALUE\" pairs can be provided, "

1051 + 'milliseconds. Multiple "KEY:VALUE" pairs can be provided, '

1052 "separated by spaces. Allowed request level metric names are "

967 - "\"ttft\", \"tpot\", \"e2el\". For more context on the definition of "

1053 + '"ttft", "tpot", "e2el". For more context on the definition of '

1054 "goodput, refer to DistServe paper: https://arxiv.org/pdf/2401.09670 "

969 - "and the blog: https://hao-ai-lab.github.io/blogs/distserve")

1055 + "and the blog: https://hao-ai-lab.github.io/blogs/distserve",

1056 + )

1057

1058 # group for dataset specific arguments

1059 sonnet_group = parser.add_argument_group("sonnet dataset options")

@@ -974,22 +1061,19 @@ if __name__ == "__main__":

1061 "--sonnet-input-len",

1062 type=int,

1063 default=550,

977 - help=

978 - "Number of input tokens per request, used only for sonnet dataset.",

1064 + help="Number of input tokens per request, used only for sonnet dataset.",

1065 )

1066 sonnet_group.add_argument(

1067 "--sonnet-output-len",

1068 type=int,

1069 default=150,

984 - help=

985 - "Number of output tokens per request, used only for sonnet dataset.",

1070 + help="Number of output tokens per request, used only for sonnet dataset.",

1071 )

1072 sonnet_group.add_argument(

1073 "--sonnet-prefix-len",

1074 type=int,

1075 default=200,

991 - help=

992 - "Number of prefix tokens per request, used only for sonnet dataset.",

1076 + help="Number of prefix tokens per request, used only for sonnet dataset.",

1077 )

1078

1079 sharegpt_group = parser.add_argument_group("sharegpt dataset options")

@@ -998,22 +1082,21 @@ if __name__ == "__main__":

1082 type=int,

1083 default=None,

1084 help="Output length for each request. Overrides the output length "

1001 - "from the ShareGPT dataset.")

1085 + "from the ShareGPT dataset.",

1086 + )

1087

1088 random_group = parser.add_argument_group("random dataset options")

1089 random_group.add_argument(

1090 "--random-input-len",

1091 type=int,

1092 default=1024,

1008 - help=

1009 - "Number of input tokens per request, used only for random sampling.",

1093 + help="Number of input tokens per request, used only for random sampling.",

1094 )

1095 random_group.add_argument(

1096 "--random-output-len",

1097 type=int,

1098 default=128,

1015 - help=

1016 - "Number of output tokens per request, used only for random sampling.",

1099 + help="Number of output tokens per request, used only for random sampling.",

1100 )

1101 random_group.add_argument(

1102 "--random-range-ratio",

@@ -1028,23 +1111,23 @@ if __name__ == "__main__":

1111 "--random-prefix-len",

1112 type=int,

1113 default=0,

1031 - help=("Number of fixed prefix tokens before the random context "

1032 - "in a request. "

1033 - "The total input length is the sum of `random-prefix-len` and "

1034 - "a random "

1035 - "context length sampled from [input_len * (1 - range_ratio), "

1036 - "input_len * (1 + range_ratio)]."),

1114 + help=(

1115 + "Number of fixed prefix tokens before the random context "

1116 + "in a request. "

1117 + "The total input length is the sum of `random-prefix-len` and "

1118 + "a random "

1119 + "context length sampled from [input_len * (1 - range_ratio), "

1120 + "input_len * (1 + range_ratio)]."

1121 + ),

1122 )

1123

1124 hf_group = parser.add_argument_group("hf dataset options")

1040 - hf_group.add_argument("--hf-subset",

1041 - type=str,

1042 - default=None,

1043 - help="Subset of the HF dataset.")

1044 - hf_group.add_argument("--hf-split",

1045 - type=str,

1046 - default=None,

1047 - help="Split of the HF dataset.")

1125 + hf_group.add_argument(

1126 + "--hf-subset", type=str, default=None, help="Subset of the HF dataset."

1127 + )

1128 + hf_group.add_argument(

1129 + "--hf-split", type=str, default=None, help="Split of the HF dataset."

1130 + )

1131 hf_group.add_argument(

1132 "--hf-output-len",

1133 type=int,

@@ -1058,52 +1141,58 @@ if __name__ == "__main__":

1141 "--top-p",

1142 type=float,

1143 default=None,

1061 - help="Top-p sampling parameter. Only has effect on openai-compatible "

1062 - "backends.")

1144 + help="Top-p sampling parameter. Only has effect on openai-compatible backends.",

1145 + )

1146 sampling_group.add_argument(

1147 "--top-k",

1148 type=int,

1149 default=None,

1067 - help="Top-k sampling parameter. Only has effect on openai-compatible "

1068 - "backends.")

1150 + help="Top-k sampling parameter. Only has effect on openai-compatible backends.",

1151 + )

1152 sampling_group.add_argument(

1153 "--min-p",

1154 type=float,

1155 default=None,

1073 - help="Min-p sampling parameter. Only has effect on openai-compatible "

1074 - "backends.")

1156 + help="Min-p sampling parameter. Only has effect on openai-compatible backends.",

1157 + )

1158 sampling_group.add_argument(

1159 "--temperature",

1160 type=float,

1161 default=None,

1162 help="Temperature sampling parameter. Only has effect on "

1163 "openai-compatible backends. If not specified, default to greedy "

1081 - "decoding (i.e. temperature==0.0).")

1164 + "decoding (i.e. temperature==0.0).",

1165 + )

1166

1167 parser.add_argument(

1084 - '--tokenizer-mode',

1168 + "--tokenizer-mode",

1169 type=str,

1170 default="auto",

1087 - choices=['auto', 'slow', 'mistral', 'custom'],

1171 + choices=["auto", "slow", "mistral", "custom"],

1172 help='The tokenizer mode.\n\n* "auto" will use the '

1173 'fast tokenizer if available.\n* "slow" will '

1090 - 'always use the slow tokenizer. \n* '

1174 + "always use the slow tokenizer. \n* "

1175 '"mistral" will always use the `mistral_common` tokenizer. \n*'

1092 - '"custom" will use --tokenizer to select the preregistered tokenizer.')

1093 -

1094 - parser.add_argument("--served-model-name",

1095 - type=str,

1096 - default=None,

1097 - help="The model name used in the API. "

1098 - "If not specified, the model name will be the "

1099 - "same as the ``--model`` argument. ")

1100 -

1101 - parser.add_argument("--lora-modules",

1102 - nargs='+',

1103 - default=None,

1104 - help="A subset of LoRA module names passed in when "

1105 - "launching the server. For each request, the "

1106 - "script chooses a LoRA module at random.")

1176 + '"custom" will use --tokenizer to select the preregistered tokenizer.',

1177 + )

1178 +

1179 + parser.add_argument(

1180 + "--served-model-name",

1181 + type=str,

1182 + default=None,

1183 + help="The model name used in the API. "

1184 + "If not specified, the model name will be the "

1185 + "same as the ``--model`` argument. ",

1186 + )

1187 +

1188 + parser.add_argument(

1189 + "--lora-modules",

1190 + nargs="+",

1191 + default=None,

1192 + help="A subset of LoRA module names passed in when "

1193 + "launching the server. For each request, the "

1194 + "script chooses a LoRA module at random.",

1195 + )

1196

1197 args = parser.parse_args()

1198

benchmarks/benchmark_serving_structured_output.py

+263 -226

@@ -19,6 +19,7 @@ On the client side, run:

19 --endpoint /generate_stream

20 to the end of the command above.

21 """

22 +

23 import argparse

24 import asyncio

25 import copy

@@ -36,11 +37,15 @@ from typing import Optional

37 import datasets

38 import numpy as np

39 import pandas as pd

39 -from backend_request_func import (ASYNC_REQUEST_FUNCS, RequestFuncInput,

40 - RequestFuncOutput)

40 from tqdm.asyncio import tqdm

41 from transformers import PreTrainedTokenizerBase

43 +from backend_request_func import (

44 + ASYNC_REQUEST_FUNCS,

45 + RequestFuncInput,

46 + RequestFuncOutput,

47 +)

48 +

49 try:

50 from vllm.transformers_utils.tokenizer import get_tokenizer

51 except ImportError:

@@ -52,7 +57,8 @@ except ImportError:

57 from argparse import ArgumentParser as FlexibleArgumentParser

59 from vllm.v1.structured_output.backend_xgrammar import (

55 - has_xgrammar_unsupported_json_features)

60 + has_xgrammar_unsupported_json_features,

61 +)

63 MILLISECONDS_TO_SECONDS_CONVERSION = 1000

@@ -98,6 +104,7 @@ class SampleRequest:

104 prompt_len: The length of the prompt in tokens.

105 expected_output_len: The expected length of the output in tokens.

106 """

107 +

108 prompt: str

109 prompt_len: int

110 expected_output_len: int

@@ -106,32 +113,28 @@ class SampleRequest:

113 completion: str = None

114

115

109 -def sample_requests(tokenizer: PreTrainedTokenizerBase,

110 - args: argparse.Namespace) -> list[SampleRequest]:

111 - if args.dataset == 'json' or args.dataset == 'json-unique':

116 +def sample_requests(

117 + tokenizer: PreTrainedTokenizerBase, args: argparse.Namespace

118 +) -> list[SampleRequest]:

119 + if args.dataset == "json" or args.dataset == "json-unique":

120 if args.json_schema_path is None:

121 dir_path = os.path.dirname(os.path.realpath(__file__))

114 - args.json_schema_path = os.path.join(dir_path,

115 - "structured_schemas",

116 - "structured_schema_1.json")

122 + args.json_schema_path = os.path.join(

123 + dir_path, "structured_schemas", "structured_schema_1.json"

124 + )

125 json_schemas = []

126 with open(args.json_schema_path) as f:

127 schema = json.load(f)

128

121 - if args.dataset == 'json-unique':

122 - json_schemas = [

123 - copy.deepcopy(schema) for _ in range(args.num_prompts)

124 - ]

129 + if args.dataset == "json-unique":

130 + json_schemas = [copy.deepcopy(schema) for _ in range(args.num_prompts)]

131 for i in range(len(json_schemas)):

132 if "properties" not in json_schemas[i]:

133 json_schemas[i]["properties"] = {}

128 - json_schemas[i]["properties"][

129 - f"__optional_field_{uuid.uuid4()}"] = {

130 - "type":

131 - "string",

132 - "description":

133 - "An unique optional field to avoid cached schemas"

134 - }

134 + json_schemas[i]["properties"][f"__optional_field_{uuid.uuid4()}"] = {

135 + "type": "string",

136 + "description": "An unique optional field to avoid cached schemas",

137 + }

138 else:

139 json_schemas = [schema] * args.num_prompts

140

@@ -142,11 +145,13 @@ def sample_requests(tokenizer: PreTrainedTokenizerBase,

145 return json_schemas[index % len(json_schemas)]

146

147 requests = [

145 - SampleRequest(prompt=gen_prompt(i),

146 - prompt_len=len(tokenizer(gen_prompt(i)).input_ids),

147 - expected_output_len=args.output_len,

148 - schema=get_schema(i),

149 - structure_type=args.structure_type)

148 + SampleRequest(

149 + prompt=gen_prompt(i),

150 + prompt_len=len(tokenizer(gen_prompt(i)).input_ids),

151 + expected_output_len=args.output_len,

152 + schema=get_schema(i),

153 + structure_type=args.structure_type,

154 + )

155 for i in range(args.num_prompts)

156 ]

157

@@ -170,11 +175,13 @@ def sample_requests(tokenizer: PreTrainedTokenizerBase,

175 input_len = len(tokenizer(prompt).input_ids)

176 print(f"Input length of the prompt: {input_len} tokens")

177 requests = [

173 - SampleRequest(prompt=prompt,

174 - prompt_len=input_len,

175 - expected_output_len=args.output_len,

176 - schema=schema,

177 - structure_type=args.structure_type)

178 + SampleRequest(

179 + prompt=prompt,

180 + prompt_len=input_len,

181 + expected_output_len=args.output_len,

182 + schema=schema,

183 + structure_type=args.structure_type,

184 + )

185 for _ in range(args.num_prompts)

186 ]

187

@@ -188,11 +195,13 @@ def sample_requests(tokenizer: PreTrainedTokenizerBase,

195 input_len = len(tokenizer(prompt).input_ids)

196 print(f"Input length of the prompt: {input_len} tokens")

197 requests = [

191 - SampleRequest(prompt=prompt,

192 - prompt_len=input_len,

193 - expected_output_len=args.output_len,

194 - schema=regex,

195 - structure_type=args.structure_type)

198 + SampleRequest(

199 + prompt=prompt,

200 + prompt_len=input_len,

201 + expected_output_len=args.output_len,

202 + schema=regex,

203 + structure_type=args.structure_type,

204 + )

205 for _ in range(args.num_prompts)

206 ]

207

@@ -203,48 +212,55 @@ def sample_requests(tokenizer: PreTrainedTokenizerBase,

212 input_len = len(tokenizer(prompt).input_ids)

213 print(f"Input length of the prompt: {input_len} tokens")

214 requests = [

206 - SampleRequest(prompt=prompt,

207 - prompt_len=input_len,

208 - expected_output_len=args.output_len,

209 - schema=choice,

210 - structure_type=args.structure_type)

215 + SampleRequest(

216 + prompt=prompt,

217 + prompt_len=input_len,

218 + expected_output_len=args.output_len,

219 + schema=choice,

220 + structure_type=args.structure_type,

221 + )

222 for _ in range(args.num_prompts)

223 ]

224

225 elif args.dataset == "xgrammar_bench":

226 requests: list[SampleRequest] = []

216 - dataset = datasets.load_dataset("NousResearch/json-mode-eval",

217 - split="train")

227 + dataset = datasets.load_dataset("NousResearch/json-mode-eval", split="train")

228 full_dataset_len = len(dataset)

229

230 def _filter_func(item):

231 import json

232 +

233 schema = json.loads(item["schema"])

234 return not has_xgrammar_unsupported_json_features(schema)

235

236 dataset = dataset.filter(_filter_func)

237 num_filtered_out = full_dataset_len - len(dataset)

227 - print(f"dataset has {len(dataset)} entries after filtering "

228 - f"out {num_filtered_out} entries with unsupported features")

238 + print(

239 + f"dataset has {len(dataset)} entries after filtering "

240 + f"out {num_filtered_out} entries with unsupported features"

241 + )

242 len_dataset = len(dataset)

243 for data_point_idx in range(args.num_prompts):

244 idx = data_point_idx

245 while idx >= len_dataset:

246 idx -= len_dataset

247 schema = dataset["schema"][idx]

235 - prompt = tokenizer.apply_chat_template(dataset["prompt"][idx],

236 - tokenize=False,

237 - add_generation_prompt=True)

248 + prompt = tokenizer.apply_chat_template(

249 + dataset["prompt"][idx], tokenize=False, add_generation_prompt=True

250 + )

251 input_len = len(tokenizer(prompt).input_ids)

252 completion = dataset["completion"][idx]

253

254 requests.append(

242 - SampleRequest(prompt=prompt,

243 - prompt_len=input_len,

244 - expected_output_len=args.output_len,

245 - schema=schema,

246 - structure_type=args.structure_type,

247 - completion=completion))

255 + SampleRequest(

256 + prompt=prompt,

257 + prompt_len=input_len,

258 + expected_output_len=args.output_len,

259 + schema=schema,

260 + structure_type=args.structure_type,

261 + completion=completion,

262 + )

263 + )

264

265 return requests

266

@@ -276,7 +292,8 @@ async def get_request(

292

293 # Calculate scale parameter theta to maintain the desired request_rate.

294 assert burstiness > 0, (

279 - f"A positive burstiness factor is expected, but given {burstiness}.")

295 + f"A positive burstiness factor is expected, but given {burstiness}."

296 + )

297 theta = 1.0 / (request_rate * burstiness)

298

299 for i, request in enumerate(input_requests):

@@ -318,8 +335,8 @@ def calculate_metrics(

335 # multiple output tokens may be bundled together

336 # Note : this may inflate the output token count slightly

337 output_len = len(

321 - tokenizer(outputs[i].generated_text,

322 - add_special_tokens=False).input_ids)

338 + tokenizer(outputs[i].generated_text, add_special_tokens=False).input_ids

339 + )

340 actual_output_lens.append(output_len)

341 total_input += input_requests[i].prompt_len

342 tpot = 0

@@ -343,16 +360,19 @@ def calculate_metrics(

360

361 if "ttft" in goodput_config_dict:

362 valid_metrics.append(ttfts)

346 - slo_values.append(goodput_config_dict["ttft"] /

347 - MILLISECONDS_TO_SECONDS_CONVERSION)

363 + slo_values.append(

364 + goodput_config_dict["ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION

365 + )

366 if "tpot" in goodput_config_dict:

367 valid_metrics.append(all_tpots)

350 - slo_values.append(goodput_config_dict["tpot"] /

351 - MILLISECONDS_TO_SECONDS_CONVERSION)

368 + slo_values.append(

369 + goodput_config_dict["tpot"] / MILLISECONDS_TO_SECONDS_CONVERSION

370 + )

371 if "e2el" in goodput_config_dict:

372 valid_metrics.append(e2els)

354 - slo_values.append(goodput_config_dict["e2el"] /

355 - MILLISECONDS_TO_SECONDS_CONVERSION)

373 + slo_values.append(

374 + goodput_config_dict["e2el"] / MILLISECONDS_TO_SECONDS_CONVERSION

375 + )

376

377 for req_metric in zip(*valid_metrics):

378 is_good_req = all([s >= r for s, r in zip(slo_values, req_metric)])

@@ -363,7 +383,8 @@ def calculate_metrics(

383 warnings.warn(

384 "All requests failed. This is likely due to a misconfiguration "

385 "on the benchmark arguments.",

366 - stacklevel=2)

386 + stacklevel=2,

387 + )

388 metrics = BenchmarkMetrics(

389 completed=completed,

390 total_input=total_input,

@@ -372,27 +393,31 @@ def calculate_metrics(

393 request_goodput=good_completed / dur_s,

394 output_throughput=sum(actual_output_lens) / dur_s,

395 total_token_throughput=(total_input + sum(actual_output_lens)) / dur_s,

375 - mean_ttft_ms=np.mean(ttfts or 0) *

376 - 1000, # ttfts is empty if streaming is not supported by backend

396 + mean_ttft_ms=np.mean(ttfts or 0)

397 + * 1000, # ttfts is empty if streaming is not supported by backend

398 std_ttft_ms=np.std(ttfts or 0) * 1000,

399 median_ttft_ms=np.median(ttfts or 0) * 1000,

379 - percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000)

380 - for p in selected_percentiles],

400 + percentiles_ttft_ms=[

401 + (p, np.percentile(ttfts or 0, p) * 1000) for p in selected_percentiles

402 + ],

403 mean_tpot_ms=np.mean(tpots or 0) * 1000,

404 std_tpot_ms=np.std(tpots or 0) * 1000,

405 median_tpot_ms=np.median(tpots or 0) * 1000,

384 - percentiles_tpot_ms=[(p, np.percentile(tpots or 0, p) * 1000)

385 - for p in selected_percentiles],

406 + percentiles_tpot_ms=[

407 + (p, np.percentile(tpots or 0, p) * 1000) for p in selected_percentiles

408 + ],

409 mean_itl_ms=np.mean(itls or 0) * 1000,

410 std_itl_ms=np.std(itls or 0) * 1000,

411 median_itl_ms=np.median(itls or 0) * 1000,

389 - percentiles_itl_ms=[(p, np.percentile(itls or 0, p) * 1000)

390 - for p in selected_percentiles],

412 + percentiles_itl_ms=[

413 + (p, np.percentile(itls or 0, p) * 1000) for p in selected_percentiles

414 + ],

415 mean_e2el_ms=np.mean(e2els or 0) * 1000,

416 std_e2el_ms=np.std(e2els or 0) * 1000,

417 median_e2el_ms=np.median(e2els or 0) * 1000,

394 - percentiles_e2el_ms=[(p, np.percentile(e2els or 0, p) * 1000)

395 - for p in selected_percentiles],

418 + percentiles_e2el_ms=[

419 + (p, np.percentile(e2els or 0, p) * 1000) for p in selected_percentiles

420 + ],

421 )

422

423 return metrics, actual_output_lens

@@ -429,12 +454,13 @@ async def benchmark(

454

455 print("Starting initial single prompt test run...")

456 structured_output_req_idx = random.sample(

432 - range(len(input_requests)),

433 - int(len(input_requests) * structured_output_ratio))

457 + range(len(input_requests)), int(len(input_requests) * structured_output_ratio)

458 + )

459

460 test_request = input_requests[0]

436 - test_req_extra_body = (prepare_extra_body(test_request)

437 - if 0 in structured_output_req_idx else None)

461 + test_req_extra_body = (

462 + prepare_extra_body(test_request) if 0 in structured_output_req_idx else None

463 + )

464 test_input = RequestFuncInput(

465 model=model_id,

466 prompt=test_request.prompt,

@@ -448,7 +474,8 @@ async def benchmark(

474 if not test_output.success:

475 raise ValueError(

476 "Initial test run failed - Please make sure benchmark arguments "

451 - f"are correctly specified. Error: {test_output.error}")

477 + f"are correctly specified. Error: {test_output.error}"

478 + )

479 else:

480 print("Initial test run completed. Starting main benchmark run...")

481

@@ -467,10 +494,7 @@ async def benchmark(

494 if profile_output.success:

495 print("Profiler started")

496

470 - if burstiness == 1.0:

471 - distribution = "Poisson process"

472 - else:

473 - distribution = "Gamma distribution"

497 + distribution = "Poisson process" if burstiness == 1.0 else "Gamma distribution"

498

499 print(f"Traffic request rate: {request_rate}")

500 print(f"Burstiness factor: {burstiness} ({distribution})")

@@ -482,24 +506,21 @@ async def benchmark(

506 # and it will simplify the code in limited_request_func.

507 # semaphore = (asyncio.Semaphore(max_concurrency)

508 # if max_concurrency else contextlib.nullcontext())

485 - semaphore = (asyncio.Semaphore(max_concurrency)

486 - if max_concurrency else None)

509 + semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None

510

511 async def limited_request_func(request_func_input, pbar):

512 if semaphore is None:

490 - return await request_func(request_func_input=request_func_input,

491 - pbar=pbar)

513 + return await request_func(request_func_input=request_func_input, pbar=pbar)

514 async with semaphore:

493 - return await request_func(request_func_input=request_func_input,

494 - pbar=pbar)

515 + return await request_func(request_func_input=request_func_input, pbar=pbar)

516

517 benchmark_start_time = time.perf_counter()

518 tasks: list[asyncio.Task] = []

519 expected: list[str] = []

499 - async for i, request in get_request(input_requests, request_rate,

500 - burstiness):

501 - extra_body = prepare_extra_body(

502 - request) if i in structured_output_req_idx else None

520 + async for i, request in get_request(input_requests, request_rate, burstiness):

521 + extra_body = (

522 + prepare_extra_body(request) if i in structured_output_req_idx else None

523 + )

524 request_func_input = RequestFuncInput(

525 model=model_id,

526 prompt=request.prompt,

@@ -512,8 +533,9 @@ async def benchmark(

533 expected.append(request.completion)

534 tasks.append(

535 asyncio.create_task(

515 - limited_request_func(request_func_input=request_func_input,

516 - pbar=pbar)))

536 + limited_request_func(request_func_input=request_func_input, pbar=pbar)

537 + )

538 + )

539 outputs: list[RequestFuncOutput] = await asyncio.gather(*tasks)

540

541 if profile:

@@ -545,54 +567,58 @@ async def benchmark(

567 goodput_config_dict=goodput_config_dict,

568 )

569

548 - print("{s:{c}^{n}}".format(s=' Serving Benchmark Result ', n=50, c='='))

570 + print("{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="="))

571 print("{:<40} {:<10}".format("Successful requests:", metrics.completed))

550 - print("{:<40} {:<10.2f}".format("Benchmark duration (s):",

551 - benchmark_duration))

572 + print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration))

573 print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input))

553 - print("{:<40} {:<10}".format("Total generated tokens:",

554 - metrics.total_output))

555 - print("{:<40} {:<10.2f}".format("Request throughput (req/s):",

556 - metrics.request_throughput))

574 + print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output))

575 + print(

576 + "{:<40} {:<10.2f}".format(

577 + "Request throughput (req/s):", metrics.request_throughput

578 + )

579 + )

580 if goodput_config_dict:

558 - print("{:<40} {:<10.2f}".format("Request goodput (req/s):",

559 - metrics.request_goodput))

560 - print("{:<40} {:<10.2f}".format("Output token throughput (tok/s):",

561 - metrics.output_throughput))

562 - print("{:<40} {:<10.2f}".format("Total Token throughput (tok/s):",

563 - metrics.total_token_throughput))

581 + print(

582 + "{:<40} {:<10.2f}".format(

583 + "Request goodput (req/s):", metrics.request_goodput

584 + )

585 + )

586 + print(

587 + "{:<40} {:<10.2f}".format(

588 + "Output token throughput (tok/s):", metrics.output_throughput

589 + )

590 + )

591 + print(

592 + "{:<40} {:<10.2f}".format(

593 + "Total Token throughput (tok/s):", metrics.total_token_throughput

594 + )

595 + )

596

597 result = {

566 - "duration":

567 - benchmark_duration,

568 - "completed":

569 - metrics.completed,

570 - "total_input_tokens":

571 - metrics.total_input,

572 - "total_output_tokens":

573 - metrics.total_output,

574 - "request_throughput":

575 - metrics.request_throughput,

576 - "output_throughput":

577 - metrics.output_throughput,

578 - "total_token_throughput":

579 - metrics.total_token_throughput,

580 - "ttft_description":

581 - pd.Series([output.ttft for output in outputs]).describe().to_dict(),

582 - "tpot_description":

583 - pd.Series([output.tpot for output in outputs]).describe().to_dict(),

598 + "duration": benchmark_duration,

599 + "completed": metrics.completed,

600 + "total_input_tokens": metrics.total_input,

601 + "total_output_tokens": metrics.total_output,

602 + "request_throughput": metrics.request_throughput,

603 + "output_throughput": metrics.output_throughput,

604 + "total_token_throughput": metrics.total_token_throughput,

605 + "ttft_description": pd.Series([output.ttft for output in outputs])

606 + .describe()

607 + .to_dict(),

608 + "tpot_description": pd.Series([output.tpot for output in outputs])

609 + .describe()

610 + .to_dict(),

611 "input_lens": [output.prompt_len for output in outputs],

585 - "output_lens":

586 - actual_output_lens,

612 + "output_lens": actual_output_lens,

613 "ttfts": [output.ttft for output in outputs],

614 "itls": [output.itl for output in outputs],

615 "errors": [output.error for output in outputs],

616 }

617

592 - ret = [{

593 - 'generated': output.generated_text,

594 - 'expected': gt

595 - } for output, gt in zip(outputs, expected)]

618 + ret = [

619 + {"generated": output.generated_text, "expected": gt}

620 + for output, gt in zip(outputs, expected)

621 + ]

622

623 def process_one_metric(

624 # E.g., "ttft"

@@ -606,29 +632,35 @@ async def benchmark(

632 # metric.

633 if metric_attribute_name not in selected_percentile_metrics:

634 return

609 - print("{s:{c}^{n}}".format(s=metric_header, n=50, c='-'))

610 - print("{:<40} {:<10.2f}".format(

611 - f"Mean {metric_name} (ms):",

612 - getattr(metrics, f"mean_{metric_attribute_name}_ms")))

613 - print("{:<40} {:<10.2f}".format(

614 - f"Median {metric_name} (ms):",

615 - getattr(metrics, f"median_{metric_attribute_name}_ms")))

635 + print("{s:{c}^{n}}".format(s=metric_header, n=50, c="-"))

636 + print(

637 + "{:<40} {:<10.2f}".format(

638 + f"Mean {metric_name} (ms):",

639 + getattr(metrics, f"mean_{metric_attribute_name}_ms"),

640 + )

641 + )

642 + print(

643 + "{:<40} {:<10.2f}".format(

644 + f"Median {metric_name} (ms):",

645 + getattr(metrics, f"median_{metric_attribute_name}_ms"),

646 + )

647 + )

648 result[f"mean_{metric_attribute_name}_ms"] = getattr(

617 - metrics, f"mean_{metric_attribute_name}_ms")

649 + metrics, f"mean_{metric_attribute_name}_ms"

650 + )

651 result[f"median_{metric_attribute_name}_ms"] = getattr(

619 - metrics, f"median_{metric_attribute_name}_ms")

652 + metrics, f"median_{metric_attribute_name}_ms"

653 + )

654 result[f"std_{metric_attribute_name}_ms"] = getattr(

621 - metrics, f"std_{metric_attribute_name}_ms")

622 - for p, value in getattr(metrics,

623 - f"percentiles_{metric_attribute_name}_ms"):

655 + metrics, f"std_{metric_attribute_name}_ms"

656 + )

657 + for p, value in getattr(metrics, f"percentiles_{metric_attribute_name}_ms"):

658 p_word = str(int(p)) if int(p) == p else str(p)

625 - print("{:<40} {:<10.2f}".format(f"P{p_word} {metric_name} (ms):",

626 - value))

659 + print("{:<40} {:<10.2f}".format(f"P{p_word} {metric_name} (ms):", value))

660 result[f"p{p_word}_{metric_attribute_name}_ms"] = value

661

662 process_one_metric("ttft", "TTFT", "Time to First Token")

630 - process_one_metric("tpot", "TPOT",

631 - "Time per Output Token (excl. 1st token)")

663 + process_one_metric("tpot", "TPOT", "Time per Output Token (excl. 1st token)")

664 process_one_metric("itl", "ITL", "Inter-token Latency")

665 process_one_metric("e2el", "E2EL", "End-to-end Latency")

666

@@ -638,13 +670,13 @@ async def benchmark(

670

671

672 def evaluate(ret, args):

641 -

673 def _eval_correctness_json(expected, actual):

674 # extract json string from string using regex

675 import re

645 - actual = actual.replace('\n', '').replace(' ', '').strip()

676 +

677 + actual = actual.replace("\n", "").replace(" ", "").strip()

678 try:

647 - actual = re.search(r'\{.*\}', actual).group()

679 + actual = re.search(r"\{.*\}", actual).group()

680 actual = json.loads(actual)

681 except Exception:

682 return False

@@ -656,28 +688,32 @@ def evaluate(ret, args):

688

689 def _eval_correctness_regex(expected, actual):

690 import re

691 +

692 return re.match(args.regex, actual) is not None

693

694 def _eval_correctness(expected, actual):

662 - if args.structure_type == 'guided_json':

695 + if args.structure_type == "guided_json":

696 return _eval_correctness_json(expected, actual)

664 - elif args.structure_type == 'guided_regex':

697 + elif args.structure_type == "guided_regex":

698 return _eval_correctness_regex(expected, actual)

666 - elif args.structure_type == 'guided_choice':

699 + elif args.structure_type == "guided_choice":

700 return _eval_correctness_choice(expected, actual)

701 else:

702 return None

703

704 scores = []

705 for res in ret:

673 - score = _eval_correctness(res['expected'], res['generated'])

674 - res['correctness'] = score

706 + score = _eval_correctness(res["expected"], res["generated"])

707 + res["correctness"] = score

708 scores.append(score)

709

710 not_none_scores = [score for score in scores if score is not None]

711

679 - return (sum(not_none_scores) / len(not_none_scores) *

680 - 100) if len(not_none_scores) > 0 else None

712 + return (

713 + (sum(not_none_scores) / len(not_none_scores) * 100)

714 + if len(not_none_scores) > 0

715 + else None

716 + )

717

718

719 def parse_goodput(slo_pairs):

@@ -689,9 +725,10 @@ def parse_goodput(slo_pairs):

725 except ValueError as err:

726 raise argparse.ArgumentTypeError(

727 "Invalid format found for service level objectives. "

692 - "Specify service level objectives for goodput as \"KEY:VALUE\" "

728 + 'Specify service level objectives for goodput as "KEY:VALUE" '

729 "pairs, where the key is a metric name, and the value is a "

694 - "number in milliseconds.") from err

730 + "number in milliseconds."

731 + ) from err

732 return goodput_config_dict

733

734

@@ -705,12 +742,14 @@ def check_goodput_args(args):

742 raise ValueError(

743 f"Invalid metric name found, {slo_name}: {slo_val}. "

744 "The service level objective name should be one of "

708 - f"{str(VALID_NAMES)}. ")

745 + f"{str(VALID_NAMES)}. "

746 + )

747 if slo_val < 0:

748 raise ValueError(

749 f"Invalid value found, {slo_name}: {slo_val}. "

750 "The service level objective value should be "

713 - "non-negative.")

751 + "non-negative."

752 + )

753 return goodput_config_dict

754

755

@@ -736,19 +775,19 @@ def main(args: argparse.Namespace):

775 tokenizer_mode=args.tokenizer_mode,

776 )

777

739 - if args.dataset == 'grammar':

740 - args.structure_type = 'guided_grammar'

741 - elif args.dataset == 'regex':

742 - args.structure_type = 'guided_regex'

743 - elif args.dataset == 'choice':

744 - args.structure_type = 'guided_choice'

778 + if args.dataset == "grammar":

779 + args.structure_type = "guided_grammar"

780 + elif args.dataset == "regex":

781 + args.structure_type = "guided_regex"

782 + elif args.dataset == "choice":

783 + args.structure_type = "guided_choice"

784 else:

746 - args.structure_type = 'guided_json'

785 + args.structure_type = "guided_json"

786

787 if args.no_structured_output:

788 args.structured_output_ratio = 0

789 if args.save_results:

751 - result_file_name = f'{args.structured_output_ratio}guided'

790 + result_file_name = f"{args.structured_output_ratio}guided"

791 result_file_name += f"_{backend}"

792 result_file_name += f"_{args.request_rate}qps"

793 result_file_name += f"_{args.model.split('/')[-1]}"

@@ -776,36 +815,29 @@ def main(args: argparse.Namespace):

815 disable_tqdm=args.disable_tqdm,

816 profile=args.profile,

817 selected_percentile_metrics=args.percentile_metrics.split(","),

779 - selected_percentiles=[

780 - float(p) for p in args.metric_percentiles.split(",")

781 - ],

818 + selected_percentiles=[float(p) for p in args.metric_percentiles.split(",")],

819 ignore_eos=args.ignore_eos,

820 max_concurrency=args.max_concurrency,

821 structured_output_ratio=args.structured_output_ratio,

822 goodput_config_dict=goodput_config_dict,

786 - ))

823 + )

824 + )

825

826 # Save config and results to json

827 score = evaluate(ret, args)

790 - print("correct_rate(%)", score, '\n')

828 + print("correct_rate(%)", score, "\n")

829 if args.save_results:

830 results = {

793 - "backend":

794 - backend,

795 - "model_id":

796 - model_id,

797 - "tokenizer_id":

798 - tokenizer_id,

799 - "num_prompts":

800 - args.num_prompts,

801 - "request_rate":

802 - args.request_rate if args.request_rate < float("inf") else "inf",

803 - "burstiness":

804 - args.burstiness,

805 - "max_concurrency":

806 - args.max_concurrency,

807 - "correct_rate(%)":

808 - score

831 + "backend": backend,

832 + "model_id": model_id,

833 + "tokenizer_id": tokenizer_id,

834 + "num_prompts": args.num_prompts,

835 + "request_rate": args.request_rate

836 + if args.request_rate < float("inf")

837 + else "inf",

838 + "burstiness": args.burstiness,

839 + "max_concurrency": args.max_concurrency,

840 + "correct_rate(%)": score,

841 }

842 results = {"outputs": ret, **results, **benchmark_result}

843

@@ -814,13 +846,14 @@ def main(args: argparse.Namespace):

846 result_file_name = args.result_filename

847 if args.result_dir:

848 result_file_name = os.path.join(args.result_dir, result_file_name)

817 - with open(result_file_name, "w", encoding='utf-8') as outfile:

849 + with open(result_file_name, "w", encoding="utf-8") as outfile:

850 json.dump(results, outfile, indent=4)

851

852

853 if __name__ == "__main__":

854 parser = FlexibleArgumentParser(

823 - description="Benchmark the online serving throughput.")

855 + description="Benchmark the online serving throughput."

856 + )

857 parser.add_argument(

858 "--backend",

859 type=str,

@@ -842,16 +875,14 @@ if __name__ == "__main__":

875 default="/v1/completions",

876 help="API endpoint.",

877 )

845 - parser.add_argument("--dataset",

846 - default='json',

847 - choices=[

848 - 'json', 'json-unique', 'grammar', 'regex',

849 - 'choice', 'xgrammar_bench'

850 - ])

851 - parser.add_argument("--json-schema-path",

852 - type=str,

853 - default=None,

854 - help="Path to json schema.")

878 + parser.add_argument(

879 + "--dataset",

880 + default="json",

881 + choices=["json", "json-unique", "grammar", "regex", "choice", "xgrammar_bench"],

882 + )

883 + parser.add_argument(

884 + "--json-schema-path", type=str, default=None, help="Path to json schema."

885 + )

886 parser.add_argument(

887 "--max-concurrency",

888 type=int,

@@ -863,7 +894,8 @@ if __name__ == "__main__":

894 "initiated, this argument will control how many are actually allowed "

895 "to execute at a time. This means that when used in combination, the "

896 "actual request rate may be lower than specified with --request-rate, "

866 - "if the server is not processing requests fast enough to keep up.")

897 + "if the server is not processing requests fast enough to keep up.",

898 + )

899 parser.add_argument(

900 "--model",

901 type=str,

@@ -873,15 +905,13 @@ if __name__ == "__main__":

905 parser.add_argument(

906 "--tokenizer",

907 type=str,

876 - help=

877 - "Name or path of the tokenizer, if not using the default tokenizer.", # noqa: E501

908 + help="Name or path of the tokenizer, if not using the default tokenizer.", # noqa: E501

909 )

910 parser.add_argument(

911 "--tokenizer-mode",

912 type=str,

913 default="auto",

883 - help=

884 - "Name or path of the tokenizer, if not using the default tokenizer.", # noqa: E501

914 + help="Name or path of the tokenizer, if not using the default tokenizer.", # noqa: E501

915 )

916 parser.add_argument(

917 "--num-prompts",

@@ -958,44 +988,51 @@ if __name__ == "__main__":

988 "--ignore-eos",

989 action="store_true",

990 help="Set ignore_eos flag when sending the benchmark request."

961 - "Warning: ignore_eos is not supported in deepspeed_mii and tgi.")

991 + "Warning: ignore_eos is not supported in deepspeed_mii and tgi.",

992 + )

993 parser.add_argument(

994 "--percentile-metrics",

995 type=str,

996 default="ttft,tpot,itl",

997 help="Comma-separated list of selected metrics to report percentils. "

998 "This argument specifies the metrics to report percentiles. "

968 - "Allowed metric names are \"ttft\", \"tpot\", \"itl\", \"e2el\". "

969 - "Default value is \"ttft,tpot,itl\".")

999 + 'Allowed metric names are "ttft", "tpot", "itl", "e2el". '

1000 + 'Default value is "ttft,tpot,itl".',

1001 + )

1002 parser.add_argument(

1003 "--metric-percentiles",

1004 type=str,

1005 default="99",

1006 help="Comma-separated list of percentiles for selected metrics. "

975 - "To report 25-th, 50-th, and 75-th percentiles, use \"25,50,75\". "

976 - "Default value is \"99\". "

977 - "Use \"--percentile-metrics\" to select metrics.",

1007 + 'To report 25-th, 50-th, and 75-th percentiles, use "25,50,75". '

1008 + 'Default value is "99". '

1009 + 'Use "--percentile-metrics" to select metrics.',

1010 )

1011 parser.add_argument(

1012 "--goodput",

1013 nargs="+",

1014 required=False,

983 - help="Specify service level objectives for goodput as \"KEY:VALUE\" "

1015 + help='Specify service level objectives for goodput as "KEY:VALUE" '

1016 "pairs, where the key is a metric name, and the value is in "

985 - "milliseconds. Multiple \"KEY:VALUE\" pairs can be provided, "

1017 + 'milliseconds. Multiple "KEY:VALUE" pairs can be provided, '

1018 "separated by spaces. Allowed request level metric names are "

987 - "\"ttft\", \"tpot\", \"e2el\". For more context on the definition of "

1019 + '"ttft", "tpot", "e2el". For more context on the definition of '

1020 "goodput, refer to DistServe paper: https://arxiv.org/pdf/2401.09670 "

989 - "and the blog: https://hao-ai-lab.github.io/blogs/distserve")

990 -

991 - parser.add_argument("--no-structured-output",

992 - action='store_true',

993 - default=False,

994 - help="Whether to disable JSON decoding or not.")

995 - parser.add_argument("--structured-output-ratio",

996 - type=float,

997 - default=1.0,

998 - help="Ratio of Structured Outputs requests")

1021 + "and the blog: https://hao-ai-lab.github.io/blogs/distserve",

1022 + )

1023 +

1024 + parser.add_argument(

1025 + "--no-structured-output",

1026 + action="store_true",

1027 + default=False,

1028 + help="Whether to disable JSON decoding or not.",

1029 + )

1030 + parser.add_argument(

1031 + "--structured-output-ratio",

1032 + type=float,

1033 + default=1.0,

1034 + help="Ratio of Structured Outputs requests",

1035 + )

1036

1037 args = parser.parse_args()

1038 main(args)

benchmarks/benchmark_throughput.py

+257 -180

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 """Benchmark offline inference throughput."""

3 +

4 import argparse

5 import dataclasses

6 import json

@@ -11,18 +12,25 @@ from typing import Any, Optional, Union

13 import torch

14 import uvloop

14 -from benchmark_dataset import (AIMODataset, BurstGPTDataset,

15 - ConversationDataset, InstructCoderDataset,

16 - RandomDataset, SampleRequest, ShareGPTDataset,

17 - SonnetDataset, VisionArenaDataset)

18 -from benchmark_utils import convert_to_pytorch_benchmark_format, write_to_json

15 from tqdm import tqdm

20 -from transformers import (AutoModelForCausalLM, AutoTokenizer,

21 - PreTrainedTokenizerBase)

22 -

16 +from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerBase

17 +

18 +from benchmark_dataset import (

19 + AIMODataset,

20 + BurstGPTDataset,

21 + ConversationDataset,

22 + InstructCoderDataset,

23 + RandomDataset,

24 + SampleRequest,

25 + ShareGPTDataset,

26 + SonnetDataset,

27 + VisionArenaDataset,

28 +)

29 +from benchmark_utils import convert_to_pytorch_benchmark_format, write_to_json

30 from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs

31 from vllm.entrypoints.openai.api_server import (

25 - build_async_engine_client_from_engine_args)

32 + build_async_engine_client_from_engine_args,

33 +)

34 from vllm.inputs import TextPrompt, TokensPrompt

35 from vllm.lora.request import LoRARequest

36 from vllm.outputs import RequestOutput

@@ -37,23 +45,30 @@ def run_vllm(

45 disable_detokenize: bool = False,

46 ) -> tuple[float, Optional[list[RequestOutput]]]:

47 from vllm import LLM, SamplingParams

48 +

49 llm = LLM(**dataclasses.asdict(engine_args))

50 assert all(

42 - llm.llm_engine.model_config.max_model_len >= (

43 - request.prompt_len + request.expected_output_len)

44 - for request in requests), (

45 - "Please ensure that max_model_len is greater than the sum of"

46 - " prompt_len and expected_output_len for all requests.")

51 + llm.llm_engine.model_config.max_model_len

52 + >= (request.prompt_len + request.expected_output_len)

53 + for request in requests

54 + ), (

55 + "Please ensure that max_model_len is greater than the sum of"

56 + " prompt_len and expected_output_len for all requests."

57 + )

58 # Add the requests to the engine.

59 prompts: list[Union[TextPrompt, TokensPrompt]] = []

60 sampling_params: list[SamplingParams] = []

61 for request in requests:

62 prompts.append(

52 - TokensPrompt(prompt_token_ids=request.prompt["prompt_token_ids"],

53 - multi_modal_data=request.multi_modal_data)

54 - if "prompt_token_ids" in request.prompt else \

55 - TextPrompt(prompt=request.prompt,

56 - multi_modal_data=request.multi_modal_data))

63 + TokensPrompt(

64 + prompt_token_ids=request.prompt["prompt_token_ids"],

65 + multi_modal_data=request.multi_modal_data,

66 + )

67 + if "prompt_token_ids" in request.prompt

68 + else TextPrompt(

69 + prompt=request.prompt, multi_modal_data=request.multi_modal_data

70 + )

71 + )

72 sampling_params.append(

73 SamplingParams(

74 n=n,

@@ -62,7 +77,8 @@ def run_vllm(

77 ignore_eos=True,

78 max_tokens=request.expected_output_len,

79 detokenize=not disable_detokenize,

65 - ))

80 + )

81 + )

82 lora_requests: Optional[list[LoRARequest]] = None

83 if engine_args.enable_lora:

84 lora_requests = [request.lora_request for request in requests]

@@ -72,10 +88,9 @@ def run_vllm(

88 outputs = None

89 if not use_beam_search:

90 start = time.perf_counter()

75 - outputs = llm.generate(prompts,

76 - sampling_params,

77 - lora_request=lora_requests,

78 - use_tqdm=True)

91 + outputs = llm.generate(

92 + prompts, sampling_params, lora_request=lora_requests, use_tqdm=True

93 + )

94 end = time.perf_counter()

95 else:

96 assert lora_requests is None, "BeamSearch API does not support LoRA"

@@ -91,30 +106,35 @@ def run_vllm(

106 beam_width=n,

107 max_tokens=output_len,

108 ignore_eos=True,

94 - ))

109 + ),

110 + )

111 end = time.perf_counter()

112 return end - start, outputs

113

114

115 def run_vllm_chat(

100 - requests: list[SampleRequest],

101 - n: int,

102 - engine_args: EngineArgs,

103 - disable_detokenize: bool = False) -> tuple[float, list[RequestOutput]]:

116 + requests: list[SampleRequest],

117 + n: int,

118 + engine_args: EngineArgs,

119 + disable_detokenize: bool = False,

120 +) -> tuple[float, list[RequestOutput]]:

121 """

122 Run vLLM chat benchmark. This function is recommended ONLY for benchmarking

123 multimodal models as it properly handles multimodal inputs and chat

124 formatting. For non-multimodal models, use run_vllm() instead.

125 """

126 from vllm import LLM, SamplingParams

127 +

128 llm = LLM(**dataclasses.asdict(engine_args))

129

130 assert all(

113 - llm.llm_engine.model_config.max_model_len >= (

114 - request.prompt_len + request.expected_output_len)

115 - for request in requests), (

116 - "Please ensure that max_model_len is greater than the sum of "

117 - "prompt_len and expected_output_len for all requests.")

131 + llm.llm_engine.model_config.max_model_len

132 + >= (request.prompt_len + request.expected_output_len)

133 + for request in requests

134 + ), (

135 + "Please ensure that max_model_len is greater than the sum of "

136 + "prompt_len and expected_output_len for all requests."

137 + )

138

139 prompts = []

140 sampling_params: list[SamplingParams] = []

@@ -128,7 +148,8 @@ def run_vllm_chat(

148 ignore_eos=True,

149 max_tokens=request.expected_output_len,

150 detokenize=not disable_detokenize,

131 - ))

151 + )

152 + )

153 start = time.perf_counter()

154 outputs = llm.chat(prompts, sampling_params, use_tqdm=True)

155 end = time.perf_counter()

@@ -145,14 +166,17 @@ async def run_vllm_async(

166 from vllm import SamplingParams

167

168 async with build_async_engine_client_from_engine_args(

148 - engine_args, disable_frontend_multiprocessing) as llm:

169 + engine_args, disable_frontend_multiprocessing

170 + ) as llm:

171 model_config = await llm.get_model_config()

172 assert all(

151 - model_config.max_model_len >= (request.prompt_len +

152 - request.expected_output_len)

153 - for request in requests), (

154 - "Please ensure that max_model_len is greater than the sum of"

155 - " prompt_len and expected_output_len for all requests.")

173 + model_config.max_model_len

174 + >= (request.prompt_len + request.expected_output_len)

175 + for request in requests

176 + ), (

177 + "Please ensure that max_model_len is greater than the sum of"

178 + " prompt_len and expected_output_len for all requests."

179 + )

180

181 # Add the requests to the engine.

182 prompts: list[Union[TextPrompt, TokensPrompt]] = []

@@ -160,11 +184,15 @@ async def run_vllm_async(

184 lora_requests: list[Optional[LoRARequest]] = []

185 for request in requests:

186 prompts.append(

163 - TokensPrompt(prompt_token_ids=request.prompt["prompt_token_ids"],

164 - multi_modal_data=request.multi_modal_data)

165 - if "prompt_token_ids" in request.prompt else \

166 - TextPrompt(prompt=request.prompt,

167 - multi_modal_data=request.multi_modal_data))

187 + TokensPrompt(

188 + prompt_token_ids=request.prompt["prompt_token_ids"],

189 + multi_modal_data=request.multi_modal_data,

190 + )

191 + if "prompt_token_ids" in request.prompt

192 + else TextPrompt(

193 + prompt=request.prompt, multi_modal_data=request.multi_modal_data

194 + )

195 + )

196 sampling_params.append(

197 SamplingParams(

198 n=n,

@@ -173,17 +201,16 @@ async def run_vllm_async(

201 ignore_eos=True,

202 max_tokens=request.expected_output_len,

203 detokenize=not disable_detokenize,

176 - ))

204 + )

205 + )

206 lora_requests.append(request.lora_request)

207

208 generators = []

209 start = time.perf_counter()

181 - for i, (prompt, sp,

182 - lr) in enumerate(zip(prompts, sampling_params, lora_requests)):

183 - generator = llm.generate(prompt,

184 - sp,

185 - lora_request=lr,

186 - request_id=f"test{i}")

210 + for i, (prompt, sp, lr) in enumerate(

211 + zip(prompts, sampling_params, lora_requests)

212 + ):

213 + generator = llm.generate(prompt, sp, lora_request=lr, request_id=f"test{i}")

214 generators.append(generator)

215 all_gens = merge_async_iterators(*generators)

216 async for i, res in all_gens:

@@ -202,7 +229,8 @@ def run_hf(

229 disable_detokenize: bool = False,

230 ) -> float:

231 llm = AutoModelForCausalLM.from_pretrained(

205 - model, torch_dtype=torch.float16, trust_remote_code=trust_remote_code)

232 + model, torch_dtype=torch.float16, trust_remote_code=trust_remote_code

233 + )

234 if llm.config.model_type == "llama":

235 # To enable padding in the HF backend.

236 tokenizer.pad_token = tokenizer.eos_token

@@ -225,14 +253,15 @@ def run_hf(

253 # Check if we can add more requests to the batch.

254 next_prompt_len = requests[i + 1].prompt_len

255 next_output_len = requests[i + 1].expected_output_len

228 - if (max(max_prompt_len, next_prompt_len) +

229 - max(max_output_len, next_output_len)) <= 2048:

256 + if (

257 + max(max_prompt_len, next_prompt_len)

258 + + max(max_output_len, next_output_len)

259 + ) <= 2048:

260 # We can add more requests to the batch.

261 continue

262

263 # Generate the sequences.

234 - input_ids = tokenizer(batch, return_tensors="pt",

235 - padding=True).input_ids

264 + input_ids = tokenizer(batch, return_tensors="pt", padding=True).input_ids

265 llm_outputs = llm.generate(

266 input_ids=input_ids.cuda(),

267 do_sample=True,

@@ -262,6 +291,7 @@ def run_mii(

291 output_len: int,

292 ) -> float:

293 from mii import client, serve

294 +

295 llm = serve(model, tensor_parallel=tensor_parallel_size)

296 prompts = [request.prompt for request in requests]

297

@@ -273,8 +303,9 @@ def run_mii(

303 return end - start

304

305

276 -def save_to_pytorch_benchmark_format(args: argparse.Namespace,

277 - results: dict[str, Any]) -> None:

306 +def save_to_pytorch_benchmark_format(

307 + args: argparse.Namespace, results: dict[str, Any]

308 +) -> None:

309 pt_records = convert_to_pytorch_benchmark_format(

310 args=args,

311 metrics={

@@ -282,9 +313,9 @@ def save_to_pytorch_benchmark_format(args: argparse.Namespace,

313 "tokens_per_second": [results["tokens_per_second"]],

314 },

315 extra_info={

285 - k: results[k]

286 - for k in ["elapsed_time", "num_requests", "total_num_tokens"]

287 - })

316 + k: results[k] for k in ["elapsed_time", "num_requests", "total_num_tokens"]

317 + },

318 + )

319 if pt_records:

320 # Don't use json suffix here as we don't want CI to pick it up

321 pt_file = f"{os.path.splitext(args.output_json)[0]}.pytorch.json"

@@ -316,7 +347,8 @@ def get_requests(args, tokenizer):

347 sample_kwargs["enable_multimodal_chat"] = True

348 elif args.dataset_name == "sonnet":

349 assert tokenizer.chat_template or tokenizer.default_chat_template, (

319 - "Tokenizer/model must have chat template for sonnet dataset.")

350 + "Tokenizer/model must have chat template for sonnet dataset."

351 + )

352 dataset_cls = SonnetDataset

353 sample_kwargs["prefix_len"] = args.prefix_len

354 sample_kwargs["return_prompt_formatted"] = True

@@ -325,21 +357,21 @@ def get_requests(args, tokenizer):

357 elif args.dataset_name == "hf":

358 if args.dataset_path in VisionArenaDataset.SUPPORTED_DATASET_PATHS:

359 dataset_cls = VisionArenaDataset

328 - common_kwargs['dataset_subset'] = None

329 - common_kwargs['dataset_split'] = "train"

360 + common_kwargs["dataset_subset"] = None

361 + common_kwargs["dataset_split"] = "train"

362 sample_kwargs["enable_multimodal_chat"] = True

363 elif args.dataset_path in InstructCoderDataset.SUPPORTED_DATASET_PATHS:

364 dataset_cls = InstructCoderDataset

333 - common_kwargs['dataset_split'] = "train"

365 + common_kwargs["dataset_split"] = "train"

366 elif args.dataset_path in ConversationDataset.SUPPORTED_DATASET_PATHS:

367 dataset_cls = ConversationDataset

336 - common_kwargs['dataset_subset'] = args.hf_subset

337 - common_kwargs['dataset_split'] = args.hf_split

368 + common_kwargs["dataset_subset"] = args.hf_subset

369 + common_kwargs["dataset_split"] = args.hf_split

370 sample_kwargs["enable_multimodal_chat"] = True

371 elif args.dataset_path in AIMODataset.SUPPORTED_DATASET_PATHS:

372 dataset_cls = AIMODataset

341 - common_kwargs['dataset_subset'] = None

342 - common_kwargs['dataset_split'] = "train"

373 + common_kwargs["dataset_subset"] = None

374 + common_kwargs["dataset_split"] = "train"

375 else:

376 raise ValueError(f"Unknown dataset name: {args.dataset_name}")

377 # Remove None values

@@ -354,10 +386,10 @@ def main(args: argparse.Namespace):

386 random.seed(args.seed)

387 # Sample the requests.

388 tokenizer = AutoTokenizer.from_pretrained(

357 - args.tokenizer, trust_remote_code=args.trust_remote_code)

389 + args.tokenizer, trust_remote_code=args.trust_remote_code

390 + )

391 requests = get_requests(args, tokenizer)

359 - is_multi_modal = any(request.multi_modal_data is not None

360 - for request in requests)

392 + is_multi_modal = any(request.multi_modal_data is not None for request in requests)

393 request_outputs: Optional[list[RequestOutput]] = None

394 if args.backend == "vllm":

395 if args.async_engine:

@@ -368,23 +400,34 @@ def main(args: argparse.Namespace):

400 AsyncEngineArgs.from_cli_args(args),

401 args.disable_frontend_multiprocessing,

402 args.disable_detokenize,

371 - ))

403 + )

404 + )

405 else:

406 elapsed_time, request_outputs = run_vllm(

374 - requests, args.n, EngineArgs.from_cli_args(args),

375 - args.disable_detokenize)

407 + requests,

408 + args.n,

409 + EngineArgs.from_cli_args(args),

410 + args.disable_detokenize,

411 + )

412 elif args.backend == "hf":

413 assert args.tensor_parallel_size == 1

378 - elapsed_time = run_hf(requests, args.model, tokenizer, args.n,

379 - args.hf_max_batch_size, args.trust_remote_code,

380 - args.disable_detokenize)

414 + elapsed_time = run_hf(

415 + requests,

416 + args.model,

417 + tokenizer,

418 + args.n,

419 + args.hf_max_batch_size,

420 + args.trust_remote_code,

421 + args.disable_detokenize,

422 + )

423 elif args.backend == "mii":

382 - elapsed_time = run_mii(requests, args.model, args.tensor_parallel_size,

383 - args.output_len)

424 + elapsed_time = run_mii(

425 + requests, args.model, args.tensor_parallel_size, args.output_len

426 + )

427 elif args.backend == "vllm-chat":

428 elapsed_time, request_outputs = run_vllm_chat(

386 - requests, args.n, EngineArgs.from_cli_args(args),

387 - args.disable_detokenize)

429 + requests, args.n, EngineArgs.from_cli_args(args), args.disable_detokenize

430 + )

431 else:

432 raise ValueError(f"Unknown backend: {args.backend}")

433

@@ -396,28 +439,31 @@ def main(args: argparse.Namespace):

439 for ro in request_outputs:

440 if not isinstance(ro, RequestOutput):

441 continue

399 - total_prompt_tokens += len(

400 - ro.prompt_token_ids) if ro.prompt_token_ids else 0

401 - total_output_tokens += sum(

402 - len(o.token_ids) for o in ro.outputs if o)

442 + total_prompt_tokens += (

443 + len(ro.prompt_token_ids) if ro.prompt_token_ids else 0

444 + )

445 + total_output_tokens += sum(len(o.token_ids) for o in ro.outputs if o)

446 total_num_tokens = total_prompt_tokens + total_output_tokens

447 else:

405 - total_num_tokens = sum(r.prompt_len + r.expected_output_len

406 - for r in requests)

448 + total_num_tokens = sum(r.prompt_len + r.expected_output_len for r in requests)

449 total_output_tokens = sum(r.expected_output_len for r in requests)

450 total_prompt_tokens = total_num_tokens - total_output_tokens

451

452 if is_multi_modal and args.backend != "vllm-chat":

411 - print("\033[91mWARNING\033[0m: Multi-modal request with "

412 - f"{args.backend} backend detected. The "

413 - "following metrics are not accurate because image tokens are not"

414 - " counted. See vllm-project/vllm/issues/9778 for details.")

453 + print(

454 + "\033[91mWARNING\033[0m: Multi-modal request with "

455 + f"{args.backend} backend detected. The "

456 + "following metrics are not accurate because image tokens are not"

457 + " counted. See vllm-project/vllm/issues/9778 for details."

458 + )

459 # TODO(vllm-project/vllm/issues/9778): Count multi-modal token length.

460 # vllm-chat backend counts the image tokens now

461

418 - print(f"Throughput: {len(requests) / elapsed_time:.2f} requests/s, "

419 - f"{total_num_tokens / elapsed_time:.2f} total tokens/s, "

420 - f"{total_output_tokens / elapsed_time:.2f} output tokens/s")

462 + print(

463 + f"Throughput: {len(requests) / elapsed_time:.2f} requests/s, "

464 + f"{total_num_tokens / elapsed_time:.2f} total tokens/s, "

465 + f"{total_output_tokens / elapsed_time:.2f} output tokens/s"

466 + )

467 print(f"Total num prompt tokens: {total_prompt_tokens}")

468 print(f"Total num output tokens: {total_output_tokens}")

469

@@ -445,7 +491,8 @@ def validate_args(args):

491 warnings.warn(

492 "The '--dataset' argument will be deprecated in the next release. "

493 "Please use '--dataset-name' and '--dataset-path' instead.",

448 - stacklevel=2)

494 + stacklevel=2,

495 + )

496 args.dataset_path = args.dataset

497

498 if not getattr(args, "tokenizer", None):

@@ -458,9 +505,8 @@ def validate_args(args):

505

506 # === Dataset Configuration ===

507 if not args.dataset and not args.dataset_path:

461 - print(

462 - "When dataset path is not set, it will default to random dataset")

463 - args.dataset_name = 'random'

508 + print("When dataset path is not set, it will default to random dataset")

509 + args.dataset_name = "random"

510 if args.input_len is None:

511 raise ValueError("input_len must be provided for a random dataset")

512

@@ -468,41 +514,55 @@ def validate_args(args):

514 # --hf-subset and --hf-split: only used

515 # when dataset_name is 'hf'

516 if args.dataset_name != "hf" and (

471 - getattr(args, "hf_subset", None) is not None

472 - or getattr(args, "hf_split", None) is not None):

473 - warnings.warn("--hf-subset and --hf-split will be ignored \

517 + getattr(args, "hf_subset", None) is not None

518 + or getattr(args, "hf_split", None) is not None

519 + ):

520 + warnings.warn(

521 + "--hf-subset and --hf-split will be ignored \

522 since --dataset-name is not 'hf'.",

475 - stacklevel=2)

523 + stacklevel=2,

524 + )

525 elif args.dataset_name == "hf":

526 if args.dataset_path in (

478 - VisionArenaDataset.SUPPORTED_DATASET_PATHS.keys()

479 - | ConversationDataset.SUPPORTED_DATASET_PATHS):

480 - assert args.backend == "vllm-chat", f"{args.dataset_path} needs to use vllm-chat as the backend." #noqa: E501

481 - elif args.dataset_path in (InstructCoderDataset.SUPPORTED_DATASET_PATHS

482 - | AIMODataset.SUPPORTED_DATASET_PATHS):

483 - assert args.backend == "vllm", f"{args.dataset_path} needs to use vllm as the backend." #noqa: E501

527 + VisionArenaDataset.SUPPORTED_DATASET_PATHS.keys()

528 + | ConversationDataset.SUPPORTED_DATASET_PATHS

529 + ):

530 + assert args.backend == "vllm-chat", (

531 + f"{args.dataset_path} needs to use vllm-chat as the backend."

532 + ) # noqa: E501

533 + elif args.dataset_path in (

534 + InstructCoderDataset.SUPPORTED_DATASET_PATHS

535 + | AIMODataset.SUPPORTED_DATASET_PATHS

536 + ):

537 + assert args.backend == "vllm", (

538 + f"{args.dataset_path} needs to use vllm as the backend."

539 + ) # noqa: E501

540 else:

485 - raise ValueError(

486 - f"{args.dataset_path} is not supported by hf dataset.")

541 + raise ValueError(f"{args.dataset_path} is not supported by hf dataset.")

542

543 # --random-range-ratio: only used when dataset_name is 'random'

489 - if args.dataset_name != 'random' and args.random_range_ratio is not None:

490 - warnings.warn("--random-range-ratio will be ignored since \

544 + if args.dataset_name != "random" and args.random_range_ratio is not None:

545 + warnings.warn(

546 + "--random-range-ratio will be ignored since \

547 --dataset-name is not 'random'.",

492 - stacklevel=2)

548 + stacklevel=2,

549 + )

550

551 # --prefix-len: only used when dataset_name is 'random', 'sonnet', or not

552 # set.

496 - if args.dataset_name not in {"random", "sonnet", None

497 - } and args.prefix_len is not None:

498 - warnings.warn("--prefix-len will be ignored since --dataset-name\

553 + if (

554 + args.dataset_name not in {"random", "sonnet", None}

555 + and args.prefix_len is not None

556 + ):

557 + warnings.warn(

558 + "--prefix-len will be ignored since --dataset-name\

559 is not 'random', 'sonnet', or not set.",

500 - stacklevel=2)

560 + stacklevel=2,

561 + )

562

563 # === LoRA Settings ===

564 if getattr(args, "enable_lora", False) and args.backend != "vllm":

504 - raise ValueError(

505 - "LoRA benchmarking is only supported for vLLM backend")

565 + raise ValueError("LoRA benchmarking is only supported for vLLM backend")

566 if getattr(args, "enable_lora", False) and args.lora_path is None:

567 raise ValueError("LoRA path must be provided when enable_lora is True")

568

@@ -512,8 +572,10 @@ def validate_args(args):

572 if args.backend != "hf" and args.hf_max_batch_size is not None:

573 raise ValueError("HF max batch size is only for HF backend.")

574

515 - if args.backend in {"hf", "mii"} and getattr(args, "quantization",

516 - None) is not None:

575 + if (

576 + args.backend in {"hf", "mii"}

577 + and getattr(args, "quantization", None) is not None

578 + ):

579 raise ValueError("Quantization is only for vLLM backend.")

580

581 if args.backend == "mii" and args.dtype != "auto":

@@ -521,29 +583,32 @@ def validate_args(args):

583 if args.backend == "mii" and args.n != 1:

584 raise ValueError("n must be 1 for MII backend.")

585 if args.backend == "mii" and args.tokenizer != args.model:

524 - raise ValueError(

525 - "Tokenizer must be the same as the model for MII backend.")

586 + raise ValueError("Tokenizer must be the same as the model for MII backend.")

587

588 # --data-parallel is not supported currently.

589 # https://github.com/vllm-project/vllm/issues/16222

590 if args.data_parallel_size > 1:

591 raise ValueError(

592 "Data parallel is not supported in offline benchmark, \

532 - please use benchmark serving instead")

593 + please use benchmark serving instead"

594 + )

595

596

597 if __name__ == "__main__":

598 parser = FlexibleArgumentParser(description="Benchmark the throughput.")

537 - parser.add_argument("--backend",

538 - type=str,

539 - choices=["vllm", "hf", "mii", "vllm-chat"],

540 - default="vllm")

599 + parser.add_argument(

600 + "--backend",

601 + type=str,

602 + choices=["vllm", "hf", "mii", "vllm-chat"],

603 + default="vllm",

604 + )

605 parser.add_argument(

606 "--dataset-name",

607 type=str,

608 choices=["sharegpt", "random", "sonnet", "burstgpt", "hf"],

609 help="Name of the dataset to benchmark on.",

546 - default="sharegpt")

610 + default="sharegpt",

611 + )

612 parser.add_argument(

613 "--dataset",

614 type=str,

@@ -551,57 +616,70 @@ if __name__ == "__main__":

616 help="Path to the ShareGPT dataset, will be deprecated in\

617 the next release. The dataset is expected to "

618 "be a json in form of list[dict[..., conversations: "

554 - "list[dict[..., value: <prompt_or_response>]]]]")

555 - parser.add_argument("--dataset-path",

556 - type=str,

557 - default=None,

558 - help="Path to the dataset")

559 - parser.add_argument("--input-len",

560 - type=int,

561 - default=None,

562 - help="Input prompt length for each request")

563 - parser.add_argument("--output-len",

564 - type=int,

565 - default=None,

566 - help="Output length for each request. Overrides the "

567 - "output length from the dataset.")

568 - parser.add_argument("--n",

569 - type=int,

570 - default=1,

571 - help="Number of generated sequences per prompt.")

572 - parser.add_argument("--num-prompts",

573 - type=int,

574 - default=1000,

575 - help="Number of prompts to process.")

576 - parser.add_argument("--hf-max-batch-size",

577 - type=int,

578 - default=None,

579 - help="Maximum batch size for HF backend.")

619 + "list[dict[..., value: <prompt_or_response>]]]]",

620 + )

621 + parser.add_argument(

622 + "--dataset-path", type=str, default=None, help="Path to the dataset"

623 + )

624 + parser.add_argument(

625 + "--input-len",

626 + type=int,

627 + default=None,

628 + help="Input prompt length for each request",

629 + )

630 + parser.add_argument(

631 + "--output-len",

632 + type=int,

633 + default=None,

634 + help="Output length for each request. Overrides the "

635 + "output length from the dataset.",

636 + )

637 + parser.add_argument(

638 + "--n", type=int, default=1, help="Number of generated sequences per prompt."

639 + )

640 parser.add_argument(

581 - '--output-json',

641 + "--num-prompts", type=int, default=1000, help="Number of prompts to process."

642 + )

643 + parser.add_argument(

644 + "--hf-max-batch-size",

645 + type=int,

646 + default=None,

647 + help="Maximum batch size for HF backend.",

648 + )

649 + parser.add_argument(

650 + "--output-json",

651 type=str,

652 default=None,

584 - help='Path to save the throughput results in JSON format.')

585 - parser.add_argument("--async-engine",

586 - action='store_true',

587 - default=False,

588 - help="Use vLLM async engine rather than LLM class.")

589 - parser.add_argument("--disable-frontend-multiprocessing",

590 - action='store_true',

591 - default=False,

592 - help="Disable decoupled async engine frontend.")

653 + help="Path to save the throughput results in JSON format.",

654 + )

655 + parser.add_argument(

656 + "--async-engine",

657 + action="store_true",

658 + default=False,

659 + help="Use vLLM async engine rather than LLM class.",

660 + )

661 + parser.add_argument(

662 + "--disable-frontend-multiprocessing",

663 + action="store_true",

664 + default=False,

665 + help="Disable decoupled async engine frontend.",

666 + )

667 parser.add_argument(

668 "--disable-detokenize",

669 action="store_true",

596 - help=("Do not detokenize the response (i.e. do not include "

597 - "detokenization time in the measurement)"))

670 + help=(

671 + "Do not detokenize the response (i.e. do not include "

672 + "detokenization time in the measurement)"

673 + ),

674 + )

675 # LoRA

676 parser.add_argument(

677 "--lora-path",

678 type=str,

679 default=None,

680 help="Path to the LoRA adapters to use. This can be an absolute path, "

604 - "a relative path, or a Hugging Face model identifier.")

681 + "a relative path, or a Hugging Face model identifier.",

682 + )

683 parser.add_argument(

684 "--prefix-len",

685 type=int,

@@ -615,7 +693,8 @@ if __name__ == "__main__":

693 f"prefix_len (default: {SonnetDataset.DEFAULT_PREFIX_LEN}) "

694 "controls how much of the input is fixed lines versus "

695 "random lines, but the total input length remains approximately "

618 - "input_len tokens.")

696 + "input_len tokens.",

697 + )

698 # random dataset

699 parser.add_argument(

700 "--random-range-ratio",

@@ -629,14 +708,12 @@ if __name__ == "__main__":

708 )

709

710 # hf dtaset

632 - parser.add_argument("--hf-subset",

633 - type=str,

634 - default=None,

635 - help="Subset of the HF dataset.")

636 - parser.add_argument("--hf-split",

637 - type=str,

638 - default=None,

639 - help="Split of the HF dataset.")

711 + parser.add_argument(

712 + "--hf-subset", type=str, default=None, help="Subset of the HF dataset."

713 + )

714 + parser.add_argument(

715 + "--hf-split", type=str, default=None, help="Split of the HF dataset."

716 + )

717

718 parser = AsyncEngineArgs.add_cli_args(parser)

719 args = parser.parse_args()

benchmarks/benchmark_utils.py

+7 -8

@@ -7,9 +7,9 @@ import os

7 from typing import Any

10 -def convert_to_pytorch_benchmark_format(args: argparse.Namespace,

11 - metrics: dict[str, list],

12 - extra_info: dict[str, Any]) -> list:

10 +def convert_to_pytorch_benchmark_format(

11 + args: argparse.Namespace, metrics: dict[str, list], extra_info: dict[str, Any]

12 +) -> list:

13 """

14 Save the benchmark results in the format used by PyTorch OSS benchmark with

15 on metric per record

@@ -37,12 +37,12 @@ def convert_to_pytorch_benchmark_format(args: argparse.Namespace,

37 },

38 }

40 - tp = record["benchmark"]["extra_info"]["args"].get(

41 - "tensor_parallel_size")

40 + tp = record["benchmark"]["extra_info"]["args"].get("tensor_parallel_size")

41 # Save tensor_parallel_size parameter if it's part of the metadata

42 if not tp and "tensor_parallel_size" in extra_info:

44 - record["benchmark"]["extra_info"]["args"][

45 - "tensor_parallel_size"] = extra_info["tensor_parallel_size"]

43 + record["benchmark"]["extra_info"]["args"]["tensor_parallel_size"] = (

44 + extra_info["tensor_parallel_size"]

45 + )

47 records.append(record)

@@ -50,7 +50,6 @@ def convert_to_pytorch_benchmark_format(args: argparse.Namespace,

52 class InfEncoder(json.JSONEncoder):

53 -

53 def clear_inf(self, o: Any):

54 if isinstance(o, dict):

55 return {k: self.clear_inf(v) for k, v in o.items()}

benchmarks/cutlass_benchmarks/sparse_benchmarks.py

+250 -122

@@ -23,8 +23,9 @@ DEFAULT_TP_SIZES = [1]

25 # bench

26 -def bench_fn(label: str, sub_label: str, description: str, fn: Callable, *args,

27 - **kwargs) -> TMeasurement:

26 +def bench_fn(

27 + label: str, sub_label: str, description: str, fn: Callable, *args, **kwargs

28 +) -> TMeasurement:

29 min_run_time = 1

31 globals = {

@@ -41,16 +42,18 @@ def bench_fn(label: str, sub_label: str, description: str, fn: Callable, *args,

42 ).blocked_autorange(min_run_time=min_run_time)

44 -def bench_int8(dtype: torch.dtype, m: int, k: int, n: int, label: str,

45 - sub_label: str) -> Iterable[TMeasurement]:

45 +def bench_int8(

46 + dtype: torch.dtype, m: int, k: int, n: int, label: str, sub_label: str

47 +) -> Iterable[TMeasurement]:

48 assert dtype == torch.int8

49 b_compressed, e, a, b = make_rand_sparse_tensors(torch.int8, m, n, k)

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

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

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

52 + bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16)

52 - out = ops.cutlass_scaled_sparse_mm(a, b_compressed, e, scale_a, scale_b,

53 - torch.bfloat16)

54 + out = ops.cutlass_scaled_sparse_mm(

55 + a, b_compressed, e, scale_a, scale_b, torch.bfloat16

56 + )

57 out_ref = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16)

59 if not torch.allclose(out, out_ref):

@@ -63,54 +66,107 @@ def bench_int8(dtype: torch.dtype, m: int, k: int, n: int, label: str,

66 timers = []

67 # pytorch impl - bfloat16

68 timers.append(

66 - bench_fn(label, sub_label, "pytorch_bf16_bf16_bf16_matmul-no-scales",

67 - torch.mm, a.to(dtype=torch.bfloat16),

68 - b.to(dtype=torch.bfloat16)))

69 + bench_fn(

70 + label,

71 + sub_label,

72 + "pytorch_bf16_bf16_bf16_matmul-no-scales",

73 + torch.mm,

74 + a.to(dtype=torch.bfloat16),

75 + b.to(dtype=torch.bfloat16),

76 + )

77 + )

79 # pytorch impl - float16

80 timers.append(

72 - bench_fn(label, sub_label,

73 - "pytorch_fp16_fp16_fp16_matmul-no-scales", torch.mm,

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

81 + bench_fn(

82 + label,

83 + sub_label,

84 + "pytorch_fp16_fp16_fp16_matmul-no-scales",

85 + torch.mm,

86 + a.to(dtype=torch.float16),

87 + b.to(dtype=torch.float16),

88 + )

89 + )

91 # cutlass impl

92 timers.append(

78 - bench_fn(label, sub_label, "cutlass_i8_i8_bf16_scaled_mm",

79 - ops.cutlass_scaled_mm, a, b, scale_a, scale_b,

80 - torch.bfloat16))

93 + bench_fn(

94 + label,

95 + sub_label,

96 + "cutlass_i8_i8_bf16_scaled_mm",

97 + ops.cutlass_scaled_mm,

98 + a,

99 + b,

100 + scale_a,

101 + scale_b,

102 + torch.bfloat16,

103 + )

104 + )

105

106 # cutlass with bias

107 timers.append(

84 - bench_fn(label, sub_label, "cutlass_i8_i8_bf16_scaled_mm_bias",

85 - ops.cutlass_scaled_mm, a, b, scale_a, scale_b, torch.bfloat16,

86 - bias))

108 + bench_fn(

109 + label,

110 + sub_label,

111 + "cutlass_i8_i8_bf16_scaled_mm_bias",

112 + ops.cutlass_scaled_mm,

113 + a,

114 + b,

115 + scale_a,

116 + scale_b,

117 + torch.bfloat16,

118 + bias,

119 + )

120 + )

121

122 # cutlass sparse impl

123 timers.append(

90 - bench_fn(label, sub_label, "cutlass_i8_i8_bf16_scaled_sparse_mm",

91 - ops.cutlass_scaled_sparse_mm, a, b_compressed, e, scale_a,

92 - scale_b, torch.bfloat16))

124 + bench_fn(

125 + label,

126 + sub_label,

127 + "cutlass_i8_i8_bf16_scaled_sparse_mm",

128 + ops.cutlass_scaled_sparse_mm,

129 + a,

130 + b_compressed,

131 + e,

132 + scale_a,

133 + scale_b,

134 + torch.bfloat16,

135 + )

136 + )

137

138 # cutlass sparse with bias

139 timers.append(

96 - bench_fn(label, sub_label, "cutlass_i8_i8_bf16_scaled_sparse_mm_bias",

97 - ops.cutlass_scaled_sparse_mm, a, b_compressed, e, scale_a,

98 - scale_b, torch.bfloat16, bias))

140 + bench_fn(

141 + label,

142 + sub_label,

143 + "cutlass_i8_i8_bf16_scaled_sparse_mm_bias",

144 + ops.cutlass_scaled_sparse_mm,

145 + a,

146 + b_compressed,

147 + e,

148 + scale_a,

149 + scale_b,

150 + torch.bfloat16,

151 + bias,

152 + )

153 + )

154

155 return timers

156

157

103 -def bench_fp8(dtype: torch.dtype, m: int, k: int, n: int, label: str,

104 - sub_label: str) -> Iterable[TMeasurement]:

158 +def bench_fp8(

159 + dtype: torch.dtype, m: int, k: int, n: int, label: str, sub_label: str

160 +) -> Iterable[TMeasurement]:

161 assert dtype == torch.float8_e4m3fn

106 - b_compressed, e, a, b = make_rand_sparse_tensors(torch.float8_e4m3fn, m, n,

107 - k)

162 + b_compressed, e, a, b = make_rand_sparse_tensors(torch.float8_e4m3fn, m, n, k)

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

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

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

165 + bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16)

166

112 - out = ops.cutlass_scaled_sparse_mm(a, b_compressed, e, scale_a, scale_b,

113 - torch.bfloat16)

167 + out = ops.cutlass_scaled_sparse_mm(

168 + a, b_compressed, e, scale_a, scale_b, torch.bfloat16

169 + )

170 out_ref = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16)

171

172 if not torch.allclose(out, out_ref):

@@ -124,97 +180,165 @@ def bench_fp8(dtype: torch.dtype, m: int, k: int, n: int, label: str,

180

181 # pytorch impl w. bf16

182 timers.append(

127 - bench_fn(label, sub_label, "pytorch_bf16_bf16_bf16_matmul-no-scales",

128 - torch.mm, a.to(dtype=torch.bfloat16, device="cuda"),

129 - b.to(dtype=torch.bfloat16, device="cuda")))

183 + bench_fn(

184 + label,

185 + sub_label,

186 + "pytorch_bf16_bf16_bf16_matmul-no-scales",

187 + torch.mm,

188 + a.to(dtype=torch.bfloat16, device="cuda"),

189 + b.to(dtype=torch.bfloat16, device="cuda"),

190 + )

191 + )

192

193 # pytorch impl: bf16 output, without fp8 fast accum

194 timers.append(

133 - bench_fn(label,

134 - sub_label,

135 - "pytorch_fp8_fp8_bf16_scaled_mm",

136 - torch._scaled_mm,

137 - a,

138 - b,

139 - scale_a=scale_a,

140 - scale_b=scale_b,

141 - out_dtype=torch.bfloat16))

195 + bench_fn(

196 + label,

197 + sub_label,

198 + "pytorch_fp8_fp8_bf16_scaled_mm",

199 + torch._scaled_mm,

200 + a,

201 + b,

202 + scale_a=scale_a,

203 + scale_b=scale_b,

204 + out_dtype=torch.bfloat16,

205 + )

206 + )

207

208 # pytorch impl: bf16 output, with fp8 fast accum

209 timers.append(

145 - bench_fn(label,

146 - sub_label,

147 - "pytorch_fp8_fp8_bf16_scaled_mm_fast_accum",

148 - torch._scaled_mm,

149 - a,

150 - b,

151 - scale_a=scale_a,

152 - scale_b=scale_b,

153 - out_dtype=torch.bfloat16,

154 - use_fast_accum=True))

210 + bench_fn(

211 + label,

212 + sub_label,

213 + "pytorch_fp8_fp8_bf16_scaled_mm_fast_accum",

214 + torch._scaled_mm,

215 + a,

216 + b,

217 + scale_a=scale_a,

218 + scale_b=scale_b,

219 + out_dtype=torch.bfloat16,

220 + use_fast_accum=True,

221 + )

222 + )

223

224 # pytorch impl: fp16 output, without fp8 fast accum

225 timers.append(

158 - bench_fn(label,

159 - sub_label,

160 - "pytorch_fp8_fp8_fp16_scaled_mm",

161 - torch._scaled_mm,

162 - a,

163 - b,

164 - scale_a=scale_a,

165 - scale_b=scale_b,

166 - out_dtype=torch.float16))

226 + bench_fn(

227 + label,

228 + sub_label,

229 + "pytorch_fp8_fp8_fp16_scaled_mm",

230 + torch._scaled_mm,

231 + a,

232 + b,

233 + scale_a=scale_a,

234 + scale_b=scale_b,

235 + out_dtype=torch.float16,

236 + )

237 + )

238

239 # pytorch impl: fp16 output, with fp8 fast accum

240 timers.append(

170 - bench_fn(label,

171 - sub_label,

172 - "pytorch_fp8_fp8_fp16_scaled_mm_fast_accum",

173 - torch._scaled_mm,

174 - a,

175 - b,

176 - scale_a=scale_a,

177 - scale_b=scale_b,

178 - out_dtype=torch.float16,

179 - use_fast_accum=True))

241 + bench_fn(

242 + label,

243 + sub_label,

244 + "pytorch_fp8_fp8_fp16_scaled_mm_fast_accum",

245 + torch._scaled_mm,

246 + a,

247 + b,

248 + scale_a=scale_a,

249 + scale_b=scale_b,

250 + out_dtype=torch.float16,

251 + use_fast_accum=True,

252 + )

253 + )

254

255 # cutlass impl: bf16 output

256 timers.append(

183 - bench_fn(label, sub_label, "cutlass_fp8_fp8_bf16_scaled_mm",

184 - ops.cutlass_scaled_mm, a, b, scale_a, scale_b,

185 - torch.bfloat16))

257 + bench_fn(

258 + label,

259 + sub_label,

260 + "cutlass_fp8_fp8_bf16_scaled_mm",

261 + ops.cutlass_scaled_mm,

262 + a,

263 + b,

264 + scale_a,

265 + scale_b,

266 + torch.bfloat16,

267 + )

268 + )

269

270 # cutlass impl: bf16 output

271 timers.append(

189 - bench_fn(label, sub_label, "cutlass_fp8_fp8_bf16_scaled_sparse_mm",

190 - ops.cutlass_scaled_sparse_mm, a, b_compressed, e, scale_a,

191 - scale_b, torch.bfloat16))

272 + bench_fn(

273 + label,

274 + sub_label,

275 + "cutlass_fp8_fp8_bf16_scaled_sparse_mm",

276 + ops.cutlass_scaled_sparse_mm,

277 + a,

278 + b_compressed,

279 + e,

280 + scale_a,

281 + scale_b,

282 + torch.bfloat16,

283 + )

284 + )

285

286 # cutlass impl: fp16 output

287 timers.append(

195 - bench_fn(label, sub_label, "cutlass_fp8_fp8_fp16_scaled_sparse_mm",

196 - ops.cutlass_scaled_sparse_mm, a, b_compressed, e, scale_a,

197 - scale_b, torch.float16))

288 + bench_fn(

289 + label,

290 + sub_label,

291 + "cutlass_fp8_fp8_fp16_scaled_sparse_mm",

292 + ops.cutlass_scaled_sparse_mm,

293 + a,

294 + b_compressed,

295 + e,

296 + scale_a,

297 + scale_b,

298 + torch.float16,

299 + )

300 + )

301

302 # cutlass impl: bf16 output, with bias

303 timers.append(

201 - bench_fn(label, sub_label,

202 - "cutlass_fp8_fp8_bf16_scaled_sparse_mm_bias",

203 - ops.cutlass_scaled_sparse_mm, a, b_compressed, e, scale_a,

204 - scale_b, torch.bfloat16, bias))

304 + bench_fn(

305 + label,

306 + sub_label,

307 + "cutlass_fp8_fp8_bf16_scaled_sparse_mm_bias",

308 + ops.cutlass_scaled_sparse_mm,

309 + a,

310 + b_compressed,

311 + e,

312 + scale_a,

313 + scale_b,

314 + torch.bfloat16,

315 + bias,

316 + )

317 + )

318

319 # cutlass impl: fp16 output, with bias

320 timers.append(

208 - bench_fn(label, sub_label,

209 - "cutlass_fp8_fp8_fp16_scaled_sparse_mm_bias",

210 - ops.cutlass_scaled_sparse_mm, a, b_compressed, e, scale_a,

211 - scale_b, torch.float16, bias.to(dtype=torch.float16)))

321 + bench_fn(

322 + label,

323 + sub_label,

324 + "cutlass_fp8_fp8_fp16_scaled_sparse_mm_bias",

325 + ops.cutlass_scaled_sparse_mm,

326 + a,

327 + b_compressed,

328 + e,

329 + scale_a,

330 + scale_b,

331 + torch.float16,

332 + bias.to(dtype=torch.float16),

333 + )

334 + )

335

336 return timers

337

338

216 -def bench(dtype: torch.dtype, m: int, k: int, n: int, label: str,

217 - sub_label: str) -> Iterable[TMeasurement]:

339 +def bench(

340 + dtype: torch.dtype, m: int, k: int, n: int, label: str, sub_label: str

341 +) -> Iterable[TMeasurement]:

342 if dtype == torch.int8:

343 return bench_int8(dtype, m, k, n, label, sub_label)

344 if dtype == torch.float8_e4m3fn:

@@ -228,12 +352,12 @@ def print_timers(timers: Iterable[TMeasurement]):

352 compare.print()

353

354

231 -def run(dtype: torch.dtype,

232 - MKNs: Iterable[tuple[int, int, int]]) -> Iterable[TMeasurement]:

355 +def run(

356 + dtype: torch.dtype, MKNs: Iterable[tuple[int, int, int]]

357 +) -> Iterable[TMeasurement]:

358 results = []

359 for m, k, n in MKNs:

235 - timers = bench(dtype, m, k, n, f"scaled-{dtype}-gemm",

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

360 + timers = bench(dtype, m, k, n, f"scaled-{dtype}-gemm", f"MKN=({m}x{k}x{n})")

361 print_timers(timers)

362 results.extend(timers)

363

@@ -241,10 +365,12 @@ def run(dtype: torch.dtype,

365

366

367 # output makers

244 -def make_output(data: Iterable[TMeasurement],

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

246 - base_description: str,

247 - timestamp=None):

368 +def make_output(

369 + data: Iterable[TMeasurement],

370 + MKNs: Iterable[tuple[int, int, int]],

371 + base_description: str,

372 + timestamp=None,

373 +):

374 print(f"== All Results {base_description} ====")

375 print_timers(data)

376

@@ -258,8 +384,7 @@ def make_output(data: Iterable[TMeasurement],

384

385

386 def run_square_bench(args):

261 - dim_sizes = list(

262 - range(args.dim_start, args.dim_end + 1, args.dim_increment))

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

388 MKNs = list(zip(dim_sizes, dim_sizes, dim_sizes))

389 data = run(args.dtype, MKNs)

390

@@ -319,7 +444,7 @@ def run_model_bench(args):

444 pkl.dump(all_data, f)

445

446

322 -if __name__ == '__main__':

447 +if __name__ == "__main__":

448

449 def to_torch_dtype(dt):

450 if dt == "int8":

@@ -344,12 +469,15 @@ Benchmark Cutlass GEMM.

469 Output:

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

471 """, # noqa: E501

347 - formatter_class=argparse.RawTextHelpFormatter)

348 -

349 - parser.add_argument("--dtype",

350 - type=to_torch_dtype,

351 - required=True,

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

472 + formatter_class=argparse.RawTextHelpFormatter,

473 + )

474 +

475 + parser.add_argument(

476 + "--dtype",

477 + type=to_torch_dtype,

478 + required=True,

479 + help="Available options are ['int8', 'fp8']",

480 + )

481 subparsers = parser.add_subparsers(dest="cmd")

482

483 square_parser = subparsers.add_parser("square_bench")

@@ -368,19 +496,19 @@ Benchmark Cutlass GEMM.

496 range_parser.set_defaults(func=run_range_bench)

497

498 model_parser = subparsers.add_parser("model_bench")

371 - model_parser.add_argument("--models",

372 - nargs="+",

373 - type=str,

374 - default=DEFAULT_MODELS,

375 - choices=WEIGHT_SHAPES.keys())

376 - model_parser.add_argument("--tp-sizes",

377 - nargs="+",

378 - type=int,

379 - default=DEFAULT_TP_SIZES)

380 - model_parser.add_argument("--batch-sizes",

381 - nargs="+",

382 - type=int,

383 - default=DEFAULT_BATCH_SIZES)

499 + model_parser.add_argument(

500 + "--models",

501 + nargs="+",

502 + type=str,

503 + default=DEFAULT_MODELS,

504 + choices=WEIGHT_SHAPES.keys(),

505 + )

506 + model_parser.add_argument(

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

508 + )

509 + model_parser.add_argument(

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

511 + )

512 model_parser.set_defaults(func=run_model_bench)

513

514 args = parser.parse_args()

benchmarks/cutlass_benchmarks/utils.py

+17 -16

@@ -10,8 +10,9 @@ import vllm._custom_ops as ops

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

12 finfo = torch.finfo(torch.float8_e4m3fn)

13 - return torch.round(tensor.clamp(

14 - min=finfo.min, max=finfo.max)).to(dtype=torch.float8_e4m3fn)

13 + return torch.round(tensor.clamp(min=finfo.min, max=finfo.max)).to(

14 + dtype=torch.float8_e4m3fn

15 + )

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

@@ -26,10 +27,11 @@ def to_fp16(tensor: torch.Tensor) -> torch.Tensor:

27 return tensor.to(dtype=torch.float16)

29 -def make_rand_tensors(dtype: torch.dtype, m: int, n: int,

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

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

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

30 +def make_rand_tensors(

31 + dtype: torch.dtype, m: int, n: int, k: int

32 +) -> tuple[torch.Tensor, torch.Tensor]:

33 + a = torch.randn((m, k), device="cuda") * 5

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

36 if dtype == torch.int8:

37 return to_int8(a), to_int8(b)

@@ -49,9 +51,7 @@ def prune_to_2_4(tensor):

52 # Create binary mask

53 mask = torch.zeros_like(reshaped)

52 - mask.scatter_(dim=1,

53 - index=indices,

54 - src=torch.ones_like(indices, dtype=mask.dtype))

54 + mask.scatter_(dim=1, index=indices, src=torch.ones_like(indices, dtype=mask.dtype))

56 # Apply mask and reshape back

57 pruned = reshaped * mask

@@ -62,10 +62,11 @@ def prune_to_2_4(tensor):

62 return pruned.reshape(original_shape)

65 -def make_rand_sparse_tensors(dtype: torch.dtype, m: int, n: int,

66 - k: int) -> tuple[torch.Tensor, torch.Tensor]:

67 - a = torch.randn((m, k), device='cuda') * 5

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

65 +def make_rand_sparse_tensors(

66 + dtype: torch.dtype, m: int, n: int, k: int

67 +) -> tuple[torch.Tensor, torch.Tensor]:

68 + a = torch.randn((m, k), device="cuda") * 5

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

71 b = prune_to_2_4(b.t()).t()

@@ -86,9 +87,9 @@ def make_rand_sparse_tensors(dtype: torch.dtype, m: int, n: int,

87 return b_compressed, e, a, b

89 -def make_n_rand_sparse_tensors(num_tensors: int, dtype: torch.dtype,

90 - m: int, n: int, k: int) -> \

91 - tuple[Iterable[torch.Tensor], Iterable[torch.Tensor]]:

90 +def make_n_rand_sparse_tensors(

91 + num_tensors: int, dtype: torch.dtype, m: int, n: int, k: int

92 +) -> tuple[Iterable[torch.Tensor], Iterable[torch.Tensor]]:

93 ABs = []

94 for _ in range(num_tensors):

95 b_comp, e, a, b = make_rand_sparse_tensors(dtype, m, n, k)

benchmarks/cutlass_benchmarks/w8a8_benchmarks.py

+140 -134

@@ -16,7 +16,8 @@ from weight_shapes import WEIGHT_SHAPES

17 from vllm import _custom_ops as ops

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

19 - w8a8_block_fp8_matmul)

19 + w8a8_block_fp8_matmul,

20 +)

21 from vllm.utils import FlexibleArgumentParser

23 DEFAULT_MODELS = list(WEIGHT_SHAPES.keys())

@@ -25,8 +26,9 @@ DEFAULT_TP_SIZES = [1]

28 # bench

28 -def bench_fn(label: str, sub_label: str, description: str, fn: Callable, *args,

29 - **kwargs) -> TMeasurement:

29 +def bench_fn(

30 + label: str, sub_label: str, description: str, fn: Callable, *args, **kwargs

31 +) -> TMeasurement:

32 min_run_time = 1

34 globals = {

@@ -44,45 +46,48 @@ def bench_fn(label: str, sub_label: str, description: str, fn: Callable, *args,

48 def bench_int8(

47 - dtype: torch.dtype,

48 - m: int,

49 - k: int,

50 - n: int,

51 - label: str,

52 - sub_label: str,

53 - bench_kernels: Optional[list[str]] = None) -> Iterable[TMeasurement]:

49 + dtype: torch.dtype,

50 + m: int,

51 + k: int,

52 + n: int,

53 + label: str,

54 + sub_label: str,

55 + bench_kernels: Optional[list[str]] = None,

56 +) -> Iterable[TMeasurement]:

57 """Benchmark INT8-based kernels."""

58 assert dtype == torch.int8

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

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

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

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

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

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

62 + bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16)

63 + azp = torch.zeros((m,), device="cuda", dtype=torch.int32)

64 + azp_adj = torch.zeros((n,), device="cuda", dtype=torch.int32)

66 bench_fns = {

64 - "pytorch_bf16_bf16_bf16_matmul-no-scales":

65 - lambda: torch.mm(a.to(dtype=torch.bfloat16), b.to(dtype=torch.bfloat16)

66 - ),

67 - "pytorch_fp16_fp16_fp16_matmul-no-scales":

68 - lambda: torch.mm(a.to(dtype=torch.float16), b.to(dtype=torch.float16)),

69 - "cutlass_i8_i8_bf16_scaled_mm":

70 - lambda: ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16),

71 - "cutlass_i8_i8_bf16_scaled_mm_bias":

72 - lambda: ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16,

73 - bias),

74 - "cutlass_i8_i8_bf16_scaled_mm_azp":

75 - lambda: ops.cutlass_scaled_mm_azp(a, b, scale_a, scale_b, torch.

76 - bfloat16, azp_adj),

77 - "cutlass_i8_i8_bf16_scaled_mm_azp_bias":

78 - lambda: ops.cutlass_scaled_mm_azp(a, b, scale_a, scale_b, torch.

79 - bfloat16, azp_adj, None, bias),

80 - "cutlass_i8_i8_bf16_scaled_mm_azp_pt":

81 - lambda: ops.cutlass_scaled_mm_azp(a, b, scale_a, scale_b, torch.

82 - bfloat16, azp_adj, azp),

83 - "cutlass_i8_i8_bf16_scaled_mm_azp_pt_bias":

84 - lambda: ops.cutlass_scaled_mm_azp(a, b, scale_a, scale_b, torch.

85 - bfloat16, azp_adj, azp, bias),

67 + "pytorch_bf16_bf16_bf16_matmul-no-scales": lambda: torch.mm(

68 + a.to(dtype=torch.bfloat16), b.to(dtype=torch.bfloat16)

69 + ),

70 + "pytorch_fp16_fp16_fp16_matmul-no-scales": lambda: torch.mm(

71 + a.to(dtype=torch.float16), b.to(dtype=torch.float16)

72 + ),

73 + "cutlass_i8_i8_bf16_scaled_mm": lambda: ops.cutlass_scaled_mm(

74 + a, b, scale_a, scale_b, torch.bfloat16

75 + ),

76 + "cutlass_i8_i8_bf16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm(

77 + a, b, scale_a, scale_b, torch.bfloat16, bias

78 + ),

79 + "cutlass_i8_i8_bf16_scaled_mm_azp": lambda: ops.cutlass_scaled_mm_azp(

80 + a, b, scale_a, scale_b, torch.bfloat16, azp_adj

81 + ),

82 + "cutlass_i8_i8_bf16_scaled_mm_azp_bias": lambda: ops.cutlass_scaled_mm_azp(

83 + a, b, scale_a, scale_b, torch.bfloat16, azp_adj, None, bias

84 + ),

85 + "cutlass_i8_i8_bf16_scaled_mm_azp_pt": lambda: ops.cutlass_scaled_mm_azp(

86 + a, b, scale_a, scale_b, torch.bfloat16, azp_adj, azp

87 + ),

88 + "cutlass_i8_i8_bf16_scaled_mm_azp_pt_bias": lambda: ops.cutlass_scaled_mm_azp(

89 + a, b, scale_a, scale_b, torch.bfloat16, azp_adj, azp, bias

90 + ),

91 }

93 timers = []

@@ -96,73 +101,65 @@ def bench_int8(

101

102

103 def bench_fp8(

99 - dtype: torch.dtype,

100 - m: int,

101 - k: int,

102 - n: int,

103 - label: str,

104 - sub_label: str,

105 - bench_kernels: Optional[list[str]] = None) -> Iterable[TMeasurement]:

104 + dtype: torch.dtype,

105 + m: int,

106 + k: int,

107 + n: int,

108 + label: str,

109 + sub_label: str,

110 + bench_kernels: Optional[list[str]] = None,

111 +) -> Iterable[TMeasurement]:

112 """Benchmark FP8-based kernels."""

113 assert dtype == torch.float8_e4m3fn

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

115 a_cont = a.contiguous()

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

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

112 - block_scale_a = torch.rand((m, k // 128),

113 - device="cuda",

114 - dtype=torch.float32)

115 - block_scale_b = torch.rand((k // 128, n // 128),

116 - device="cuda",

117 - dtype=torch.float32)

118 + block_scale_a = torch.rand((m, k // 128), device="cuda", dtype=torch.float32)

119 + block_scale_b = torch.rand((k // 128, n // 128), device="cuda", dtype=torch.float32)

120 block_scale_a_M_major = block_scale_a.t().contiguous().t()

121 block_scale_b_K_major = block_scale_b.t().contiguous().t()

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

122 + bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16)

123

124 print(m, k, n)

125

126 bench_fns = {

125 - "pytorch_bf16_bf16_bf16_matmul-no-scales":

126 - lambda: torch.mm(a.to(dtype=torch.bfloat16), b.to(dtype=torch.bfloat16)

127 - ),

128 - "pytorch_fp16_fp16_fp16_matmul-no-scales":

129 - lambda: torch.mm(a.to(dtype=torch.float16), b.to(dtype=torch.float16)),

130 - "pytorch_fp8_fp8_fp16_scaled_mm":

131 - lambda: torch._scaled_mm(

132 - a, b, scale_a, scale_b, out_dtype=torch.float16),

133 - "pytorch_fp8_fp8_fp16_scaled_mm_fast_accum":

134 - lambda: torch._scaled_mm(a,

135 - b,

136 - scale_a,

137 - scale_b,

138 - out_dtype=torch.float16,

139 - use_fast_accum=True),

140 - "pytorch_fp8_fp8_bf16_scaled_mm":

141 - lambda: torch._scaled_mm(

142 - a, b, scale_a, scale_b, out_dtype=torch.bfloat16),

143 - "pytorch_fp8_fp8_bf16_scaled_mm_fast_accum":

144 - lambda: torch._scaled_mm(a,

145 - b,

146 - scale_a,

147 - scale_b,

148 - out_dtype=torch.bfloat16,

149 - use_fast_accum=True),

150 - "cutlass_fp8_fp8_bf16_scaled_mm":

151 - lambda: ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16),

152 - "cutlass_fp8_fp8_fp16_scaled_mm":

153 - lambda: ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.float16),

154 - "cutlass_fp8_fp8_bf16_scaled_mm_bias":

155 - lambda: ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16,

156 - bias),

157 - "cutlass_fp8_fp8_fp16_scaled_mm_bias":

158 - lambda: ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.float16,

159 - bias.to(dtype=torch.float16)),

160 - "triton_fp8_fp8_fp16_scaled_mm_blockwise":

161 - lambda: w8a8_block_fp8_matmul(a_cont, b.t(), block_scale_a,

162 - block_scale_b.t(), (128, 128)),

163 - "cutlass_fp8_fp8_fp16_scaled_mm_blockwise":

164 - lambda: ops.cutlass_scaled_mm(a, b, block_scale_a_M_major,

165 - block_scale_b_K_major, torch.float16),

127 + "pytorch_bf16_bf16_bf16_matmul-no-scales": lambda: torch.mm(

128 + a.to(dtype=torch.bfloat16), b.to(dtype=torch.bfloat16)

129 + ),

130 + "pytorch_fp16_fp16_fp16_matmul-no-scales": lambda: torch.mm(

131 + a.to(dtype=torch.float16), b.to(dtype=torch.float16)

132 + ),

133 + "pytorch_fp8_fp8_fp16_scaled_mm": lambda: torch._scaled_mm(

134 + a, b, scale_a, scale_b, out_dtype=torch.float16

135 + ),

136 + "pytorch_fp8_fp8_fp16_scaled_mm_fast_accum": lambda: torch._scaled_mm(

137 + a, b, scale_a, scale_b, out_dtype=torch.float16, use_fast_accum=True

138 + ),

139 + "pytorch_fp8_fp8_bf16_scaled_mm": lambda: torch._scaled_mm(

140 + a, b, scale_a, scale_b, out_dtype=torch.bfloat16

141 + ),

142 + "pytorch_fp8_fp8_bf16_scaled_mm_fast_accum": lambda: torch._scaled_mm(

143 + a, b, scale_a, scale_b, out_dtype=torch.bfloat16, use_fast_accum=True

144 + ),

145 + "cutlass_fp8_fp8_bf16_scaled_mm": lambda: ops.cutlass_scaled_mm(

146 + a, b, scale_a, scale_b, torch.bfloat16

147 + ),

148 + "cutlass_fp8_fp8_fp16_scaled_mm": lambda: ops.cutlass_scaled_mm(

149 + a, b, scale_a, scale_b, torch.float16

150 + ),

151 + "cutlass_fp8_fp8_bf16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm(

152 + a, b, scale_a, scale_b, torch.bfloat16, bias

153 + ),

154 + "cutlass_fp8_fp8_fp16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm(

155 + a, b, scale_a, scale_b, torch.float16, bias.to(dtype=torch.float16)

156 + ),

157 + "triton_fp8_fp8_fp16_scaled_mm_blockwise": lambda: w8a8_block_fp8_matmul(

158 + a_cont, b.t(), block_scale_a, block_scale_b.t(), (128, 128)

159 + ),

160 + "cutlass_fp8_fp8_fp16_scaled_mm_blockwise": lambda: ops.cutlass_scaled_mm(

161 + a, b, block_scale_a_M_major, block_scale_b_K_major, torch.float16

162 + ),

163 }

164

165 timers = []

@@ -175,13 +172,15 @@ def bench_fp8(

172 return timers

173

174

178 -def bench(dtype: torch.dtype,

179 - m: int,

180 - k: int,

181 - n: int,

182 - label: str,

183 - sub_label: str,

184 - bench_kernels: Optional[list[str]] = None) -> Iterable[TMeasurement]:

175 +def bench(

176 + dtype: torch.dtype,

177 + m: int,

178 + k: int,

179 + n: int,

180 + label: str,

181 + sub_label: str,

182 + bench_kernels: Optional[list[str]] = None,

183 +) -> Iterable[TMeasurement]:

184 if dtype == torch.int8:

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

186 if dtype == torch.float8_e4m3fn:

@@ -195,27 +194,33 @@ def print_timers(timers: Iterable[TMeasurement]):

194 compare.print()

195

196

198 -def run(dtype: torch.dtype,

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

200 - bench_kernels: Optional[list[str]] = None) -> Iterable[TMeasurement]:

197 +def run(

198 + dtype: torch.dtype,

199 + MKNs: Iterable[tuple[int, int, int]],

200 + bench_kernels: Optional[list[str]] = None,

201 +) -> Iterable[TMeasurement]:

202 results = []

203 for m, k, n in MKNs:

203 - timers = bench(dtype,

204 - m,

205 - k,

206 - n,

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

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

209 - bench_kernels=bench_kernels)

204 + timers = bench(

205 + dtype,

206 + m,

207 + k,

208 + n,

209 + f"scaled-{dtype}-gemm",

210 + f"MKN=({m}x{k}x{n})",

211 + bench_kernels=bench_kernels,

212 + )

213 print_timers(timers)

214 results.extend(timers)

215 return results

216

217

215 -def make_output(data: Iterable[TMeasurement],

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

217 - base_description: str,

218 - timestamp=None):

218 +def make_output(

219 + data: Iterable[TMeasurement],

220 + MKNs: Iterable[tuple[int, int, int]],

221 + base_description: str,

222 + timestamp=None,

223 +):

224 print(f"== All Results {base_description} ====")

225 print_timers(data)

226

@@ -226,8 +231,7 @@ def make_output(data: Iterable[TMeasurement],

231

232

233 def run_square_bench(args):

229 - dim_sizes = list(

230 - range(args.dim_start, args.dim_end + 1, args.dim_increment))

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

235 MKNs = list(zip(dim_sizes, dim_sizes, dim_sizes))

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

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

@@ -285,7 +289,7 @@ def run_model_bench(args):

289 pkl.dump(all_data, f)

290

291

288 -if __name__ == '__main__':

292 +if __name__ == "__main__":

293

294 def to_torch_dtype(dt):

295 if dt == "int8":

@@ -310,19 +314,21 @@ Benchmark Cutlass GEMM.

314 Output:

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

316 """, # noqa: E501

313 - formatter_class=argparse.RawTextHelpFormatter)

317 + formatter_class=argparse.RawTextHelpFormatter,

318 + )

319

315 - parser.add_argument("--dtype",

316 - type=to_torch_dtype,

317 - required=True,

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

320 + parser.add_argument(

321 + "--dtype",

322 + type=to_torch_dtype,

323 + required=True,

324 + help="Available options are ['int8', 'fp8']",

325 + )

326 parser.add_argument(

327 "--kernels",

328 nargs="+",

329 type=str,

330 default=None,

324 - help=

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

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

332 )

333

334 subparsers = parser.add_subparsers(dest="cmd")

@@ -343,19 +349,19 @@ Benchmark Cutlass GEMM.

349 range_parser.set_defaults(func=run_range_bench)

350

351 model_parser = subparsers.add_parser("model_bench")

346 - model_parser.add_argument("--models",

347 - nargs="+",

348 - type=str,

349 - default=DEFAULT_MODELS,

350 - choices=WEIGHT_SHAPES.keys())

351 - model_parser.add_argument("--tp-sizes",

352 - nargs="+",

353 - type=int,

354 - default=DEFAULT_TP_SIZES)

355 - model_parser.add_argument("--batch-sizes",

356 - nargs="+",

357 - type=int,

358 - default=DEFAULT_BATCH_SIZES)

352 + model_parser.add_argument(

353 + "--models",

354 + nargs="+",

355 + type=str,

356 + default=DEFAULT_MODELS,

357 + choices=WEIGHT_SHAPES.keys(),

358 + )

359 + model_parser.add_argument(

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

361 + )

362 + model_parser.add_argument(

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

364 + )

365 model_parser.set_defaults(func=run_model_bench)

366

367 args = parser.parse_args()

benchmarks/cutlass_benchmarks/weight_shapes.py

+1 -1

@@ -42,4 +42,4 @@ WEIGHT_SHAPES = {

42 ([8192, 57344], 1),

43 ([28672, 8192], 0),

44 ],

45 -}

\ No newline at end of file

45 +}

benchmarks/disagg_benchmarks/disagg_prefill_proxy_server.py

+13 -14

@@ -12,39 +12,37 @@ app = Quart(__name__)

13 async def forward_request(url, data):

14 async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:

15 - headers = {

16 - "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"

17 - }

18 - async with session.post(url=url, json=data,

19 - headers=headers) as response:

15 + headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}

16 + async with session.post(url=url, json=data, headers=headers) as response:

17 if response.status == 200:

18 # if response.headers.get('Transfer-Encoding') == 'chunked':

19 if True:

23 - async for chunk_bytes in response.content.iter_chunked(

24 - 1024):

20 + async for chunk_bytes in response.content.iter_chunked(1024):

21 yield chunk_bytes

22 else:

23 content = await response.read()

24 yield content

31 -@app.route('/v1/completions', methods=['POST'])

27 +@app.route("/v1/completions", methods=["POST"])

28 async def handle_request():

29 try:

30 original_request_data = await request.get_json()

32 prefill_request = original_request_data.copy()

33 # change max_tokens = 1 to let it only do prefill

38 - prefill_request['max_tokens'] = 1

34 + prefill_request["max_tokens"] = 1

36 # finish prefill

41 - async for _ in forward_request('http://localhost:8100/v1/completions',

42 - prefill_request):

37 + async for _ in forward_request(

38 + "http://localhost:8100/v1/completions", prefill_request

39 + ):

40 continue

42 # return decode

46 - generator = forward_request('http://localhost:8200/v1/completions',

47 - original_request_data)

43 + generator = forward_request(

44 + "http://localhost:8200/v1/completions", original_request_data

45 + )

46 response = await make_response(generator)

47 response.timeout = None

@@ -53,11 +51,12 @@ async def handle_request():

51 except Exception as e:

52 import sys

53 import traceback

54 +

55 exc_info = sys.exc_info()

56 print("Error occurred in disagg prefill proxy server")

57 print(e)

58 print("".join(traceback.format_exception(*exc_info)))

62 -if __name__ == '__main__':

61 +if __name__ == "__main__":

62 app.run(port=8000)

benchmarks/disagg_benchmarks/round_robin_proxy.py

+10 -10

@@ -8,7 +8,6 @@ from aiohttp import web

10 class RoundRobinProxy:

11 -

11 def __init__(self, target_ports):

12 self.target_ports = target_ports

13 self.port_cycle = itertools.cycle(self.target_ports)

@@ -21,14 +20,15 @@ class RoundRobinProxy:

20 try:

21 # Forward the request

22 async with session.request(

24 - method=request.method,

25 - url=target_url,

26 - headers=request.headers,

27 - data=request.content,

23 + method=request.method,

24 + url=target_url,

25 + headers=request.headers,

26 + data=request.content,

27 ) as response:

28 # Start sending the response

30 - resp = web.StreamResponse(status=response.status,

31 - headers=response.headers)

29 + resp = web.StreamResponse(

30 + status=response.status, headers=response.headers

31 + )

32 await resp.prepare(request)

34 # Stream the response content

@@ -45,11 +45,11 @@ class RoundRobinProxy:

45 async def main():

46 proxy = RoundRobinProxy([8100, 8200])

47 app = web.Application()

48 - app.router.add_route('*', '/{path:.*}', proxy.handle_request)

48 + app.router.add_route("*", "/{path:.*}", proxy.handle_request)

50 runner = web.AppRunner(app)

51 await runner.setup()

52 - site = web.TCPSite(runner, 'localhost', 8000)

52 + site = web.TCPSite(runner, "localhost", 8000)

53 await site.start()

55 print("Proxy server started on http://localhost:8000")

@@ -58,5 +58,5 @@ async def main():

58 await asyncio.Event().wait()

61 -if __name__ == '__main__':

61 +if __name__ == "__main__":

62 asyncio.run(main())

benchmarks/disagg_benchmarks/visualize_benchmark_results.py

+21 -23

@@ -6,43 +6,41 @@ import matplotlib.pyplot as plt

6 import pandas as pd

8 if __name__ == "__main__":

9 -

9 data = []

11 - for name in ['disagg_prefill', 'chunked_prefill']:

10 + for name in ["disagg_prefill", "chunked_prefill"]:

11 for qps in [2, 4, 6, 8]:

12 with open(f"results/{name}-qps-{qps}.json") as f:

13 x = json.load(f)

15 - x['name'] = name

16 - x['qps'] = qps

14 + x["name"] = name

15 + x["qps"] = qps

16 data.append(x)

18 df = pd.DataFrame.from_dict(data)

20 - dis_df = df[df['name'] == 'disagg_prefill']

21 - chu_df = df[df['name'] == 'chunked_prefill']

19 + dis_df = df[df["name"] == "disagg_prefill"]

20 + chu_df = df[df["name"] == "chunked_prefill"]

23 - plt.style.use('bmh')

24 - plt.rcParams['font.size'] = 20

22 + plt.style.use("bmh")

23 + plt.rcParams["font.size"] = 20

25 for key in [

27 - 'mean_ttft_ms', 'median_ttft_ms', 'p99_ttft_ms', 'mean_itl_ms',

28 - 'median_itl_ms', 'p99_itl_ms'

26 + "mean_ttft_ms",

27 + "median_ttft_ms",

28 + "p99_ttft_ms",

29 + "mean_itl_ms",

30 + "median_itl_ms",

31 + "p99_itl_ms",

32 ]:

30 -

33 fig, ax = plt.subplots(figsize=(11, 7))

32 - plt.plot(dis_df['qps'],

33 - dis_df[key],

34 - label='disagg_prefill',

35 - marker='o',

36 - linewidth=4)

37 - plt.plot(chu_df['qps'],

38 - chu_df[key],

39 - label='chunked_prefill',

40 - marker='o',

41 - linewidth=4)

34 + plt.plot(

35 + dis_df["qps"], dis_df[key], label="disagg_prefill", marker="o", linewidth=4

36 + )

37 + plt.plot(

38 + chu_df["qps"], chu_df[key], label="chunked_prefill", marker="o", linewidth=4

39 + )

40 ax.legend()

44 - ax.set_xlabel('QPS')

42 + ax.set_xlabel("QPS")

43 ax.set_ylabel(key)

44 ax.set_ylim(bottom=0)

47 - fig.savefig(f'results/{key}.png')

45 + fig.savefig(f"results/{key}.png")

46 plt.close(fig)

benchmarks/fused_kernels/layernorm_rms_benchmarks.py

+96 -45

@@ -24,10 +24,12 @@ class bench_params_t:

24 dtype: torch.dtype

26 def description(self):

27 - return (f'N {self.num_tokens} '

28 - f'x D {self.hidden_size} '

29 - f'x R {self.add_residual} '

30 - f'x DT {self.dtype}')

27 + return (

28 + f"N {self.num_tokens} "

29 + f"x D {self.hidden_size} "

30 + f"x R {self.add_residual} "

31 + f"x DT {self.dtype}"

32 + )

35 def get_bench_params() -> list[bench_params_t]:

@@ -38,15 +40,19 @@ def get_bench_params() -> list[bench_params_t]:

40 DTYPES = [torch.bfloat16, torch.float]

42 combinations = product(NUM_TOKENS, HIDDEN_SIZES, ADD_RESIDUAL, DTYPES)

41 - bench_params = list(map(lambda x: \

42 - bench_params_t(x[0], x[1], x[2], x[3]), combinations))

43 + bench_params = list(

44 + map(lambda x: bench_params_t(x[0], x[1], x[2], x[3]), combinations)

45 + )

46 return bench_params

49 # Reference impls

47 -def unfused_int8_impl(rms_norm_layer: RMSNorm, x: torch.Tensor,

48 - residual: Optional[torch.Tensor],

49 - quant_dtype: torch.dtype):

50 +def unfused_int8_impl(

51 + rms_norm_layer: RMSNorm,

52 + x: torch.Tensor,

53 + residual: Optional[torch.Tensor],

54 + quant_dtype: torch.dtype,

55 +):

56 # Norm

57 torch_out = None

58 if residual is None:

@@ -58,9 +64,12 @@ def unfused_int8_impl(rms_norm_layer: RMSNorm, x: torch.Tensor,

64 torch_out, _, _ = ops.scaled_int8_quant(torch_out)

61 -def unfused_fp8_impl(rms_norm_layer: RMSNorm, x: torch.Tensor,

62 - residual: Optional[torch.Tensor],

63 - quant_dtype: torch.dtype):

67 +def unfused_fp8_impl(

68 + rms_norm_layer: RMSNorm,

69 + x: torch.Tensor,

70 + residual: Optional[torch.Tensor],

71 + quant_dtype: torch.dtype,

72 +):

73 # Norm

74 torch_out = None

75 if residual is None:

@@ -73,22 +82,27 @@ def unfused_fp8_impl(rms_norm_layer: RMSNorm, x: torch.Tensor,

84 def fused_impl(

76 - rms_norm_layer: RMSNorm, # this stores the weights

77 - x: torch.Tensor,

78 - residual: Optional[torch.Tensor],

79 - quant_dtype: torch.dtype):

80 - out, _ = ops.rms_norm_dynamic_per_token_quant(x,

81 - rms_norm_layer.weight,

82 - 1e-6,

83 - quant_dtype,

84 - residual=residual)

85 + rms_norm_layer: RMSNorm, # this stores the weights

86 + x: torch.Tensor,

87 + residual: Optional[torch.Tensor],

88 + quant_dtype: torch.dtype,

89 +):

90 + out, _ = ops.rms_norm_dynamic_per_token_quant(

91 + x, rms_norm_layer.weight, 1e-6, quant_dtype, residual=residual

92 + )

95 # Bench functions

88 -def bench_fn(rms_norm_layer: RMSNorm, x: torch.Tensor, residual: torch.Tensor,

89 - quant_dtype: torch.dtype, label: str, sub_label: str,

90 - fn: Callable, description: str) -> TMeasurement:

91 -

96 +def bench_fn(

97 + rms_norm_layer: RMSNorm,

98 + x: torch.Tensor,

99 + residual: torch.Tensor,

100 + quant_dtype: torch.dtype,

101 + label: str,

102 + sub_label: str,

103 + fn: Callable,

104 + description: str,

105 +) -> TMeasurement:

106 min_run_time = 1

107

108 globals = {

@@ -106,43 +120,81 @@ def bench_fn(rms_norm_layer: RMSNorm, x: torch.Tensor, residual: torch.Tensor,

120 description=description,

121 ).blocked_autorange(min_run_time=min_run_time)

122

109 -def bench(params: bench_params_t, label: str, sub_label: str) \

110 - -> Iterable[TMeasurement]:

123

124 +def bench(params: bench_params_t, label: str, sub_label: str) -> Iterable[TMeasurement]:

125 # Make inputs

126 layer = RMSNorm(params.hidden_size, 1e-6).to(dtype=params.dtype)

127 # Make weights

128 layer.weight.data.normal_(mean=1.0, std=0.1)

129 # Make inputs

130 scale = 1 / params.hidden_size

118 - x = torch.randn(params.num_tokens,

119 - params.hidden_size,

120 - dtype=params.dtype,

121 - device='cuda') * scale

122 - residual = (torch.randn_like(x) * scale).to(device='cuda') \

123 - if params.add_residual else None

131 + x = (

132 + torch.randn(

133 + params.num_tokens, params.hidden_size, dtype=params.dtype, device="cuda"

134 + )

135 + * scale

136 + )

137 + residual = (

138 + (torch.randn_like(x) * scale).to(device="cuda") if params.add_residual else None

139 + )

140

141 timers = []

142

143 # unfused int8 impl.

144 timers.append(

129 - bench_fn(layer, x, residual, torch.int8, label, sub_label,

130 - unfused_int8_impl, "unfused_int8_impl"))

145 + bench_fn(

146 + layer,

147 + x,

148 + residual,

149 + torch.int8,

150 + label,

151 + sub_label,

152 + unfused_int8_impl,

153 + "unfused_int8_impl",

154 + )

155 + )

156

157 # unfused fp8 impl.

158 timers.append(

134 - bench_fn(layer, x, residual, torch.float8_e4m3fn, label, sub_label,

135 - unfused_fp8_impl, "unfused_fp8_impl"))

159 + bench_fn(

160 + layer,

161 + x,

162 + residual,

163 + torch.float8_e4m3fn,

164 + label,

165 + sub_label,

166 + unfused_fp8_impl,

167 + "unfused_fp8_impl",

168 + )

169 + )

170

171 # fused int8 impl.

172 timers.append(

139 - bench_fn(layer, x, residual, torch.int8, label, sub_label, fused_impl,

140 - "fused_int8_impl"))

173 + bench_fn(

174 + layer,

175 + x,

176 + residual,

177 + torch.int8,

178 + label,

179 + sub_label,

180 + fused_impl,

181 + "fused_int8_impl",

182 + )

183 + )

184

185 # fused fp8 impl.

186 timers.append(

144 - bench_fn(layer, x, residual, torch.float8_e4m3fn, label, sub_label,

145 - fused_impl, "fused_fp8_impl"))

187 + bench_fn(

188 + layer,

189 + x,

190 + residual,

191 + torch.float8_e4m3fn,

192 + label,

193 + sub_label,

194 + fused_impl,

195 + "fused_fp8_impl",

196 + )

197 + )

198

199 print_timers(timers)

200

@@ -157,13 +209,12 @@ def print_timers(timers: Iterable[TMeasurement]):

209

210

211 def main():

160 - torch.set_default_device('cuda')

212 + torch.set_default_device("cuda")

213 bench_params = get_bench_params()

214

215 timers = []

216 for bp in tqdm(bench_params):

165 - timers.extend(

166 - bench(bp, "rms-norm-dynamic-per-token-quant", bp.description()))

217 + timers.extend(bench(bp, "rms-norm-dynamic-per-token-quant", bp.description()))

218 print_timers(timers)

219

220 # pickle all the results

@@ -172,5 +223,5 @@ def main():

223 pkl.dump(timers, f)

224

225

175 -if __name__ == '__main__':

226 +if __name__ == "__main__":

227 main()

benchmarks/kernels/benchmark_aqlm.py

+118 -78

@@ -9,32 +9,39 @@ import torch.nn.functional as F

10 from vllm import _custom_ops as ops

11 from vllm.model_executor.layers.quantization.aqlm import (

12 - dequantize_weight, generic_dequantize_gemm, get_int_dtype,

13 - optimized_dequantize_gemm)

12 + dequantize_weight,

13 + generic_dequantize_gemm,

14 + get_int_dtype,

15 + optimized_dequantize_gemm,

16 +)

17 from vllm.utils import FlexibleArgumentParser

16 -os.environ['CUDA_VISIBLE_DEVICES'] = '0'

19 +os.environ["CUDA_VISIBLE_DEVICES"] = "0"

22 def torch_mult(

20 - input: torch.Tensor, # [..., in_features]

21 - weights: torch.Tensor,

22 - scales: torch.Tensor, # [num_out_groups, 1, 1, 1]

23 + # [..., in_features]

24 + input: torch.Tensor,

25 + weights: torch.Tensor,

26 + # [num_out_groups, 1, 1, 1]

27 + scales: torch.Tensor,

28 ) -> torch.Tensor:

29 output = F.linear(input, weights)

30 return output

33 def dequant_out_scale(

29 - input: torch.Tensor, # [..., in_features]

30 - codes: torch.IntTensor, # [num_out_groups, num_in_groups, num_codebooks]

31 - codebooks: torch.

32 - Tensor, # [num_codebooks, codebook_size, out_group_size, in_group_size]

33 - scales: torch.Tensor, # [num_out_groups, 1, 1, 1]

34 + # [..., in_features]

35 + input: torch.Tensor,

36 + # [num_out_groups, num_in_groups, num_codebooks]

37 + codes: torch.IntTensor,

38 + # [num_codebooks, codebook_size, out_group_size, in_group_size]

39 + codebooks: torch.Tensor,

40 + # [num_out_groups, 1, 1, 1]

41 + scales: torch.Tensor,

42 output_partition_sizes: torch.IntTensor,

43 bias: Optional[torch.Tensor],

44 ) -> torch.Tensor:

37 -

45 weights = ops.aqlm_dequant(codes, codebooks, output_partition_sizes)

47 if bias is None:

@@ -46,40 +53,42 @@ def dequant_out_scale(

53 flattened_output *= b_scales

54 return flattened_output.view(orig_shape)

55 else:

49 - b_scales = scales.view(scales.shape[:-3] + (-1, )).expand(

50 - -1, weights.shape[1])

56 + b_scales = scales.view(scales.shape[:-3] + (-1,)).expand(-1, weights.shape[1])

57 weights *= b_scales

58 return F.linear(input, weights, bias)

61 def dequant_weight_scale(

56 - input: torch.Tensor, # [..., in_features]

57 - codes: torch.IntTensor, # [num_out_groups, num_in_groups, num_codebooks]

58 - codebooks: torch.

59 - Tensor, # [num_codebooks, codebook_size, out_group_size, in_group_size]

60 - scales: torch.Tensor, # [num_out_groups, 1, 1, 1]

62 + # [..., in_features]

63 + input: torch.Tensor,

64 + # [num_out_groups, num_in_groups, num_codebooks]

65 + codes: torch.IntTensor,

66 + # [num_codebooks, codebook_size, out_group_size, in_group_size]

67 + codebooks: torch.Tensor,

68 + # [num_out_groups, 1, 1, 1]

69 + scales: torch.Tensor,

70 output_partition_sizes: torch.IntTensor,

71 bias: Optional[torch.Tensor],

72 ) -> torch.Tensor:

64 -

73 weights = ops.aqlm_dequant(codes, codebooks, output_partition_sizes)

67 - b_scales = scales.view(scales.shape[:-3] + (-1, )).expand(

68 - -1, weights.shape[1])

75 + b_scales = scales.view(scales.shape[:-3] + (-1,)).expand(-1, weights.shape[1])

76 weights *= b_scales

77 return F.linear(input, weights, bias)

80 def dequant_no_scale(

74 - input: torch.Tensor, # [..., in_features]

75 - codes: torch.IntTensor, # [num_out_groups, num_in_groups, num_codebooks]

76 - codebooks: torch.

77 - Tensor, # [num_codebooks, codebook_size, out_group_size, in_group_size]

78 - scales: torch.Tensor, # [num_out_groups, 1, 1, 1]

81 + # [..., in_features]

82 + input: torch.Tensor,

83 + # [num_out_groups, num_in_groups, num_codebooks]

84 + codes: torch.IntTensor,

85 + # [num_codebooks, codebook_size, out_group_size, in_group_size]

86 + codebooks: torch.Tensor,

87 + # [num_out_groups, 1, 1, 1]

88 + scales: torch.Tensor,

89 output_partition_sizes: torch.IntTensor,

90 bias: Optional[torch.Tensor],

91 ) -> torch.Tensor:

82 -

92 weights = ops.aqlm_dequant(codes, codebooks, output_partition_sizes)

94 return F.linear(input, weights, bias)

@@ -89,23 +98,26 @@ def dequant_no_scale(

98 # the generic pytorch version.

99 # Just visual comparison.

100 def dequant_test(k: int, parts: torch.Tensor, nbooks: int, bits: int) -> None:

92 -

101 n = int(parts.sum().item())

102

95 - device = torch.device('cuda:0')

103 + device = torch.device("cuda:0")

104

105 code_range = (1 << bits) // 2

106 ingroups = 8

107

100 - codes = torch.randint(-code_range,

101 - code_range,

102 - size=(n, k // ingroups, nbooks),

103 - dtype=get_int_dtype(bits),

104 - device=device)

108 + codes = torch.randint(

109 + -code_range,

110 + code_range,

111 + size=(n, k // ingroups, nbooks),

112 + dtype=get_int_dtype(bits),

113 + device=device,

114 + )

115

106 - codebooks = torch.randn(size=(parts.shape[0] * nbooks, 1 << bits, 1, 8),

107 - dtype=torch.float16,

108 - device=device)

116 + codebooks = torch.randn(

117 + size=(parts.shape[0] * nbooks, 1 << bits, 1, 8),

118 + dtype=torch.float16,

119 + device=device,

120 + )

121

122 count = 0

123 for index in range(16):

@@ -138,24 +150,25 @@ def dequant_test(k: int, parts: torch.Tensor, nbooks: int, bits: int) -> None:

150

151

152 def main():

141 -

153 parser = FlexibleArgumentParser(description="Benchmark aqlm performance.")

154

155 # Add arguments

145 - parser.add_argument("--nbooks",

146 - type=int,

147 - default=1,

148 - help="Number of codebooks (default: 1)")

149 - parser.add_argument("--bits",

150 - type=int,

151 - default=16,

152 - help="Number of bits per code element (default: 16)")

156 + parser.add_argument(

157 + "--nbooks", type=int, default=1, help="Number of codebooks (default: 1)"

158 + )

159 + parser.add_argument(

160 + "--bits",

161 + type=int,

162 + default=16,

163 + help="Number of bits per code element (default: 16)",

164 + )

165 parser.add_argument(

166 "--test",

167 type=bool,

168 default=False,

169 help="Run the decompression/dequant tester rather than benchmarking "

158 - "(default: False)")

170 + "(default: False)",

171 + )

172

173 # Parse the arguments

174 args = parser.parse_args()

@@ -165,7 +178,7 @@ def main():

178 bits = args.bits

179

180 if args.test:

168 - dequant_test(4096, torch.tensor((4096, )), nbooks, bits)

181 + dequant_test(4096, torch.tensor((4096,)), nbooks, bits)

182 return

183

184 # Otherwise, benchmark.

@@ -184,31 +197,54 @@ def main():

197 with open(filename, "w") as f:

198 sys.stdout = f

199

187 - print('m | k | n | n parts', end='')

200 + print("m | k | n | n parts", end="")

201 for method in methods:

189 - print(f" | {method.__name__.replace('_', ' ')} (µs)", end='')

190 - print('')

202 + print(f" | {method.__name__.replace('_', ' ')} (µs)", end="")

203 + print("")

204

205 # These are reasonable prefill sizes.

193 - ksandpartions = ((4096, (4096, 4096, 4096)), (4096, (4096, )),

194 - (4096, (11008, 11008)), (11008, (4096, )))

206 + ksandpartions = (

207 + (4096, (4096, 4096, 4096)),

208 + (4096, (4096,)),

209 + (4096, (11008, 11008)),

210 + (11008, (4096,)),

211 + )

212

213 # reasonable ranges for m.

214 for m in [

198 - 1, 2, 4, 8, 10, 12, 14, 16, 24, 32, 48, 52, 56, 64, 96, 112,

199 - 128, 256, 512, 1024, 1536, 2048, 3072, 4096

215 + 1,

216 + 2,

217 + 4,

218 + 8,

219 + 10,

220 + 12,

221 + 14,

222 + 16,

223 + 24,

224 + 32,

225 + 48,

226 + 52,

227 + 56,

228 + 64,

229 + 96,

230 + 112,

231 + 128,

232 + 256,

233 + 512,

234 + 1024,

235 + 1536,

236 + 2048,

237 + 3072,

238 + 4096,

239 ]:

201 - print(f'{m}', file=sys.__stdout__)

240 + print(f"{m}", file=sys.__stdout__)

241 for ksp in ksandpartions:

203 - run_grid(m, ksp[0], torch.tensor(ksp[1]), nbooks, bits,

204 - methods)

242 + run_grid(m, ksp[0], torch.tensor(ksp[1]), nbooks, bits, methods)

243

244 sys.stdout = sys.__stdout__

245

246

209 -def run_grid(m: int, k: int, parts: torch.Tensor, nbooks: int, bits: int,

210 - methods):

211 -

247 +def run_grid(m: int, k: int, parts: torch.Tensor, nbooks: int, bits: int, methods):

248 # I didn't see visible improvements from increasing these, but feel free :)

249 num_warmup_trials = 1

250 num_trials = 1

@@ -229,7 +265,7 @@ def run_grid(m: int, k: int, parts: torch.Tensor, nbooks: int, bits: int,

265 )

266

267 n = parts.sum().item()

232 - print(f'{m} | {k} | {n} | {parts.tolist()}', end='')

268 + print(f"{m} | {k} | {n} | {parts.tolist()}", end="")

269

270 for method in methods:

271 best_time_us = 1e20

@@ -249,32 +285,36 @@ def run_grid(m: int, k: int, parts: torch.Tensor, nbooks: int, bits: int,

285 if kernel_dur_us < best_time_us:

286 best_time_us = kernel_dur_us

287

252 - print(f' | {kernel_dur_us:.0f}', end='')

288 + print(f" | {kernel_dur_us:.0f}", end="")

289

254 - print('')

290 + print("")

291

292

257 -def run_timing(num_calls: int, m: int, k: int, parts: torch.Tensor,

258 - nbooks: int, bits: int, method) -> float:

259 -

293 +def run_timing(

294 + num_calls: int, m: int, k: int, parts: torch.Tensor, nbooks: int, bits: int, method

295 +) -> float:

296 n = int(parts.sum().item())

297

262 - device = torch.device('cuda:0')

298 + device = torch.device("cuda:0")

299

300 input = torch.randn((1, m, k), dtype=torch.float16, device=device)

301

302 code_range = (1 << bits) // 2

303 ingroups = 8

304

269 - codes = torch.randint(-code_range,

270 - code_range,

271 - size=(n, k // ingroups, nbooks),

272 - dtype=get_int_dtype(bits),

273 - device=device)

274 -

275 - codebooks = torch.randn(size=(parts.shape[0] * nbooks, 1 << bits, 1, 8),

276 - dtype=torch.float16,

277 - device=device)

305 + codes = torch.randint(

306 + -code_range,

307 + code_range,

308 + size=(n, k // ingroups, nbooks),

309 + dtype=get_int_dtype(bits),

310 + device=device,

311 + )

312 +

313 + codebooks = torch.randn(

314 + size=(parts.shape[0] * nbooks, 1 << bits, 1, 8),

315 + dtype=torch.float16,

316 + device=device,

317 + )

318

319 scales = torch.randn(size=(n, 1, 1, 1), dtype=torch.float16, device=device)

320

benchmarks/kernels/benchmark_bitblas.py

+33 -28

@@ -3,27 +3,33 @@

3 # Licensed under the MIT License.

5 from vllm.model_executor.layers.quantization.utils.bitblas_utils import (

6 - MINIMUM_BITBLAS_VERSION)

6 + MINIMUM_BITBLAS_VERSION,

7 +)

9 try:

10 import bitblas

11 +

12 if bitblas.__version__ < MINIMUM_BITBLAS_VERSION:

11 - raise ImportError("bitblas version is wrong. Please "

12 - f"install bitblas>={MINIMUM_BITBLAS_VERSION}")

13 + raise ImportError(

14 + "bitblas version is wrong. Please "

15 + f"install bitblas>={MINIMUM_BITBLAS_VERSION}"

16 + )

17 except ImportError as e:

18 bitblas_import_exception = e

15 - raise ValueError("Trying to use the bitblas backend, but could not import"

16 - f"with the following error: {bitblas_import_exception}. "

17 - "Please install bitblas through the following command: "

18 - f"`pip install bitblas>={MINIMUM_BITBLAS_VERSION}`"

19 - ) from bitblas_import_exception

19 + raise ValueError(

20 + "Trying to use the bitblas backend, but could not import"

21 + f"with the following error: {bitblas_import_exception}. "

22 + "Please install bitblas through the following command: "

23 + f"`pip install bitblas>={MINIMUM_BITBLAS_VERSION}`"

24 + ) from bitblas_import_exception

26 from bitblas import Matmul, MatmulConfig, auto_detect_nvidia_target

28 from vllm.utils import FlexibleArgumentParser

30 parser = FlexibleArgumentParser(

26 - description="Benchmark BitBLAS int4 on a specific target.")

31 + description="Benchmark BitBLAS int4 on a specific target."

32 +)

34 # Add arguments to the parser

35 parser.add_argument(

@@ -32,10 +38,9 @@ parser.add_argument(

38 default=auto_detect_nvidia_target(),

39 help="Specify the target device for benchmarking.",

40 )

35 -parser.add_argument("--group_size",

36 - type=int,

37 - default=None,

38 - help="Group size for grouped quantization.")

41 +parser.add_argument(

42 + "--group_size", type=int, default=None, help="Group size for grouped quantization."

43 +)

44 parser.add_argument(

45 "--A_dtype",

46 type=str,

@@ -82,17 +87,17 @@ parser.add_argument(

87 choices=["nt", "nn"],

88 help="Matrix layout, 'nt' for non-transpose A and transpose W.",

89 )

85 -parser.add_argument("--with_bias",

86 - action="store_true",

87 - help="Include bias in the benchmark.")

90 +parser.add_argument(

91 + "--with_bias", action="store_true", help="Include bias in the benchmark."

92 +)

93 parser.add_argument(

94 "--with_scaling",

95 action="store_true",

96 help="Include scaling factor in the quantization.",

97 )

93 -parser.add_argument("--with_zeros",

94 - action="store_true",

95 - help="Include zeros in the quantization.")

98 +parser.add_argument(

99 + "--with_zeros", action="store_true", help="Include zeros in the quantization."

100 +)

101 parser.add_argument(

102 "--zeros_mode",

103 type=str,

@@ -170,8 +175,7 @@ shapes = [

175 ]

176

177 # Build test shapes with all the shared arguments

173 -test_shapes = [(MatmulConfig, Matmul, (*shape, *shared_args))

174 - for shape in shapes]

178 +test_shapes = [(MatmulConfig, Matmul, (*shape, *shared_args)) for shape in shapes]

179

180 benchmark_sets = []

181 benchmark_sets.extend(test_shapes)

@@ -206,12 +210,12 @@ for config_key, values in benchmark_results.items():

210 func_name = args_split[0]

211 input_args_str = "-".join(args_split[1:])

212 col_widths[0] = max(col_widths[0], len(func_name) + 2, len(headers[0]) + 2)

209 - col_widths[1] = max(col_widths[1],

210 - len(input_args_str) + 2,

211 - len(headers[1]) + 2)

212 - col_widths[2] = max(col_widths[2],

213 - len(f"{values['BitBLAS_top20_latency']:.3f} ms") + 2,

214 - len(headers[2]) + 2)

213 + col_widths[1] = max(col_widths[1], len(input_args_str) + 2, len(headers[1]) + 2)

214 + col_widths[2] = max(

215 + col_widths[2],

216 + len(f"{values['BitBLAS_top20_latency']:.3f} ms") + 2,

217 + len(headers[2]) + 2,

218 + )

219 # break only if you want to measure widths from a single example;

220 # otherwise, let it loop over all items.

221

@@ -232,5 +236,6 @@ for config_key, values in benchmark_results.items():

236 f"{values['BitBLAS_top20_latency']:.3f} ms",

237 ]

238 row_str = "".join(

235 - [str(cell).ljust(col_widths[idx]) for idx, cell in enumerate(row)])

239 + [str(cell).ljust(col_widths[idx]) for idx, cell in enumerate(row)]

240 + )

241 print(row_str)

benchmarks/kernels/benchmark_cutlass_fp4_moe.py

+249 -168

@@ -5,6 +5,7 @@ kernel. The cutlass_moe_fp4 kernel takes in fp4 quantized weights and 16-bit

5 activations. The triton_moe kernel takes in fp8 weights(tensor scaled to fp8)

6 and 16-bit activations.

7 """

8 +

9 import nvtx

10 import torch

11 import torch.utils.benchmark as benchmark

@@ -12,8 +13,7 @@ import torch.utils.benchmark as benchmark

13 from vllm import _custom_ops as ops

14 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

15 from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp4

15 -from vllm.model_executor.layers.fused_moe.fused_moe import (fused_experts,

16 - fused_topk)

16 +from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk

17 from vllm.scalar_type import scalar_types

18 from vllm.utils import FlexibleArgumentParser

@@ -38,19 +38,27 @@ FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max

39 def to_fp8(tensor: torch.Tensor):

40 finfo = torch.finfo(torch.float8_e4m3fn)

41 - return torch.round(tensor.clamp(

42 - min=finfo.min, max=finfo.max)).to(dtype=torch.float8_e4m3fn)

41 + return torch.round(tensor.clamp(min=finfo.min, max=finfo.max)).to(

42 + dtype=torch.float8_e4m3fn

43 + )

45 -def bench_run(results: list[benchmark.Measurement], model: str,

46 - num_experts: int, topk: int, per_act_token: bool,

47 - per_out_ch: bool, mkn: tuple[int, int, int]):

46 +def bench_run(

47 + results: list[benchmark.Measurement],

48 + model: str,

49 + num_experts: int,

50 + topk: int,

51 + per_act_token: bool,

52 + per_out_ch: bool,

53 + mkn: tuple[int, int, int],

54 +):

55 label = "NVFP4 Blockscaled CUTLASS MOE vs FP8 Tensor Scaled Triton"

57 sub_label = (

51 - "{}, num_experts={}, topk={}, per_act_token={} per_out_ch={}, "

52 - "MKN=({})".format(model, num_experts, topk, per_act_token, per_out_ch,

53 - mkn))

58 + "{}, num_experts={}, topk={}, per_act_token={} per_out_ch={}, MKN=({})".format(

59 + model, num_experts, topk, per_act_token, per_out_ch, mkn

60 + )

61 + )

63 print(f"Testing: {sub_label}")

@@ -64,18 +72,12 @@ def bench_run(results: list[benchmark.Measurement], model: str,

73 _, a_fp8_scale = ops.scaled_fp8_quant(a)

67 - w1_fp8q = torch.empty((num_experts, 2 * n, k),

68 - device=device,

69 - dtype=torch.float8_e4m3fn)

70 - w2_fp8q = torch.empty((num_experts, k, n),

71 - device=device,

72 - dtype=torch.float8_e4m3fn)

73 - w1_fp8scale = torch.empty((num_experts, 1, 1),

74 - device=device,

75 - dtype=torch.float32)

76 - w2_fp8scale = torch.empty((num_experts, 1, 1),

77 - device=device,

78 - dtype=torch.float32)

75 + w1_fp8q = torch.empty(

76 + (num_experts, 2 * n, k), device=device, dtype=torch.float8_e4m3fn

77 + )

78 + w2_fp8q = torch.empty((num_experts, k, n), device=device, dtype=torch.float8_e4m3fn)

79 + w1_fp8scale = torch.empty((num_experts, 1, 1), device=device, dtype=torch.float32)

80 + w2_fp8scale = torch.empty((num_experts, 1, 1), device=device, dtype=torch.float32)

82 for expert in range(num_experts):

83 w1_fp8q[expert], w1_fp8scale[expert] = ops.scaled_fp8_quant(w1[expert])

@@ -91,26 +93,24 @@ def bench_run(results: list[benchmark.Measurement], model: str,

93 topk_weights, topk_ids = fused_topk(a, score, topk, renormalize=False)

95 quant_blocksize = 16

94 - w1_blockscale = torch.empty((num_experts, 2 * n, k // quant_blocksize),

95 - device=device,

96 - dtype=torch.float8_e4m3fn)

97 - w2_blockscale = torch.empty((num_experts, k, n // quant_blocksize),

98 - device=device,

99 - dtype=torch.float8_e4m3fn)

96 + w1_blockscale = torch.empty(

97 + (num_experts, 2 * n, k // quant_blocksize),

98 + device=device,

99 + dtype=torch.float8_e4m3fn,

100 + )

101 + w2_blockscale = torch.empty(

102 + (num_experts, k, n // quant_blocksize), device=device, dtype=torch.float8_e4m3fn

103 + )

104

105 # n_b_scales = 2 * n if per_out_ch else 1

106 # k_b_scales = k if per_out_ch else 1

103 - w1_fp4 = torch.empty((num_experts, 2 * n, k // 2),

104 - device=device,

105 - dtype=torch.uint8)

106 - w2_fp4 = torch.empty((num_experts, k, n // 2),

107 - device=device,

108 - dtype=torch.uint8)

109 -

110 - w1_gs = torch.empty((num_experts, ), device=device, dtype=torch.float32)

111 - w2_gs = torch.empty((num_experts, ), device=device, dtype=torch.float32)

112 - a1_gs = torch.ones((num_experts, ), device=device, dtype=torch.float32)

113 - a2_gs = torch.ones((num_experts, ), device=device, dtype=torch.float32)

107 + w1_fp4 = torch.empty((num_experts, 2 * n, k // 2), device=device, dtype=torch.uint8)

108 + w2_fp4 = torch.empty((num_experts, k, n // 2), device=device, dtype=torch.uint8)

109 +

110 + w1_gs = torch.empty((num_experts,), device=device, dtype=torch.float32)

111 + w2_gs = torch.empty((num_experts,), device=device, dtype=torch.float32)

112 + a1_gs = torch.ones((num_experts,), device=device, dtype=torch.float32)

113 + a2_gs = torch.ones((num_experts,), device=device, dtype=torch.float32)

114

115 for expert in range(num_experts):

116 w1_e = w1[expert]

@@ -121,96 +121,141 @@ def bench_run(results: list[benchmark.Measurement], model: str,

121 w2_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax

122

123 w1_fp4[expert], w1_blockscale[expert] = ops.scaled_fp4_quant(

124 - w1_e, w1_gs[expert])

124 + w1_e, w1_gs[expert]

125 + )

126

127 w2_fp4[expert], w2_blockscale[expert] = ops.scaled_fp4_quant(

127 - w2_e, w2_gs[expert])

128 -

129 - def run_triton_moe(a: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor,

130 - topk_weights: torch.Tensor, topk_ids: torch.Tensor,

131 - w1_scale: torch.Tensor, w2_scale: torch.Tensor,

132 - a_fp8_scale: torch.Tensor, num_repeats: int):

128 + w2_e, w2_gs[expert]

129 + )

130 +

131 + def run_triton_moe(

132 + a: torch.Tensor,

133 + w1: torch.Tensor,

134 + w2: torch.Tensor,

135 + topk_weights: torch.Tensor,

136 + topk_ids: torch.Tensor,

137 + w1_scale: torch.Tensor,

138 + w2_scale: torch.Tensor,

139 + a_fp8_scale: torch.Tensor,

140 + num_repeats: int,

141 + ):

142 for _ in range(num_repeats):

134 - fused_experts(a,

135 - w1,

136 - w2,

137 - topk_weights,

138 - topk_ids,

139 - use_fp8_w8a8=True,

140 - w1_scale=w1_scale,

141 - w2_scale=w2_scale,

142 - a1_scale=a_fp8_scale)

143 -

144 - def run_cutlass_moe_fp4(a: torch.Tensor, w1_fp4: torch.Tensor,

145 - w2_fp4: torch.Tensor, w1_blockscale: torch.Tensor,

146 - w2_blockscale: torch.Tensor, w1_gs: torch.Tensor,

147 - w2_gs: torch.Tensor, a1_gs: torch.Tensor,

148 - a2_gs: torch.Tensor, topk_weights: torch.Tensor,

149 - topk_ids: torch.Tensor, m: int, n: int, k: int,

150 - e: int, device: torch.device, num_repeats: int):

143 + fused_experts(

144 + a,

145 + w1,

146 + w2,

147 + topk_weights,

148 + topk_ids,

149 + use_fp8_w8a8=True,

150 + w1_scale=w1_scale,

151 + w2_scale=w2_scale,

152 + a1_scale=a_fp8_scale,

153 + )

154 +

155 + def run_cutlass_moe_fp4(

156 + a: torch.Tensor,

157 + w1_fp4: torch.Tensor,

158 + w2_fp4: torch.Tensor,

159 + w1_blockscale: torch.Tensor,

160 + w2_blockscale: torch.Tensor,

161 + w1_gs: torch.Tensor,

162 + w2_gs: torch.Tensor,

163 + a1_gs: torch.Tensor,

164 + a2_gs: torch.Tensor,

165 + topk_weights: torch.Tensor,

166 + topk_ids: torch.Tensor,

167 + m: int,

168 + n: int,

169 + k: int,

170 + e: int,

171 + device: torch.device,

172 + num_repeats: int,

173 + ):

174 for _ in range(num_repeats):

175 with nvtx.annotate("cutlass_moe_fp4", color="green"):

153 - cutlass_moe_fp4(a=a,

154 - a1_gscale=a1_gs,

155 - a2_gscale=a2_gs,

156 - w1_fp4=w1_fp4,

157 - w1_blockscale=w1_blockscale,

158 - w1_alphas=w1_gs,

159 - w2_fp4=w2_fp4,

160 - w2_blockscale=w2_blockscale,

161 - w2_alphas=w2_gs,

162 - topk_weights=topk_weights,

163 - topk_ids=topk_ids,

164 - m=m,

165 - n=n,

166 - k=k,

167 - e=num_experts,

168 - device=device)

176 + cutlass_moe_fp4(

177 + a=a,

178 + a1_gscale=a1_gs,

179 + a2_gscale=a2_gs,

180 + w1_fp4=w1_fp4,

181 + w1_blockscale=w1_blockscale,

182 + w1_alphas=w1_gs,

183 + w2_fp4=w2_fp4,

184 + w2_blockscale=w2_blockscale,

185 + w2_alphas=w2_gs,

186 + topk_weights=topk_weights,

187 + topk_ids=topk_ids,

188 + m=m,

189 + n=n,

190 + k=k,

191 + e=num_experts,

192 + device=device,

193 + )

194

195 def run_cutlass_from_graph(

171 - a: torch.Tensor, a1_gscale: torch.Tensor, w1_fp4: torch.Tensor,

172 - w1_blockscale: torch.Tensor, w1_alphas: torch.Tensor,

173 - a2_gscale: torch.Tensor, w2_fp4: torch.Tensor,

174 - w2_blockscale: torch.Tensor, w2_alphas: torch.Tensor,

175 - topk_weights: torch.Tensor, topk_ids: torch.Tensor, m: int, n: int,

176 - k: int, e: int, device: torch.device):

196 + a: torch.Tensor,

197 + a1_gscale: torch.Tensor,

198 + w1_fp4: torch.Tensor,

199 + w1_blockscale: torch.Tensor,

200 + w1_alphas: torch.Tensor,

201 + a2_gscale: torch.Tensor,

202 + w2_fp4: torch.Tensor,

203 + w2_blockscale: torch.Tensor,

204 + w2_alphas: torch.Tensor,

205 + topk_weights: torch.Tensor,

206 + topk_ids: torch.Tensor,

207 + m: int,

208 + n: int,

209 + k: int,

210 + e: int,

211 + device: torch.device,

212 + ):

213 with set_current_vllm_config(

178 - VllmConfig(parallel_config=ParallelConfig(

179 - pipeline_parallel_size=1))):

180 - return cutlass_moe_fp4(a=a,

181 - a1_gscale=a1_gs,

182 - w1_fp4=w1_fp4,

183 - w1_blockscale=w1_blockscale,

184 - w1_alphas=w1_alphas,

185 - a2_gscale=a2_gs,

186 - w2_fp4=w2_fp4,

187 - w2_blockscale=w2_blockscale,

188 - w2_alphas=w2_alphas,

189 - topk_weights=topk_weights,

190 - topk_ids=topk_ids,

191 - m=m,

192 - n=n,

193 - k=k,

194 - e=num_experts,

195 - device=device)

196 -

197 - def run_triton_from_graph(a: torch.Tensor, w1: torch.Tensor,

198 - w2: torch.Tensor, topk_weights: torch.Tensor,

199 - topk_ids: torch.Tensor, w1_scale: torch.Tensor,

200 - w2_scale: torch.Tensor,

201 - a_fp8_scale: torch.Tensor):

214 + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))

215 + ):

216 + return cutlass_moe_fp4(

217 + a=a,

218 + a1_gscale=a1_gs,

219 + w1_fp4=w1_fp4,

220 + w1_blockscale=w1_blockscale,

221 + w1_alphas=w1_alphas,

222 + a2_gscale=a2_gs,

223 + w2_fp4=w2_fp4,

224 + w2_blockscale=w2_blockscale,

225 + w2_alphas=w2_alphas,

226 + topk_weights=topk_weights,

227 + topk_ids=topk_ids,

228 + m=m,

229 + n=n,

230 + k=k,

231 + e=num_experts,

232 + device=device,

233 + )

234 +

235 + def run_triton_from_graph(

236 + a: torch.Tensor,

237 + w1: torch.Tensor,

238 + w2: torch.Tensor,

239 + topk_weights: torch.Tensor,

240 + topk_ids: torch.Tensor,

241 + w1_scale: torch.Tensor,

242 + w2_scale: torch.Tensor,

243 + a_fp8_scale: torch.Tensor,

244 + ):

245 with set_current_vllm_config(

203 - VllmConfig(parallel_config=ParallelConfig(

204 - pipeline_parallel_size=1))):

205 - return fused_experts(a,

206 - w1,

207 - w2,

208 - topk_weights,

209 - topk_ids,

210 - use_fp8_w8a8=True,

211 - w1_scale=w1_scale,

212 - w2_scale=w2_scale,

213 - a1_scale=a_fp8_scale)

246 + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))

247 + ):

248 + return fused_experts(

249 + a,

250 + w1,

251 + w2,

252 + topk_weights,

253 + topk_ids,

254 + use_fp8_w8a8=True,

255 + w1_scale=w1_scale,

256 + w2_scale=w2_scale,

257 + a1_scale=a_fp8_scale,

258 + )

259

260 def replay_graph(graph, num_repeats):

261 for _ in range(num_repeats):

@@ -220,30 +265,39 @@ def bench_run(results: list[benchmark.Measurement], model: str,

265 cutlass_stream = torch.cuda.Stream()

266 cutlass_graph = torch.cuda.CUDAGraph()

267 with torch.cuda.graph(cutlass_graph, stream=cutlass_stream):

223 - run_cutlass_from_graph(a=a,

224 - a1_gscale=a1_gs,

225 - w1_fp4=w1_fp4,

226 - w1_blockscale=w1_blockscale,

227 - w1_alphas=w1_gs,

228 - a2_gscale=a2_gs,

229 - w2_fp4=w2_fp4,

230 - w2_blockscale=w2_blockscale,

231 - w2_alphas=w2_gs,

232 - topk_weights=topk_weights,

233 - topk_ids=topk_ids,

234 - m=m,

235 - n=n,

236 - k=k,

237 - e=num_experts,

238 - device=device)

268 + run_cutlass_from_graph(

269 + a=a,

270 + a1_gscale=a1_gs,

271 + w1_fp4=w1_fp4,

272 + w1_blockscale=w1_blockscale,

273 + w1_alphas=w1_gs,

274 + a2_gscale=a2_gs,

275 + w2_fp4=w2_fp4,

276 + w2_blockscale=w2_blockscale,

277 + w2_alphas=w2_gs,

278 + topk_weights=topk_weights,

279 + topk_ids=topk_ids,

280 + m=m,

281 + n=n,

282 + k=k,

283 + e=num_experts,

284 + device=device,

285 + )

286 torch.cuda.synchronize()

287

288 triton_stream = torch.cuda.Stream()

289 triton_graph = torch.cuda.CUDAGraph()

290 with torch.cuda.graph(triton_graph, stream=triton_stream):

244 - run_triton_from_graph(a, w1_fp8q_notransp, w2_fp8q_notransp,

245 - topk_weights, topk_ids, w1_fp8scale, w2_fp8scale,

246 - a_fp8_scale)

291 + run_triton_from_graph(

292 + a,

293 + w1_fp8q_notransp,

294 + w2_fp8q_notransp,

295 + topk_weights,

296 + topk_ids,

297 + w1_fp8scale,

298 + w2_fp8scale,

299 + a_fp8_scale,

300 + )

301 torch.cuda.synchronize()

302

303 min_run_time = 5

@@ -290,18 +344,27 @@ def bench_run(results: list[benchmark.Measurement], model: str,

344 }

345

346 # Warmup

293 - run_triton_moe(a, w1_fp8q_notransp, w2_fp8q_notransp, topk_weights,

294 - topk_ids, w1_fp8scale, w2_fp8scale, a_fp8_scale, num_warmup)

347 + run_triton_moe(

348 + a,

349 + w1_fp8q_notransp,

350 + w2_fp8q_notransp,

351 + topk_weights,

352 + topk_ids,

353 + w1_fp8scale,

354 + w2_fp8scale,

355 + a_fp8_scale,

356 + num_warmup,

357 + )

358

359 results.append(

360 benchmark.Timer(

298 - stmt=

299 - "run_triton_moe(a, w1_fp8q_notransp, w2_fp8q_notransp, topk_weights, topk_ids, w1_fp8scale, w2_fp8scale, a_fp8_scale, num_runs)", # noqa: E501

361 + stmt="run_triton_moe(a, w1_fp8q_notransp, w2_fp8q_notransp, topk_weights, topk_ids, w1_fp8scale, w2_fp8scale, a_fp8_scale, num_runs)", # noqa: E501

362 globals=globals,

363 label=label,

364 sub_label=sub_label,

365 description="triton_moe",

304 - ).blocked_autorange(min_run_time=min_run_time))

366 + ).blocked_autorange(min_run_time=min_run_time)

367 + )

368

369 # Warmup

370 replay_graph(triton_graph, num_warmup)

@@ -313,23 +376,40 @@ def bench_run(results: list[benchmark.Measurement], model: str,

376 label=label,

377 sub_label=sub_label,

378 description="triton_moe_cuda_graphs",

316 - ).blocked_autorange(min_run_time=min_run_time))

379 + ).blocked_autorange(min_run_time=min_run_time)

380 + )

381

382 # Warmup

383

320 - run_cutlass_moe_fp4(a, w1_fp4, w2_fp4, w1_blockscale, w2_blockscale, w1_gs,

321 - w2_gs, a1_gs, a2_gs, topk_weights, topk_ids, m, n, k,

322 - num_experts, device, num_warmup)

384 + run_cutlass_moe_fp4(

385 + a,

386 + w1_fp4,

387 + w2_fp4,

388 + w1_blockscale,

389 + w2_blockscale,

390 + w1_gs,

391 + w2_gs,

392 + a1_gs,

393 + a2_gs,

394 + topk_weights,

395 + topk_ids,

396 + m,

397 + n,

398 + k,

399 + num_experts,

400 + device,

401 + num_warmup,

402 + )

403

404 results.append(

405 benchmark.Timer(

326 - stmt=

327 - "run_cutlass_moe_fp4(a, w1_fp4, w2_fp4, w1_blockscale, w2_blockscale, w1_alphas, w2_alphas, a1_gscale, a2_gscale, topk_weights, topk_ids, m, n, k, e, device, num_runs)", # noqa: E501

406 + stmt="run_cutlass_moe_fp4(a, w1_fp4, w2_fp4, w1_blockscale, w2_blockscale, w1_alphas, w2_alphas, a1_gscale, a2_gscale, topk_weights, topk_ids, m, n, k, e, device, num_runs)", # noqa: E501

407 globals=globals,

408 label=label,

409 sub_label=sub_label,

410 description="cutlass_moe_fp4",

332 - ).blocked_autorange(min_run_time=min_run_time))

411 + ).blocked_autorange(min_run_time=min_run_time)

412 + )

413

414 # Warmup

415 replay_graph(cutlass_graph, num_warmup)

@@ -341,7 +421,8 @@ def bench_run(results: list[benchmark.Measurement], model: str,

421 label=label,

422 sub_label=sub_label,

423 description="cutlass_moe_fp4_cuda_graphs",

344 - ).blocked_autorange(min_run_time=min_run_time))

424 + ).blocked_autorange(min_run_time=min_run_time)

425 + )

426

427

428 def main(args):

@@ -369,8 +450,15 @@ def main(args):

450 for per_out_ch in PER_OUT_CH_OPTS:

451 for size_m in args.batch_sizes:

452 mkn = (size_m, size_k, size_n)

372 - bench_run(results, model, num_experts, topk,

373 - per_act_token, per_out_ch, mkn)

453 + bench_run(

454 + results,

455 + model,

456 + num_experts,

457 + topk,

458 + per_act_token,

459 + per_out_ch,

460 + mkn,

461 + )

462

463 compare = benchmark.Compare(results)

464 compare.print()

@@ -378,8 +466,8 @@ def main(args):

466

467 if __name__ == "__main__":

468 parser = FlexibleArgumentParser(

381 - description="Benchmark NVFP4 CUTLASS MOE across specified "

382 - "models/shapes/batches")

469 + description="Benchmark NVFP4 CUTLASS MOE across specified models/shapes/batches"

470 + )

471 parser.add_argument(

472 "--models",

473 nargs="+",

@@ -387,21 +475,14 @@ if __name__ == "__main__":

475 default=DEFAULT_MODELS,

476 choices=WEIGHT_SHAPES_MOE.keys(),

477 )

390 - parser.add_argument("--tp-sizes",

391 - nargs="+",

392 - type=int,

393 - default=DEFAULT_TP_SIZES)

394 - parser.add_argument("--batch-sizes",

395 - nargs="+",

396 - type=int,

397 - default=DEFAULT_BATCH_SIZES)

478 + parser.add_argument("--tp-sizes", nargs="+", type=int, default=DEFAULT_TP_SIZES)

479 + parser.add_argument(

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

481 + )

482 parser.add_argument("--limit-k", nargs="+", type=int, default=[])

483 parser.add_argument("--limit-n", nargs="+", type=int, default=[])

484 parser.add_argument("--limit-num-groups", nargs="+", type=int, default=[])

401 - parser.add_argument("--limit-per-act-token",

402 - nargs="+",

403 - type=int,

404 - default=[])

485 + parser.add_argument("--limit-per-act-token", nargs="+", type=int, default=[])

486 parser.add_argument("--limit-per-out-ch", nargs="+", type=int, default=[])

487

488 args = parser.parse_args()

benchmarks/kernels/benchmark_grouped_gemm_cutlass.py

+219 -144

@@ -6,14 +6,18 @@ from benchmark_shapes import WEIGHT_SHAPES_MOE

7 from vllm import _custom_ops as ops

8 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

9 -from vllm.model_executor.layers.fused_moe.fused_moe import (cutlass_moe_fp8,

10 - fused_experts,

11 - fused_topk)

9 +from vllm.model_executor.layers.fused_moe.fused_moe import (

10 + cutlass_moe_fp8,

11 + fused_experts,

12 + fused_topk,

13 +)

14 from vllm.utils import FlexibleArgumentParser

16 DEFAULT_MODELS = [

15 - "nm-testing/Mixtral-8x7B-Instruct-v0.1", "nm-testing/deepseekv2-lite",

16 - "ibm-granite/granite-3.0-1b-a400m", "ibm-granite/granite-3.0-3b-a800m"

17 + "nm-testing/Mixtral-8x7B-Instruct-v0.1",

18 + "nm-testing/deepseekv2-lite",

19 + "ibm-granite/granite-3.0-1b-a400m",

20 + "ibm-granite/granite-3.0-3b-a800m",

21 ]

22 DEFAULT_BATCH_SIZES = [1, 4, 8, 16, 32, 64, 128, 256, 512]

23 DEFAULT_TP_SIZES = [1]

@@ -24,19 +28,27 @@ PER_OUT_CH_OPTS = [False]

29 def to_fp8(tensor: torch.Tensor):

30 finfo = torch.finfo(torch.float8_e4m3fn)

27 - return torch.round(tensor.clamp(

28 - min=finfo.min, max=finfo.max)).to(dtype=torch.float8_e4m3fn)

31 + return torch.round(tensor.clamp(min=finfo.min, max=finfo.max)).to(

32 + dtype=torch.float8_e4m3fn

33 + )

31 -def bench_run(results: list[benchmark.Measurement], model: str,

32 - num_experts: int, topk: int, per_act_token: bool,

33 - per_out_ch: bool, mkn: tuple[int, int, int]):

36 +def bench_run(

37 + results: list[benchmark.Measurement],

38 + model: str,

39 + num_experts: int,

40 + topk: int,

41 + per_act_token: bool,

42 + per_out_ch: bool,

43 + mkn: tuple[int, int, int],

44 +):

45 label = "Quant Matmul"

47 sub_label = (

37 - "{}, num_experts={}, topk={}, per_act_token={} per_out_ch={}, "

38 - "MKN=({})".format(model, num_experts, topk, per_act_token, per_out_ch,

39 - mkn))

48 + "{}, num_experts={}, topk={}, per_act_token={} per_out_ch={}, MKN=({})".format(

49 + model, num_experts, topk, per_act_token, per_out_ch, mkn

50 + )

51 + )

53 print(f"Testing: {sub_label}")

@@ -50,35 +62,17 @@ def bench_run(results: list[benchmark.Measurement], model: str,

63 _, a_scale = ops.scaled_fp8_quant(a)

53 - w1_q = torch.empty((num_experts, 2 * n, k),

54 - device="cuda",

55 - dtype=torch.float8_e4m3fn)

56 - w2_q = torch.empty((num_experts, k, n),

57 - device="cuda",

58 - dtype=torch.float8_e4m3fn)

59 - w1_scale = torch.empty((num_experts, 1, 1),

60 - device="cuda",

61 - dtype=torch.float32)

62 - w2_scale = torch.empty((num_experts, 1, 1),

63 - device="cuda",

64 - dtype=torch.float32)

65 -

66 - ab_strides1 = torch.full((num_experts, ),

67 - k,

68 - device="cuda",

69 - dtype=torch.int64)

70 - c_strides1 = torch.full((num_experts, ),

71 - 2 * n,

72 - device="cuda",

73 - dtype=torch.int64)

74 - ab_strides2 = torch.full((num_experts, ),

75 - n,

76 - device="cuda",

77 - dtype=torch.int64)

78 - c_strides2 = torch.full((num_experts, ),

79 - k,

80 - device="cuda",

81 - dtype=torch.int64)

65 + w1_q = torch.empty(

66 + (num_experts, 2 * n, k), device="cuda", dtype=torch.float8_e4m3fn

67 + )

68 + w2_q = torch.empty((num_experts, k, n), device="cuda", dtype=torch.float8_e4m3fn)

69 + w1_scale = torch.empty((num_experts, 1, 1), device="cuda", dtype=torch.float32)

70 + w2_scale = torch.empty((num_experts, 1, 1), device="cuda", dtype=torch.float32)

71 +

72 + ab_strides1 = torch.full((num_experts,), k, device="cuda", dtype=torch.int64)

73 + c_strides1 = torch.full((num_experts,), 2 * n, device="cuda", dtype=torch.int64)

74 + ab_strides2 = torch.full((num_experts,), n, device="cuda", dtype=torch.int64)

75 + c_strides2 = torch.full((num_experts,), k, device="cuda", dtype=torch.int64)

77 for expert in range(num_experts):

78 w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(w1[expert])

@@ -91,82 +85,120 @@ def bench_run(results: list[benchmark.Measurement], model: str,

85 score = torch.randn((m, num_experts), device="cuda", dtype=dtype)

87 topk_weights, topk_ids, token_expert_indices = fused_topk(

94 - a, score, topk, renormalize=False)

88 + a, score, topk, renormalize=False

89 + )

96 - def run_triton_moe(a: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor,

97 - topk_weights: torch.Tensor, topk_ids: torch.Tensor,

98 - w1_scale: torch.Tensor, w2_scale: torch.Tensor,

99 - a_scale: torch.Tensor, num_repeats: int):

91 + def run_triton_moe(

92 + a: torch.Tensor,

93 + w1: torch.Tensor,

94 + w2: torch.Tensor,

95 + topk_weights: torch.Tensor,

96 + topk_ids: torch.Tensor,

97 + w1_scale: torch.Tensor,

98 + w2_scale: torch.Tensor,

99 + a_scale: torch.Tensor,

100 + num_repeats: int,

101 + ):

102 for _ in range(num_repeats):

101 - fused_experts(a,

102 - w1,

103 - w2,

104 - topk_weights,

105 - topk_ids,

106 - use_fp8_w8a8=True,

107 - w1_scale=w1_scale,

108 - w2_scale=w2_scale,

109 - a1_scale=a_scale)

110 -

111 - def run_cutlass_moe(a: torch.Tensor, a_scale: torch.Tensor,

112 - w1: torch.Tensor, w2: torch.Tensor,

113 - w1_scale: torch.Tensor, w2_scale: torch.Tensor,

114 - topk_weights: torch.Tensor, topk_ids: torch.Tensor,

115 - ab_strides1: torch.Tensor, c_strides1: torch.Tensor,

116 - ab_strides2: torch.Tensor, c_strides2: torch.Tensor,

117 - num_repeats: int):

103 + fused_experts(

104 + a,

105 + w1,

106 + w2,

107 + topk_weights,

108 + topk_ids,

109 + use_fp8_w8a8=True,

110 + w1_scale=w1_scale,

111 + w2_scale=w2_scale,

112 + a1_scale=a_scale,

113 + )

114 +

115 + def run_cutlass_moe(

116 + a: torch.Tensor,

117 + a_scale: torch.Tensor,

118 + w1: torch.Tensor,

119 + w2: torch.Tensor,

120 + w1_scale: torch.Tensor,

121 + w2_scale: torch.Tensor,

122 + topk_weights: torch.Tensor,

123 + topk_ids: torch.Tensor,

124 + ab_strides1: torch.Tensor,

125 + c_strides1: torch.Tensor,

126 + ab_strides2: torch.Tensor,

127 + c_strides2: torch.Tensor,

128 + num_repeats: int,

129 + ):

130 for _ in range(num_repeats):

119 - cutlass_moe_fp8(a,

120 - w1,

121 - w2,

122 - w1_scale,

123 - w2_scale,

124 - topk_weights,

125 - topk_ids,

126 - ab_strides1,

127 - c_strides1,

128 - ab_strides2,

129 - c_strides2,

130 - a1_scale=a_scale)

131 + cutlass_moe_fp8(

132 + a,

133 + w1,

134 + w2,

135 + w1_scale,

136 + w2_scale,

137 + topk_weights,

138 + topk_ids,

139 + ab_strides1,

140 + c_strides1,

141 + ab_strides2,

142 + c_strides2,

143 + a1_scale=a_scale,

144 + )

145

146 def run_cutlass_from_graph(

133 - a: torch.Tensor, a_scale: torch.Tensor, w1_q: torch.Tensor,

134 - w2_q: torch.Tensor, w1_scale: torch.Tensor, w2_scale: torch.Tensor,

135 - topk_weights: torch.Tensor, topk_ids: torch.Tensor,

136 - ab_strides1: torch.Tensor, c_strides1: torch.Tensor,

137 - ab_strides2: torch.Tensor, c_strides2: torch.Tensor):

147 + a: torch.Tensor,

148 + a_scale: torch.Tensor,

149 + w1_q: torch.Tensor,

150 + w2_q: torch.Tensor,

151 + w1_scale: torch.Tensor,

152 + w2_scale: torch.Tensor,

153 + topk_weights: torch.Tensor,

154 + topk_ids: torch.Tensor,

155 + ab_strides1: torch.Tensor,

156 + c_strides1: torch.Tensor,

157 + ab_strides2: torch.Tensor,

158 + c_strides2: torch.Tensor,

159 + ):

160 with set_current_vllm_config(

139 - VllmConfig(parallel_config=ParallelConfig(

140 - pipeline_parallel_size=1))):

141 - return cutlass_moe_fp8(a,

142 - w1_q,

143 - w2_q,

144 - w1_scale,

145 - w2_scale,

146 - topk_weights,

147 - topk_ids,

148 - ab_strides1,

149 - c_strides1,

150 - ab_strides2,

151 - c_strides2,

152 - a1_scale=a_scale)

153 -

154 - def run_triton_from_graph(a: torch.Tensor, w1: torch.Tensor,

155 - w2: torch.Tensor, topk_weights: torch.Tensor,

156 - topk_ids: torch.Tensor, w1_scale: torch.Tensor,

157 - w2_scale: torch.Tensor, a_scale: torch.Tensor):

161 + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))

162 + ):

163 + return cutlass_moe_fp8(

164 + a,

165 + w1_q,

166 + w2_q,

167 + w1_scale,

168 + w2_scale,

169 + topk_weights,

170 + topk_ids,

171 + ab_strides1,

172 + c_strides1,

173 + ab_strides2,

174 + c_strides2,

175 + a1_scale=a_scale,

176 + )

177 +

178 + def run_triton_from_graph(

179 + a: torch.Tensor,

180 + w1: torch.Tensor,

181 + w2: torch.Tensor,

182 + topk_weights: torch.Tensor,

183 + topk_ids: torch.Tensor,

184 + w1_scale: torch.Tensor,

185 + w2_scale: torch.Tensor,

186 + a_scale: torch.Tensor,

187 + ):

188 with set_current_vllm_config(

159 - VllmConfig(parallel_config=ParallelConfig(

160 - pipeline_parallel_size=1))):

161 - return fused_experts(a,

162 - w1,

163 - w2,

164 - topk_weights,

165 - topk_ids,

166 - use_fp8_w8a8=True,

167 - w1_scale=w1_scale,

168 - w2_scale=w2_scale,

169 - a1_scale=a_scale)

189 + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))

190 + ):

191 + return fused_experts(

192 + a,

193 + w1,

194 + w2,

195 + topk_weights,

196 + topk_ids,

197 + use_fp8_w8a8=True,

198 + w1_scale=w1_scale,

199 + w2_scale=w2_scale,

200 + a1_scale=a_scale,

201 + )

202

203 def replay_graph(graph, num_repeats):

204 for _ in range(num_repeats):

@@ -176,16 +208,35 @@ def bench_run(results: list[benchmark.Measurement], model: str,

208 cutlass_stream = torch.cuda.Stream()

209 cutlass_graph = torch.cuda.CUDAGraph()

210 with torch.cuda.graph(cutlass_graph, stream=cutlass_stream):

179 - run_cutlass_from_graph(a, a_scale, w1_q, w2_q, w1_scale, w2_scale,

180 - topk_weights, topk_ids, ab_strides1, c_strides1,

181 - ab_strides2, c_strides2)

211 + run_cutlass_from_graph(

212 + a,

213 + a_scale,

214 + w1_q,

215 + w2_q,

216 + w1_scale,

217 + w2_scale,

218 + topk_weights,

219 + topk_ids,

220 + ab_strides1,

221 + c_strides1,

222 + ab_strides2,

223 + c_strides2,

224 + )

225 torch.cuda.synchronize()

226

227 triton_stream = torch.cuda.Stream()

228 triton_graph = torch.cuda.CUDAGraph()

229 with torch.cuda.graph(triton_graph, stream=triton_stream):

187 - run_triton_from_graph(a, w1_q_notransp, w2_q_notransp, topk_weights,

188 - topk_ids, w1_scale, w2_scale, a_scale)

230 + run_triton_from_graph(

231 + a,

232 + w1_q_notransp,

233 + w2_q_notransp,

234 + topk_weights,

235 + topk_ids,

236 + w1_scale,

237 + w2_scale,

238 + a_scale,

239 + )

240 torch.cuda.synchronize()

241

242 min_run_time = 5

@@ -225,18 +276,27 @@ def bench_run(results: list[benchmark.Measurement], model: str,

276 }

277

278 # Warmup

228 - run_triton_moe(a, w1_q_notransp, w2_q_notransp, topk_weights, topk_ids,

229 - w1_scale, w2_scale, a_scale, num_warmup)

279 + run_triton_moe(

280 + a,

281 + w1_q_notransp,

282 + w2_q_notransp,

283 + topk_weights,

284 + topk_ids,

285 + w1_scale,

286 + w2_scale,

287 + a_scale,

288 + num_warmup,

289 + )

290

291 results.append(

292 benchmark.Timer(

233 - stmt=

234 - "run_triton_moe(a, w1_q_notransp, w2_q_notransp, topk_weights, topk_ids, w1_scale, w2_scale, a_scale, num_runs)", # noqa: E501

293 + stmt="run_triton_moe(a, w1_q_notransp, w2_q_notransp, topk_weights, topk_ids, w1_scale, w2_scale, a_scale, num_runs)", # noqa: E501

294 globals=globals,

295 label=label,

296 sub_label=sub_label,

297 description="triton_moe",

239 - ).blocked_autorange(min_run_time=min_run_time))

298 + ).blocked_autorange(min_run_time=min_run_time)

299 + )

300

301 # Warmup

302 replay_graph(triton_graph, num_warmup)

@@ -248,22 +308,35 @@ def bench_run(results: list[benchmark.Measurement], model: str,

308 label=label,

309 sub_label=sub_label,

310 description="triton_moe_cuda_graphs",

251 - ).blocked_autorange(min_run_time=min_run_time))

311 + ).blocked_autorange(min_run_time=min_run_time)

312 + )

313

314 # Warmup

254 - run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, topk_weights,

255 - topk_ids, ab_strides1, c_strides1, ab_strides2, c_strides2,

256 - num_warmup)

315 + run_cutlass_moe(

316 + a,

317 + a_scale,

318 + w1_q,

319 + w2_q,

320 + w1_scale,

321 + w2_scale,

322 + topk_weights,

323 + topk_ids,

324 + ab_strides1,

325 + c_strides1,

326 + ab_strides2,

327 + c_strides2,

328 + num_warmup,

329 + )

330

331 results.append(

332 benchmark.Timer(

260 - stmt=

261 - "run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, topk_weights, topk_ids, ab_strides1, c_strides1, ab_strides2, c_strides2, num_runs)", # noqa: E501

333 + stmt="run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, topk_weights, topk_ids, ab_strides1, c_strides1, ab_strides2, c_strides2, num_runs)", # noqa: E501

334 globals=globals,

335 label=label,

336 sub_label=sub_label,

337 description="grouped_gemm_moe",

266 - ).blocked_autorange(min_run_time=min_run_time))

338 + ).blocked_autorange(min_run_time=min_run_time)

339 + )

340

341 # Warmup

342 replay_graph(cutlass_graph, num_warmup)

@@ -275,7 +348,8 @@ def bench_run(results: list[benchmark.Measurement], model: str,

348 label=label,

349 sub_label=sub_label,

350 description="grouped_gemm_moe_cuda_graphs",

278 - ).blocked_autorange(min_run_time=min_run_time))

351 + ).blocked_autorange(min_run_time=min_run_time)

352 + )

353

354

355 def main(args):

@@ -303,8 +377,15 @@ def main(args):

377 for per_out_ch in PER_OUT_CH_OPTS:

378 for size_m in DEFAULT_BATCH_SIZES:

379 mkn = (size_m, size_k, size_n)

306 - bench_run(results, model, num_experts, topk,

307 - per_act_token, per_out_ch, mkn)

380 + bench_run(

381 + results,

382 + model,

383 + num_experts,

384 + topk,

385 + per_act_token,

386 + per_out_ch,

387 + mkn,

388 + )

389

390 compare = benchmark.Compare(results)

391 compare.print()

@@ -312,7 +393,8 @@ def main(args):

393

394 if __name__ == "__main__":

395 parser = FlexibleArgumentParser(

315 - description="Benchmark Marlin across specified models/shapes/batches")

396 + description="Benchmark Marlin across specified models/shapes/batches"

397 + )

398 parser.add_argument(

399 "--models",

400 nargs="+",

@@ -320,21 +402,14 @@ if __name__ == "__main__":

402 default=DEFAULT_MODELS,

403 choices=WEIGHT_SHAPES_MOE.keys(),

404 )

323 - parser.add_argument("--tp-sizes",

324 - nargs="+",

325 - type=int,

326 - default=DEFAULT_TP_SIZES)

327 - parser.add_argument("--batch-sizes",

328 - nargs="+",

329 - type=int,

330 - default=DEFAULT_BATCH_SIZES)

405 + parser.add_argument("--tp-sizes", nargs="+", type=int, default=DEFAULT_TP_SIZES)

406 + parser.add_argument(

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

408 + )

409 parser.add_argument("--limit-k", nargs="+", type=int, default=[])

410 parser.add_argument("--limit-n", nargs="+", type=int, default=[])

411 parser.add_argument("--limit-num-groups", nargs="+", type=int, default=[])

334 - parser.add_argument("--limit-per-act-token",

335 - nargs="+",

336 - type=int,

337 - default=[])

412 + parser.add_argument("--limit-per-act-token", nargs="+", type=int, default=[])

413 parser.add_argument("--limit-per-out-ch", nargs="+", type=int, default=[])

414

415 args = parser.parse_args()

benchmarks/kernels/benchmark_layernorm.py

+32 -28

@@ -10,14 +10,16 @@ from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, FlexibleArgumentParser

12 @torch.inference_mode()

13 -def main(num_tokens: int,

14 - hidden_size: int,

15 - add_residual: bool,

16 - dtype: torch.dtype,

17 - seed: int = 0,

18 - do_profile: bool = False,

19 - num_warmup_iters: int = 5,

20 - num_iters: int = 100) -> None:

13 +def main(

14 + num_tokens: int,

15 + hidden_size: int,

16 + add_residual: bool,

17 + dtype: torch.dtype,

18 + seed: int = 0,

19 + do_profile: bool = False,

20 + num_warmup_iters: int = 5,

21 + num_iters: int = 100,

22 +) -> None:

23 current_platform.seed_everything(seed)

24 torch.set_default_device("cuda")

@@ -56,33 +58,35 @@ def main(num_tokens: int,

58 print(f"Kernel running time: {latency * 1000000:.3f} us")

59 -if __name__ == '__main__':

60 - parser = FlexibleArgumentParser(

61 - description="Benchmark the layernorm kernel.")

61 +if __name__ == "__main__":

62 + parser = FlexibleArgumentParser(description="Benchmark the layernorm kernel.")

63 parser.add_argument("--num-tokens", type=int, default=4096)

64 parser.add_argument("--hidden-size", type=int, default=8192)

65 parser.add_argument("--add-residual", action="store_true")

65 - parser.add_argument("--dtype",

66 - type=str,

67 - choices=["half", "bfloat16", "float"],

68 - default="half")

66 + parser.add_argument(

67 + "--dtype", type=str, choices=["half", "bfloat16", "float"], default="half"

68 + )

69 parser.add_argument("--seed", type=int, default=0)

70 parser.add_argument("--profile", action="store_true")

71 parser.add_argument("--num-warmup-iters", type=int, default=5)

72 - parser.add_argument("--num-iters",

73 - type=int,

74 - default=100,

75 - help="Number of benchmark iterations. "

76 - "If --profile is set, this number is ignored")

72 + parser.add_argument(

73 + "--num-iters",

74 + type=int,

75 + default=100,

76 + help="Number of benchmark iterations. "

77 + "If --profile is set, this number is ignored",

78 + )

80 args = parser.parse_args()

81 print(args)

81 - main(num_tokens=args.num_tokens,

82 - hidden_size=args.hidden_size,

83 - add_residual=args.add_residual,

84 - dtype=STR_DTYPE_TO_TORCH_DTYPE[args.dtype],

85 - seed=args.seed,

86 - do_profile=args.profile,

87 - num_warmup_iters=args.num_warmup_iters,

88 - num_iters=args.num_iters)

83 + main(

84 + num_tokens=args.num_tokens,

85 + hidden_size=args.hidden_size,

86 + add_residual=args.add_residual,

87 + dtype=STR_DTYPE_TO_TORCH_DTYPE[args.dtype],

88 + seed=args.seed,

89 + do_profile=args.profile,

90 + num_warmup_iters=args.num_warmup_iters,

91 + num_iters=args.num_iters,

92 + )

benchmarks/kernels/benchmark_lora.py

+378 -275

@@ -20,18 +20,36 @@ from weight_shapes import WEIGHT_SHAPES

20 from vllm.triton_utils import HAS_TRITON

22 if HAS_TRITON:

23 - from vllm.lora.ops.triton_ops import (LoRAKernelMeta, lora_expand,

24 - lora_shrink)

25 - from vllm.lora.ops.triton_ops.utils import (_LORA_A_PTR_DICT,

26 - _LORA_B_PTR_DICT)

23 + from vllm.lora.ops.triton_ops import LoRAKernelMeta, lora_expand, lora_shrink

24 + from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT

26 from vllm.utils import FlexibleArgumentParser

28 DEFAULT_MODELS = list(WEIGHT_SHAPES.keys())

29 DEFAULT_TP_SIZES = [1]

30 DEFAULT_BATCH_SIZES = [

33 - 1, 16, 32, 64, 128, 192, 256, 320, 384, 448, 512, 640, 768, 896, 1024,

34 - 2048, 3072, 4096, 5120, 6144, 7168, 8192

31 + 1,

32 + 16,

33 + 32,

34 + 64,

35 + 128,

36 + 192,

37 + 256,

38 + 320,

39 + 384,

40 + 448,

41 + 512,

42 + 640,

43 + 768,

44 + 896,

45 + 1024,

46 + 2048,

47 + 3072,

48 + 4096,

49 + 5120,

50 + 6144,

51 + 7168,

52 + 8192,

53 ]

54 DEFAULT_HIDDEN_SIZES = [1024, 2048, 4096, 8192, 16384]

55 DEFAULT_LORA_RANKS = [16]

@@ -52,12 +70,9 @@ def dtype_to_str(dtype: torch.dtype):

70 raise ValueError(f"Unsupported dtype {dtype}")

55 -def make_rand_lora_weight_tensor(k: int,

56 - n: int,

57 - num_loras: int,

58 - dtype: torch.dtype,

59 - device: str = "cuda") -> torch.Tensor:

60 -

73 +def make_rand_lora_weight_tensor(

74 + k: int, n: int, num_loras: int, dtype: torch.dtype, device: str = "cuda"

75 +) -> torch.Tensor:

76 # LoRA weights column major

77 return torch.rand((num_loras, n, k), dtype=dtype).to(device)

@@ -78,18 +93,15 @@ def make_rand_tensors(

93 A = torch.rand(a_shape, dtype=a_dtype).to(device)

95 # LoRA weights column major

81 - Bs = [

82 - torch.rand(b_shape, dtype=b_dtype).to(device)

83 - for _ in range(num_slices)

84 - ]

96 + Bs = [torch.rand(b_shape, dtype=b_dtype).to(device) for _ in range(num_slices)]

98 C = torch.zeros(c_shape, dtype=c_dtype).to(device)

99 return A, Bs, C

100

101

90 -def make_prompt_lora_mapping(num_prompts: int, num_active_loras: int,

91 - sort_by_lora_id: bool,

92 - device: str) -> torch.Tensor:

102 +def make_prompt_lora_mapping(

103 + num_prompts: int, num_active_loras: int, sort_by_lora_id: bool, device: str

104 +) -> torch.Tensor:

105 """

106 All prompts are mapped to a LoRA ID in range [0, num_active_loras).

107 where 0 refers to first lora, 1 refers to second lora and so on.

@@ -97,9 +109,7 @@ def make_prompt_lora_mapping(num_prompts: int, num_active_loras: int,

109 assert num_active_loras > 0

110

111 if not sort_by_lora_id:

100 - return torch.randint(0,

101 - num_active_loras, (num_prompts, ),

102 - dtype=torch.long)

112 + return torch.randint(0, num_active_loras, (num_prompts,), dtype=torch.long)

113

114 # Divide LoRAs equally and in order.

115 part_size = num_prompts // num_active_loras

@@ -110,14 +120,18 @@ def make_prompt_lora_mapping(num_prompts: int, num_active_loras: int,

120 while len(prompt_lora_mapping) < num_prompts:

121 prompt_lora_mapping.extend([lora_id] * part_size)

122 lora_id = lora_id + 1 if lora_id + 1 < num_active_loras else lora_id

113 - return torch.tensor(prompt_lora_mapping[:num_prompts],

114 - dtype=torch.long,

115 - device=device)

116 -

117 -

118 -def make_token_lora_mapping(num_tokens: int, num_prompts: int,

119 - prompt_lora_mapping: torch.Tensor,

120 - seq_len_tensor: torch.Tensor, device: str):

123 + return torch.tensor(

124 + prompt_lora_mapping[:num_prompts], dtype=torch.long, device=device

125 + )

126 +

127 +

128 +def make_token_lora_mapping(

129 + num_tokens: int,

130 + num_prompts: int,

131 + prompt_lora_mapping: torch.Tensor,

132 + seq_len_tensor: torch.Tensor,

133 + device: str,

134 +):

135 """

136 Make token_lora_mapping from prompt_lora_mapping and seq_lens_tensor

137 """

@@ -136,11 +150,15 @@ def make_token_lora_mapping(num_tokens: int, num_prompts: int,

150 return torch.tensor(token_lora_mapping, dtype=torch.long, device=device)

151

152

139 -def ref_group_gemm(ref_out: torch.Tensor, input: torch.Tensor,

140 - lora_weights: list[torch.Tensor],

141 - seq_lens_cpu: torch.Tensor,

142 - prompt_lora_mapping_cpu: torch.Tensor, scaling: float,

143 - add_inputs: Optional[bool]):

153 +def ref_group_gemm(

154 + ref_out: torch.Tensor,

155 + input: torch.Tensor,

156 + lora_weights: list[torch.Tensor],

157 + seq_lens_cpu: torch.Tensor,

158 + prompt_lora_mapping_cpu: torch.Tensor,

159 + scaling: float,

160 + add_inputs: Optional[bool],

161 +):

162 """

163 Torch group gemm reference implementation to test correctness of

164 benchmarking operations.

@@ -149,7 +167,7 @@ def ref_group_gemm(ref_out: torch.Tensor, input: torch.Tensor,

167 out_list = []

168 current_offset = 0

169 for lora_index, b_length in zip(range(batches), seq_lens_cpu):

152 - x = input[current_offset:b_length + current_offset, :]

170 + x = input[current_offset : b_length + current_offset, :]

171 current_offset += b_length

172 w = lora_weights[prompt_lora_mapping_cpu[lora_index]]

173 result = torch.nn.functional.linear(x, w)

@@ -168,6 +186,7 @@ class OpType(Enum):

186 """

187 LoRA Ops to benchmark and its properties.

188 """

189 +

190 LORA_SHRINK = auto()

191 LORA_EXPAND = auto()

192

@@ -188,8 +207,9 @@ class OpType(Enum):

207 def num_slices(self) -> list[int]:

208 return [1, 2, 3]

209

191 - def mkn(self, batch_size: int, seq_length: int, hidden_size: int,

192 - lora_rank: int) -> tuple[int, int, int]:

210 + def mkn(

211 + self, batch_size: int, seq_length: int, hidden_size: int, lora_rank: int

212 + ) -> tuple[int, int, int]:

213 num_tokens = batch_size * seq_length

214 if self.is_shrink_fn():

215 m = num_tokens

@@ -203,7 +223,7 @@ class OpType(Enum):

223 return m, k, n

224

225 def matmul_dtypes(

206 - self, op_dtype: torch.dtype

226 + self, op_dtype: torch.dtype

227 ) -> tuple[torch.dtype, torch.dtype, torch.dtype]:

228 """

229 return a type, b type and c type for A x B = C

@@ -215,9 +235,14 @@ class OpType(Enum):

235 return torch.float32, op_dtype, op_dtype

236

237 def matmul_shapes(

218 - self, batch_size: int, seq_length: int, hidden_size: int,

219 - lora_rank: int, num_loras: int,

220 - num_slices: int) -> tuple[tuple[int], tuple[int], tuple[int]]:

238 + self,

239 + batch_size: int,

240 + seq_length: int,

241 + hidden_size: int,

242 + lora_rank: int,

243 + num_loras: int,

244 + num_slices: int,

245 + ) -> tuple[tuple[int], tuple[int], tuple[int]]:

246 """

247 Given num_slices, return the shapes of the A, B, and C matrices

248 in A x B = C, for the op_type

@@ -241,31 +266,38 @@ class OpType(Enum):

266

267 raise ValueError(f"Unrecognized optype {self}")

268

244 - def run_ref_group_gemm(self, output: torch.Tensor, input: torch.Tensor,

245 - lora_weights: list[torch.Tensor],

246 - **kwargs) -> Callable:

269 + def run_ref_group_gemm(

270 + self,

271 + output: torch.Tensor,

272 + input: torch.Tensor,

273 + lora_weights: list[torch.Tensor],

274 + **kwargs,

275 + ) -> Callable:

276 """Each benchmark operation expects the input, lora_weights and outputs

248 - in a slightly different format. Refer to self.matmul_shapes().

249 - run_ref_group_gemm accounts for those differences in executing a

250 - reference group gemm for correctness testing.

277 + in a slightly different format. Refer to self.matmul_shapes().

278 + run_ref_group_gemm accounts for those differences in executing a

279 + reference group gemm for correctness testing.

280 """

281 w_dtype = lora_weights[0].dtype

282 num_slices = len(lora_weights)

283 if self in [OpType.LORA_SHRINK]:

284 for slice_idx in range(num_slices):

256 - ref_group_gemm(ref_out=output[slice_idx, :],

257 - input=input,

258 - lora_weights=lora_weights[slice_idx],

259 - **kwargs)

285 + ref_group_gemm(

286 + ref_out=output[slice_idx, :],

287 + input=input,

288 + lora_weights=lora_weights[slice_idx],

289 + **kwargs,

290 + )

291 elif self in [OpType.LORA_EXPAND]:

292 hidden_size = lora_weights[0].shape[1]

293 for slice_idx in range(num_slices):

294 slice_offset = slice_idx * hidden_size

295 ref_group_gemm(

265 - ref_out=output[:, slice_offset:slice_offset + hidden_size],

296 + ref_out=output[:, slice_offset : slice_offset + hidden_size],

297 input=input[slice_idx].clone().to(dtype=w_dtype),

298 lora_weights=lora_weights[slice_idx],

268 - **kwargs)

299 + **kwargs,

300 + )

301 else:

302 raise ValueError(f"Unrecognized optype {self}")

303

@@ -275,6 +307,7 @@ class BenchmarkContext:

307 """

308 LoRA benchmark context

309 """

310 +

311 batch_size: int

312 hidden_size: int

313 num_loras: int

@@ -299,17 +332,18 @@ class BenchmarkContext:

332 return f"lora-{self.dtype}"

333

334 def bench_sublabel(self, op_type: OpType) -> str:

302 - m, k, n = op_type.mkn(self.batch_size, self.seq_length,

303 - self.hidden_size, self.lora_rank)

335 + m, k, n = op_type.mkn(

336 + self.batch_size, self.seq_length, self.hidden_size, self.lora_rank

337 + )

338 desc = {

305 - 'bs': self.batch_size,

306 - 'sl': self.seq_length,

307 - 'm': m,

308 - 'k': k,

309 - 'n': n,

310 - 'num_loras': self.num_loras,

311 - 'sort_by_lora': self.sort_by_lora_id,

312 - 'num_slices': self.num_slices,

339 + "bs": self.batch_size,

340 + "sl": self.seq_length,

341 + "m": m,

342 + "k": k,

343 + "n": n,

344 + "num_loras": self.num_loras,

345 + "sort_by_lora": self.sort_by_lora_id,

346 + "num_slices": self.num_slices,

347 }

348 return json.dumps(desc)

349

@@ -319,6 +353,7 @@ class BenchmarkTensors:

353 """

354 Input/Output tensors used for benchmarks

355 """

356 +

357 # matmul tensors

358 input: torch.Tensor

359 lora_weights_lst: list[torch.Tensor]

@@ -330,23 +365,29 @@ class BenchmarkTensors:

365 prompt_lora_mapping: torch.Tensor

366

367 def io_types(self) -> str:

333 - return (f"{dtype_to_str(self.input.dtype)}x"

334 - f"{dtype_to_str(self.lora_weights_lst[0].dtype)}=>"

335 - f"{dtype_to_str(self.output.dtype)}")

368 + return (

369 + f"{dtype_to_str(self.input.dtype)}x"

370 + f"{dtype_to_str(self.lora_weights_lst[0].dtype)}=>"

371 + f"{dtype_to_str(self.output.dtype)}"

372 + )

373

374 @staticmethod

338 - def make(ctx: BenchmarkContext,

339 - op_type: OpType,

340 - device: str = "cuda") -> "BenchmarkTensors":

341 -

375 + def make(

376 + ctx: BenchmarkContext, op_type: OpType, device: str = "cuda"

377 + ) -> "BenchmarkTensors":

378 # Make input / output matmul tensors.

379 a_shape, b_shape, c_shape = op_type.matmul_shapes(

344 - ctx.batch_size, ctx.seq_length, ctx.hidden_size, ctx.lora_rank,

345 - ctx.num_loras, ctx.num_slices)

380 + ctx.batch_size,

381 + ctx.seq_length,

382 + ctx.hidden_size,

383 + ctx.lora_rank,

384 + ctx.num_loras,

385 + ctx.num_slices,

386 + )

387 a_type, b_type, c_type = op_type.matmul_dtypes(ctx.dtype)

347 - input_tensor, lora_weights, output_tensor = \

348 - make_rand_tensors(a_shape, b_shape, c_shape, a_type, b_type, c_type,

349 - num_slices = ctx.num_slices)

388 + input_tensor, lora_weights, output_tensor = make_rand_tensors(

389 + a_shape, b_shape, c_shape, a_type, b_type, c_type, num_slices=ctx.num_slices

390 + )

391

392 # Make metadata tensors.

393 # Keep the metadata tensors in the CPU for further processing if needed.

@@ -356,27 +397,38 @@ class BenchmarkTensors:

397

398 # Make metadata tensors involved in correctness testing.

399 # Prepare seq lens tensor

359 - seq_len_tensor = torch.randint(ctx.seq_length, ctx.seq_length + 1,

360 - (ctx.batch_size, ))

400 + seq_len_tensor = torch.randint(

401 + ctx.seq_length, ctx.seq_length + 1, (ctx.batch_size,)

402 + )

403 assert total_tokens == seq_len_tensor.sum()

404 # Prepare prompt lora indices tensor

405 prompt_lora_indices_tensor = make_prompt_lora_mapping(

364 - ctx.batch_size, ctx.num_active_loras, ctx.sort_by_lora_id, "cpu")

406 + ctx.batch_size, ctx.num_active_loras, ctx.sort_by_lora_id, "cpu"

407 + )

408

409 # Make LoRAKernelMeta

410 token_lora_indices_tensor = make_token_lora_mapping(

368 - total_tokens, ctx.batch_size, prompt_lora_indices_tensor,

369 - seq_len_tensor, "cpu")

411 + total_tokens,

412 + ctx.batch_size,

413 + prompt_lora_indices_tensor,

414 + seq_len_tensor,

415 + "cpu",

416 + )

417 lora_kernel_meta = LoRAKernelMeta.make(

418 max_loras=ctx.num_loras,

419 max_num_tokens=token_lora_indices_tensor.size(0),

373 - device="cpu")

374 - lora_kernel_meta.prepare_tensors(

375 - token_lora_mapping=token_lora_indices_tensor)

376 -

377 - return BenchmarkTensors(input_tensor, lora_weights, output_tensor,

378 - lora_kernel_meta, seq_len_tensor,

379 - prompt_lora_indices_tensor)

420 + device="cpu",

421 + )

422 + lora_kernel_meta.prepare_tensors(token_lora_mapping=token_lora_indices_tensor)

423 +

424 + return BenchmarkTensors(

425 + input_tensor,

426 + lora_weights,

427 + output_tensor,

428 + lora_kernel_meta,

429 + seq_len_tensor,

430 + prompt_lora_indices_tensor,

431 + )

432

433 def sanity_check(self) -> None:

434 """

@@ -386,7 +438,7 @@ class BenchmarkTensors:

438 # check metadata tensors

439 assert torch.sum(self.seq_lens) == num_tokens

440 num_seqs = self.seq_lens.shape[0]

389 - #assert self.seq_start_loc.shape[0] == num_seqs

441 + # assert self.seq_start_loc.shape[0] == num_seqs

442 assert self.prompt_lora_mapping.shape[0] == num_seqs

443 assert self.lora_kernel_meta.token_lora_mapping.shape[0] == num_tokens

444

@@ -430,8 +482,11 @@ class BenchmarkTensors:

482 _, num_tokens, _, num_slices = self.metadata()

483

484 # Sanity check matrix shapes.

433 - i_shape, lw_shape, o_shape = self.input.shape, self.lora_weights_lst[

434 - 0].shape, self.output.shape

485 + i_shape, lw_shape, o_shape = (

486 + self.input.shape,

487 + self.lora_weights_lst[0].shape,

488 + self.output.shape,

489 + )

490 # Expected input shape [num_tokens, hidden_size]

491 assert len(i_shape) == 2

492 assert i_shape[0] == num_tokens

@@ -445,16 +500,17 @@ class BenchmarkTensors:

500 assert o_shape == (num_slices, num_tokens, lora_rank)

501

502 return {

448 - 'inputs': self.input,

449 - 'lora_a_weights': self.lora_weights_lst,

450 - 'output_tensor': self.output,

451 - 'token_lora_mapping': self.lora_kernel_meta.token_lora_mapping,

452 - 'token_indices_sorted_by_lora_ids':

453 - self.lora_kernel_meta.token_indices_sorted_by_lora_ids,

454 - 'num_tokens_per_lora': self.lora_kernel_meta.num_tokens_per_lora,

455 - 'lora_token_start_loc': self.lora_kernel_meta.lora_token_start_loc,

456 - 'lora_ids': self.lora_kernel_meta.active_lora_ids,

457 - 'scaling': 1.0,

503 + "inputs": self.input,

504 + "lora_a_weights": self.lora_weights_lst,

505 + "output_tensor": self.output,

506 + "token_lora_mapping": self.lora_kernel_meta.token_lora_mapping,

507 + "token_indices_sorted_by_lora_ids": (

508 + self.lora_kernel_meta.token_indices_sorted_by_lora_ids

509 + ),

510 + "num_tokens_per_lora": self.lora_kernel_meta.num_tokens_per_lora,

511 + "lora_token_start_loc": self.lora_kernel_meta.lora_token_start_loc,

512 + "lora_ids": self.lora_kernel_meta.active_lora_ids,

513 + "scaling": 1.0,

514 }

515

516 def as_lora_expand_kwargs(self, add_inputs: bool) -> dict[str, Any]:

@@ -464,8 +520,11 @@ class BenchmarkTensors:

520 _, num_tokens, _, num_slices = self.metadata()

521

522 # Sanity check matrix shapes.

467 - i_shape, lw_shape, o_shape = self.input.shape, self.lora_weights_lst[

468 - 0].shape, self.output.shape

523 + i_shape, lw_shape, o_shape = (

524 + self.input.shape,

525 + self.lora_weights_lst[0].shape,

526 + self.output.shape,

527 + )

528 # Expected input shape : [num_slices, num_tokens, lora_rank]

529 assert len(i_shape) == 3

530 assert i_shape[0] == num_slices

@@ -480,22 +539,23 @@ class BenchmarkTensors:

539 assert o_shape == (num_tokens, hidden_size * num_slices)

540

541 return {

483 - 'inputs': self.input,

484 - 'lora_b_weights': self.lora_weights_lst,

485 - 'output_tensor': self.output,

486 - 'token_lora_mapping': self.lora_kernel_meta.token_lora_mapping,

487 - 'token_indices_sorted_by_lora_ids':

488 - self.lora_kernel_meta.token_indices_sorted_by_lora_ids,

489 - 'num_tokens_per_lora': self.lora_kernel_meta.num_tokens_per_lora,

490 - 'lora_token_start_loc': self.lora_kernel_meta.lora_token_start_loc,

491 - 'lora_ids': self.lora_kernel_meta.active_lora_ids,

492 - 'offset_start': 0,

493 - 'add_inputs': add_inputs,

542 + "inputs": self.input,

543 + "lora_b_weights": self.lora_weights_lst,

544 + "output_tensor": self.output,

545 + "token_lora_mapping": self.lora_kernel_meta.token_lora_mapping,

546 + "token_indices_sorted_by_lora_ids": (

547 + self.lora_kernel_meta.token_indices_sorted_by_lora_ids

548 + ),

549 + "num_tokens_per_lora": self.lora_kernel_meta.num_tokens_per_lora,

550 + "lora_token_start_loc": self.lora_kernel_meta.lora_token_start_loc,

551 + "lora_ids": self.lora_kernel_meta.active_lora_ids,

552 + "offset_start": 0,

553 + "add_inputs": add_inputs,

554 }

555

496 - def bench_fn_kwargs(self,

497 - op_type: OpType,

498 - add_inputs: Optional[bool] = None) -> dict[str, Any]:

556 + def bench_fn_kwargs(

557 + self, op_type: OpType, add_inputs: Optional[bool] = None

558 + ) -> dict[str, Any]:

559 if op_type.is_shrink_fn():

560 assert add_inputs is None

561 else:

@@ -507,8 +567,9 @@ class BenchmarkTensors:

567 return self.as_lora_expand_kwargs(add_inputs)

568 raise ValueError(f"Unrecognized optype {self}")

569

510 - def test_correctness(self, op_type: OpType,

511 - expand_fn_add_inputs: Optional[bool]) -> bool:

570 + def test_correctness(

571 + self, op_type: OpType, expand_fn_add_inputs: Optional[bool]

572 + ) -> bool:

573 """

574 Test correctness of op_type implementation against a grouped gemm

575 reference implementation.

@@ -518,8 +579,7 @@ class BenchmarkTensors:

579 ref_output = self.output.clone()

580

581 self.output.zero_()

521 - op_type.bench_fn()(

522 - **self.bench_fn_kwargs(op_type, expand_fn_add_inputs))

582 + op_type.bench_fn()(**self.bench_fn_kwargs(op_type, expand_fn_add_inputs))

583

584 op_type.run_ref_group_gemm(

585 ref_output,

@@ -528,7 +588,8 @@ class BenchmarkTensors:

588 seq_lens_cpu=seq_lens_cpu,

589 prompt_lora_mapping_cpu=prompt_lora_mapping_cpu,

590 scaling=1.0,

531 - add_inputs=expand_fn_add_inputs)

591 + add_inputs=expand_fn_add_inputs,

592 + )

593

594 rtol, atol = {

595 torch.float16: (6e-2, 6e-2),

@@ -539,13 +600,14 @@ class BenchmarkTensors:

600 return torch.allclose(ref_output, self.output, rtol=rtol, atol=atol)

601

602

542 -def bench_optype(ctx: BenchmarkContext,

543 - arg_pool_size: int,

544 - op_type: OpType,

545 - cuda_graph_nops: Optional[int] = None,

546 - expand_fn_add_inputs: Optional[bool] = None,

547 - test_correctness: bool = False) -> TMeasurement:

548 -

603 +def bench_optype(

604 + ctx: BenchmarkContext,

605 + arg_pool_size: int,

606 + op_type: OpType,

607 + cuda_graph_nops: Optional[int] = None,

608 + expand_fn_add_inputs: Optional[bool] = None,

609 + test_correctness: bool = False,

610 +) -> TMeasurement:

611 assert arg_pool_size >= 1

612 if op_type.is_shrink_fn():

613 assert expand_fn_add_inputs is None

@@ -553,17 +615,17 @@ def bench_optype(ctx: BenchmarkContext,

615 assert expand_fn_add_inputs is not None

616

617 # BenchmarkContext -> BenchmarkTensors

556 - bench_tensors : list[BenchmarkTensors] = \

557 - [BenchmarkTensors.make(ctx, op_type) for _ in range(arg_pool_size)]

618 + bench_tensors: list[BenchmarkTensors] = [

619 + BenchmarkTensors.make(ctx, op_type) for _ in range(arg_pool_size)

620 + ]

621 for bt in bench_tensors:

622 bt.sanity_check()

623

624 # Test correctness of our implementation.

625 if test_correctness:

563 - assert all([

564 - bt.test_correctness(op_type, expand_fn_add_inputs)

565 - for bt in bench_tensors

566 - ])

626 + assert all(

627 + [bt.test_correctness(op_type, expand_fn_add_inputs) for bt in bench_tensors]

628 + )

629

630 # BenchmarkTensors -> dict (kwargs)

631 kwargs_list = [

@@ -585,40 +647,49 @@ def bench_optype(ctx: BenchmarkContext,

647 for k, v in _kwargs.items():

648 kwargs[k].values.append(v)

649

588 - describe_args = (f"add_inputs={expand_fn_add_inputs}"

589 - if expand_fn_add_inputs is not None else "")

590 - description = (

591 - f"{op_type.name}({describe_args}) ({bench_tensors[0].io_types()})")

650 + describe_args = (

651 + f"add_inputs={expand_fn_add_inputs}" if expand_fn_add_inputs is not None else ""

652 + )

653 + description = f"{op_type.name}({describe_args}) ({bench_tensors[0].io_types()})"

654

655 cuda_graph_params = None

656 if cuda_graph_nops:

657 cuda_graph_params = CudaGraphBenchParams(cuda_graph_nops)

658 timer = None

597 - with Bench(cuda_graph_params,

598 - ctx.bench_label(), ctx.bench_sublabel(op_type), description,

599 - op_type.bench_fn(), **kwargs) as bench:

659 + with Bench(

660 + cuda_graph_params,

661 + ctx.bench_label(),

662 + ctx.bench_sublabel(op_type),

663 + description,

664 + op_type.bench_fn(),

665 + **kwargs,

666 + ) as bench:

667 timer = bench.run()

668 return timer

669

670

604 -def bench_torch_mm(ctx: BenchmarkContext,

605 - arg_pool_size: int,

606 - op_type: OpType,

607 - cuda_graph_nops: Optional[int] = None) -> TMeasurement:

671 +def bench_torch_mm(

672 + ctx: BenchmarkContext,

673 + arg_pool_size: int,

674 + op_type: OpType,

675 + cuda_graph_nops: Optional[int] = None,

676 +) -> TMeasurement:

677 """

678 Benchmark basic torch.mm as a roofline.

679

680 When all the input tokens have the same LoRA ID, the LoRA kernels are just

612 - a matmul. This torch.mm benchmark serves as a roofline for that case.

681 + a matmul. This torch.mm benchmark serves as a roofline for that case.

682

683 input op_type is used in determining the m, k, n dimensions for the matmul.

684 """

685

617 - batch_size, hidden_size, lora_rank, seq_length, dtype = (ctx.batch_size,

618 - ctx.hidden_size,

619 - ctx.lora_rank,

620 - ctx.seq_length,

621 - ctx.dtype)

686 + batch_size, hidden_size, lora_rank, seq_length, dtype = (

687 + ctx.batch_size,

688 + ctx.hidden_size,

689 + ctx.lora_rank,

690 + ctx.seq_length,

691 + ctx.dtype,

692 + )

693

694 m, k, n = op_type.mkn(batch_size, seq_length, hidden_size, lora_rank)

695 # For a fairer comparison.

@@ -632,18 +703,24 @@ def bench_torch_mm(ctx: BenchmarkContext,

703 Cs.append(torch.rand((m, n), dtype=dtype).to("cuda"))

704

705 # Make torch.mm kwargs

635 - mm_kwargs = {'input': ArgPool(As), 'mat2': ArgPool(Bs), 'out': ArgPool(Cs)}

706 + mm_kwargs = {"input": ArgPool(As), "mat2": ArgPool(Bs), "out": ArgPool(Cs)}

707

708 description = (

709 f"single-lora roofline using torch.mm ({dtype_to_str(dtype)}"

710 f"x{dtype_to_str(dtype)}"

640 - f"=>{dtype_to_str(dtype)})")

711 + f"=>{dtype_to_str(dtype)})"

712 + )

713 cuda_graph_params = None

714 if cuda_graph_nops:

715 cuda_graph_params = CudaGraphBenchParams(cuda_graph_nops)

644 - with Bench(cuda_graph_params, ctx.bench_label(),

645 - ctx.bench_sublabel(op_type), description, torch.mm,

646 - **mm_kwargs) as bench:

716 + with Bench(

717 + cuda_graph_params,

718 + ctx.bench_label(),

719 + ctx.bench_sublabel(op_type),

720 + description,

721 + torch.mm,

722 + **mm_kwargs,

723 + ) as bench:

724 return bench.run()

725

726

@@ -660,8 +737,7 @@ def use_cuda_graph_recommendation() -> str:

737 """

738

739

663 -def print_timers(timers: list[TMeasurement],

664 - args: Optional[argparse.Namespace] = None):

740 +def print_timers(timers: list[TMeasurement], args: Optional[argparse.Namespace] = None):

741 compare = TBenchmark.Compare(timers)

742 compare.print()

743

@@ -670,22 +746,23 @@ def print_timers(timers: list[TMeasurement],

746 f"Note : The timings reported above is for {args.cuda_graph_nops} "

747 "consecutive invocations of the benchmarking functions. "

748 f"Please divide by {args.cuda_graph_nops} for single invocation "

673 - "timings.")

749 + "timings."

750 + )

751

675 - print("Note on Comparison with torch.mm : The torch.mm numbers are "

676 - "benchmark numbers of a simple matmul emulating the single lora "

677 - "case. It is provided as a roofline for comparing our LoRA Kernel "

678 - "implementations. It is expected that the LoRA kernels will be "

679 - "slower than torch.mm in cases where num_loras is big. But for "

680 - "small num_loras the goal should be to match the torch.mm numbers.")

752 + print(

753 + "Note on Comparison with torch.mm : The torch.mm numbers are "

754 + "benchmark numbers of a simple matmul emulating the single lora "

755 + "case. It is provided as a roofline for comparing our LoRA Kernel "

756 + "implementations. It is expected that the LoRA kernels will be "

757 + "slower than torch.mm in cases where num_loras is big. But for "

758 + "small num_loras the goal should be to match the torch.mm numbers."

759 + )

760

761

762 def run(args: argparse.Namespace, bench_ctxs: list[BenchmarkContext]):

684 -

763 if args.cuda_graph_nops is not None:

764 assert args.cuda_graph_nops > 0

687 - print(f"Benchmarking {args.cuda_graph_nops} invocations inside a CUDA "

688 - "Graph")

765 + print(f"Benchmarking {args.cuda_graph_nops} invocations inside a CUDA Graph")

766 else:

767 print(f"CUDA Graphs not enabled.\n{use_cuda_graph_recommendation()}")

768

@@ -697,21 +774,30 @@ def run(args: argparse.Namespace, bench_ctxs: list[BenchmarkContext]):

774 for bench_op in bench_ops:

775 for num_slices in bench_op.num_slices():

776 _ctx = bench_ctx.with_seq_length(seq_len).with_num_slices(

700 - num_slices)

777 + num_slices

778 + )

779 # Benchmark torch.mm as a roofline

780 seq_len_timers.append(

703 - bench_torch_mm(_ctx, args.arg_pool_size, bench_op,

704 - args.cuda_graph_nops))

781 + bench_torch_mm(

782 + _ctx, args.arg_pool_size, bench_op, args.cuda_graph_nops

783 + )

784 + )

785

786 # Benchmark bench_op

707 - expand_fn_add_inputs = [

708 - None

709 - ] if bench_op.is_shrink_fn() else args.expand_fn_add_inputs

787 + expand_fn_add_inputs = (

788 + [None] if bench_op.is_shrink_fn() else args.expand_fn_add_inputs

789 + )

790 for add_input_arg in expand_fn_add_inputs:

791 seq_len_timers.append(

712 - bench_optype(_ctx, args.arg_pool_size, bench_op,

713 - args.cuda_graph_nops, add_input_arg,

714 - args.test_correctness))

792 + bench_optype(

793 + _ctx,

794 + args.arg_pool_size,

795 + bench_op,

796 + args.cuda_graph_nops,

797 + add_input_arg,

798 + args.test_correctness,

799 + )

800 + )

801

802 print_timers(seq_len_timers)

803 timers.extend(seq_len_timers)

@@ -733,13 +819,17 @@ def run(args: argparse.Namespace, bench_ctxs: list[BenchmarkContext]):

819 pickle.dump(timers, f)

820

821

736 -def as_benchmark_contexts(hidden_sizes: list[int], lora_ranks: list[int],

737 - args: argparse.Namespace) -> list[BenchmarkContext]:

738 -

822 +def as_benchmark_contexts(

823 + hidden_sizes: list[int], lora_ranks: list[int], args: argparse.Namespace

824 +) -> list[BenchmarkContext]:

825 ctxs: list[BenchmarkContext] = []

826 for batch_size, hidden_size, lora_rank, num_loras, sort_by_lora_id in product( # noqa

741 - args.batch_sizes, list(hidden_sizes), lora_ranks, args.num_loras,

742 - args.sort_by_lora_id):

827 + args.batch_sizes,

828 + list(hidden_sizes),

829 + lora_ranks,

830 + args.num_loras,

831 + args.sort_by_lora_id,

832 + ):

833 ctxs.append(

834 BenchmarkContext(

835 batch_size=batch_size,

@@ -747,13 +837,16 @@ def as_benchmark_contexts(hidden_sizes: list[int], lora_ranks: list[int],

837 lora_rank=lora_rank,

838 num_loras=num_loras,

839 num_active_loras=args.num_active_loras

750 - if args.num_active_loras else num_loras,

840 + if args.num_active_loras

841 + else num_loras,

842 # To be filled based on the OpType to benchmark

843 seq_length=None,

844 sort_by_lora_id=sort_by_lora_id,

845 dtype=args.dtype,

846 # To be filled based on the OpType to benchmark

756 - num_slices=None))

847 + num_slices=None,

848 + )

849 + )

850

851 return ctxs

852

@@ -761,13 +854,16 @@ def as_benchmark_contexts(hidden_sizes: list[int], lora_ranks: list[int],

854 def run_list_bench(args: argparse.Namespace):

855 print(args)

856

764 - print("List bench :\n"

765 - f" Hidden Sizes {args.hidden_sizes}"

766 - f" LoRA Ranks {args.lora_ranks}")

857 + print(

858 + "List bench :\n"

859 + f" Hidden Sizes {args.hidden_sizes}"

860 + f" LoRA Ranks {args.lora_ranks}"

861 + )

862

863 # Get all benchmarking contexts

864 bench_contexts: list[BenchmarkContext] = as_benchmark_contexts(

770 - hidden_sizes=args.hidden_sizes, lora_ranks=args.lora_ranks, args=args)

865 + hidden_sizes=args.hidden_sizes, lora_ranks=args.lora_ranks, args=args

866 + )

867

868 run(args, bench_contexts)

869

@@ -776,19 +872,22 @@ def run_range_bench(args: argparse.Namespace):

872 print(args)

873

874 hidden_sizes = list(

779 - range(args.hidden_sizes_start, args.hidden_sizes_end + 1,

780 - args.hidden_sizes_increment))

875 + range(

876 + args.hidden_sizes_start,

877 + args.hidden_sizes_end + 1,

878 + args.hidden_sizes_increment,

879 + )

880 + )

881 lora_ranks = list(

782 - range(args.lora_ranks_start, args.lora_ranks_end + 1,

783 - args.lora_ranks_increment))

882 + range(args.lora_ranks_start, args.lora_ranks_end + 1, args.lora_ranks_increment)

883 + )

884

785 - print("Range bench :\n"

786 - f" Hidden Sizes {hidden_sizes}"

787 - f" LoRA Ranks {lora_ranks}")

885 + print(f"Range bench :\n Hidden Sizes {hidden_sizes} LoRA Ranks {lora_ranks}")

886

887 # Get all benchmarking contexts

888 bench_contexts: list[BenchmarkContext] = as_benchmark_contexts(

791 - hidden_sizes=hidden_sizes, lora_ranks=lora_ranks, args=args)

889 + hidden_sizes=hidden_sizes, lora_ranks=lora_ranks, args=args

890 + )

891

892 run(args, bench_contexts)

893

@@ -806,21 +905,19 @@ def run_model_bench(args: argparse.Namespace):

905 # Get all hidden sizes

906 hidden_sizes: set[int] = set()

907 for model_name, tp_size in product(args.models, args.tp_sizes):

809 - hidden_sizes = hidden_sizes.union(

810 - hidden_sizes_from_model(model_name, tp_size))

908 + hidden_sizes = hidden_sizes.union(hidden_sizes_from_model(model_name, tp_size))

909

812 - print("Model bench :\n"

813 - f" Hidden Sizes {hidden_sizes}"

814 - f" LoRA Ranks {args.lora_ranks}")

910 + print(f"Model bench :\n Hidden Sizes {hidden_sizes} LoRA Ranks {args.lora_ranks}")

911

912 # Get all benchmarking contexts

913 bench_contexts: list[BenchmarkContext] = as_benchmark_contexts(

818 - hidden_sizes=hidden_sizes, lora_ranks=args.lora_ranks, args=args)

914 + hidden_sizes=hidden_sizes, lora_ranks=args.lora_ranks, args=args

915 + )

916

917 run(args, bench_contexts)

918

919

823 -if __name__ == '__main__':

920 +if __name__ == "__main__":

921

922 def to_torch_dtype(dt):

923 if dt == "torch.float16":

@@ -830,14 +927,15 @@ if __name__ == '__main__':

927 raise ValueError("unsupported dtype")

928

929 def get_bool(s: str) -> bool:

833 - return s.lower() in ['true', '1']

930 + return s.lower() in ["true", "1"]

931

932 def add_common_command_args(p: argparse.ArgumentParser):

933 p.add_argument(

934 "--dtype",

935 type=to_torch_dtype,

936 required=True,

840 - help="Available options are ['torch.float16', 'torch.bfloat16']")

937 + help="Available options are ['torch.float16', 'torch.bfloat16']",

938 + )

939

940 p.add_argument(

941 "--arg-pool-size",

@@ -845,56 +943,66 @@ if __name__ == '__main__':

943 default=32,

944 help="Run profiles with a pool of input/output/meta tensors instead"

945 "of simply reusing the same tensors for all runs. A bigger arg-pool"

848 - "mitigates hardware caching effects during benchmarking.")

946 + "mitigates hardware caching effects during benchmarking.",

947 + )

948

949 p.add_argument(

950 "--cuda-graph-nops",

951 type=int,

853 - help=("when set profiling is done using cudagraph, "

854 - "with the given number of operations in a graph."

855 - "Note that the measurement returned is the time "

856 - "taken for N consecutive executions of the benchmarking "

857 - "functions, where N is the value of this argument."))

858 - p.add_argument("--num-loras",

859 - nargs="+",

860 - type=int,

861 - default=DEFAULT_NUM_LORAS)

862 - p.add_argument("--num-active-loras",

863 - type=int,

864 - default=None,

865 - help="Active LoRAs. When None, all LoRAs are active")

866 - p.add_argument("--sort-by-lora-id",

867 - nargs="+",

868 - type=get_bool,

869 - default=DEFAULT_SORT_BY_LORA_IDS)

870 - p.add_argument("--op-types",

871 - nargs="+",

872 - type=OpType.from_str,

873 - default=list(OpType))

874 - p.add_argument('--seq-lengths',

875 - nargs="+",

876 - type=int,

877 - default=DEFAULT_SEQ_LENGTHS)

878 - p.add_argument("--batch-sizes",

879 - nargs="+",

880 - type=int,

881 - default=DEFAULT_BATCH_SIZES)

882 - p.add_argument("--expand-fn-add-inputs",

883 - nargs="+",

884 - type=get_bool,

885 - default=DEFAULT_EXPAND_FN_ADD_INPUTS)

952 + help=(

953 + "when set profiling is done using cudagraph, "

954 + "with the given number of operations in a graph."

955 + "Note that the measurement returned is the time "

956 + "taken for N consecutive executions of the benchmarking "

957 + "functions, where N is the value of this argument."

958 + ),

959 + )

960 + p.add_argument("--num-loras", nargs="+", type=int, default=DEFAULT_NUM_LORAS)

961 + p.add_argument(

962 + "--num-active-loras",

963 + type=int,

964 + default=None,

965 + help="Active LoRAs. When None, all LoRAs are active",

966 + )

967 + p.add_argument(

968 + "--sort-by-lora-id",

969 + nargs="+",

970 + type=get_bool,

971 + default=DEFAULT_SORT_BY_LORA_IDS,

972 + )

973 + p.add_argument(

974 + "--op-types", nargs="+", type=OpType.from_str, default=list(OpType)

975 + )

976 + p.add_argument(

977 + "--seq-lengths", nargs="+", type=int, default=DEFAULT_SEQ_LENGTHS

978 + )

979 + p.add_argument(

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

981 + )

982 + p.add_argument(

983 + "--expand-fn-add-inputs",

984 + nargs="+",

985 + type=get_bool,

986 + default=DEFAULT_EXPAND_FN_ADD_INPUTS,

987 + )

988 p.add_argument(

887 - '-o',

888 - '--output-directory',

989 + "-o",

990 + "--output-directory",

991 type=str,

890 - help=("Output directory to store a the list of benchmarking"

891 - "TMeasurement objects as a pickle file"))

992 + help=(

993 + "Output directory to store a the list of benchmarking"

994 + "TMeasurement objects as a pickle file"

995 + ),

996 + )

997

998 p.add_argument(

999 "--test-correctness",

895 - action='store_true',

896 - help=("When enabled, the benchmarking functions are tested"

897 - "for correctness before the actual benchmarking"))

1000 + action="store_true",

1001 + help=(

1002 + "When enabled, the benchmarking functions are tested"

1003 + "for correctness before the actual benchmarking"

1004 + ),

1005 + )

1006

1007 parser = FlexibleArgumentParser(

1008 description=f"""

@@ -910,50 +1018,45 @@ Benchmark LoRA kernels:

1018 range_bench example:

1019 python3 benchmarks/kernels/benchmark_lora.py range_bench --arg-pool-size 32 --batch-sizes 1 16 32 --dtype torch.float16 --num-loras 1 4 --op-types lora_shrink lora_expand --seq-lengths 1 16 --sort-by-lora-id 1 --cuda-graph-nops 32 --hidden-sizes-start 1024 --hidden-sizes-end 4096 --hidden-sizes-increment 1024 --lora-ranks-start 8 --lora-ranks-end 24 --lora-ranks-increment 8

1020 """, # noqa: E501

913 - formatter_class=argparse.RawTextHelpFormatter)

1021 + formatter_class=argparse.RawTextHelpFormatter,

1022 + )

1023

1024 subparsers = parser.add_subparsers(dest="cmd", required=True)

1025

1026 list_parser = subparsers.add_parser("list_bench")

918 - list_parser.add_argument("--hidden-sizes",

919 - nargs="+",

920 - type=int,

921 - default=DEFAULT_HIDDEN_SIZES)

922 - list_parser.add_argument("--lora-ranks",

923 - nargs="+",

924 - type=int,

925 - default=DEFAULT_LORA_RANKS)

1027 + list_parser.add_argument(

1028 + "--hidden-sizes", nargs="+", type=int, default=DEFAULT_HIDDEN_SIZES

1029 + )

1030 + list_parser.add_argument(

1031 + "--lora-ranks", nargs="+", type=int, default=DEFAULT_LORA_RANKS

1032 + )

1033 add_common_command_args(list_parser)

1034 list_parser.set_defaults(func=run_list_bench)

1035

1036 range_parser = subparsers.add_parser("range_bench")

1037 range_parser.add_argument("--hidden-sizes-start", type=int, required=True)

1038 range_parser.add_argument("--hidden-sizes-end", type=int, required=True)

932 - range_parser.add_argument("--hidden-sizes-increment",

933 - type=int,

934 - required=True)

1039 + range_parser.add_argument("--hidden-sizes-increment", type=int, required=True)

1040 range_parser.add_argument("--lora-ranks-start", type=int, required=True)

1041 range_parser.add_argument("--lora-ranks-end", type=int, required=True)

937 - range_parser.add_argument("--lora-ranks-increment",

938 - type=int,

939 - required=True)

1042 + range_parser.add_argument("--lora-ranks-increment", type=int, required=True)

1043 add_common_command_args(range_parser)

1044 range_parser.set_defaults(func=run_range_bench)

1045

1046 model_parser = subparsers.add_parser("model_bench")

944 - model_parser.add_argument("--models",

945 - nargs="+",

946 - type=str,

947 - default=DEFAULT_MODELS,

948 - choices=WEIGHT_SHAPES.keys())

949 - model_parser.add_argument("--tp-sizes",

950 - nargs="+",

951 - type=int,

952 - default=DEFAULT_TP_SIZES)

953 - model_parser.add_argument("--lora-ranks",

954 - nargs="+",

955 - type=int,

956 - default=DEFAULT_LORA_RANKS)

1047 + model_parser.add_argument(

1048 + "--models",

1049 + nargs="+",

1050 + type=str,

1051 + default=DEFAULT_MODELS,

1052 + choices=WEIGHT_SHAPES.keys(),

1053 + )

1054 + model_parser.add_argument(

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

1056 + )

1057 + model_parser.add_argument(

1058 + "--lora-ranks", nargs="+", type=int, default=DEFAULT_LORA_RANKS

1059 + )

1060 add_common_command_args(model_parser)

1061 model_parser.set_defaults(func=run_model_bench)

1062

benchmarks/kernels/benchmark_machete.py

+212 -157

@@ -20,12 +20,18 @@ from weight_shapes import WEIGHT_SHAPES

21 from vllm import _custom_ops as ops

22 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

23 - GPTQ_MARLIN_MAX_PARALLEL, GPTQ_MARLIN_MIN_THREAD_N, marlin_permute_scales,

24 - marlin_zero_points)

23 + GPTQ_MARLIN_MAX_PARALLEL,

24 + GPTQ_MARLIN_MIN_THREAD_N,

25 + marlin_permute_scales,

26 + marlin_zero_points,

27 +)

28 from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (

26 - MarlinWorkspace)

29 + MarlinWorkspace,

30 +)

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

28 - pack_rows, quantize_weights)

32 + pack_rows,

33 + quantize_weights,

34 +)

35 from vllm.scalar_type import ScalarType, scalar_types

36 from vllm.utils import FlexibleArgumentParser

@@ -82,12 +88,14 @@ def rand_data(shape, dtype=torch.float16, scale=1):

88 return torch.randint(-15, 15, shape, dtype=dtype, device="cuda")

85 -def quantize_and_pack(atype: torch.dtype,

86 - w: torch.Tensor,

87 - wtype: ScalarType,

88 - stype: Optional[torch.dtype],

89 - group_size: Optional[int],

90 - zero_points: bool = False):

91 +def quantize_and_pack(

92 + atype: torch.dtype,

93 + w: torch.Tensor,

94 + wtype: ScalarType,

95 + stype: Optional[torch.dtype],

96 + group_size: Optional[int],

97 + zero_points: bool = False,

98 +):

99 assert wtype.is_integer(), "TODO: support floating point weights"

100

101 w_ref, w_q, w_s, w_zp = quantize_weights(

@@ -96,21 +104,24 @@ def quantize_and_pack(atype: torch.dtype,

104 group_size=group_size,

105 zero_points=zero_points,

106 # to match how the kernel applies zps

99 - ref_zero_points_after_scales=True)

107 + ref_zero_points_after_scales=True,

108 + )

109

110 w_q = pack_rows(w_q, wtype.size_bits, *w_q.shape)

111 return w_ref, w_q, w_s, w_zp

112

113

105 -def create_bench_tensors(shape: tuple[int, int, int], types: TypeConfig,

106 - group_size: Optional[int]) -> list[BenchmarkTensors]:

114 +def create_bench_tensors(

115 + shape: tuple[int, int, int], types: TypeConfig, group_size: Optional[int]

116 +) -> list[BenchmarkTensors]:

117 m, n, k = shape

118

119 # we want to make sure that weights don't fit into L2 cache between runs so

120 # we construct enough weights to exceed L2 cache, which is 50mb on a H100

121 # so we target total weight size > 2*50mb

112 - num_weights = math.ceil(2 * 50 * 1024**2 * 8 /

113 - (k * n * types.weight_type.size_bits))

122 + num_weights = math.ceil(

123 + 2 * 50 * 1024**2 * 8 / (k * n * types.weight_type.size_bits)

124 + )

125

126 a = rand_data((m, k), types.act_type, scale=5)

127

@@ -124,8 +135,13 @@ def create_bench_tensors(shape: tuple[int, int, int], types: TypeConfig,

135 w = w.to(torch.float16)

136

137 w_ref, w_q_packed, w_s, w_zp = quantize_and_pack(

127 - a.dtype, w, types.weight_type, types.group_scale_type, group_size,

128 - types.group_zero_type is not None)

138 + a.dtype,

139 + w,

140 + types.weight_type,

141 + types.group_scale_type,

142 + group_size,

143 + types.group_zero_type is not None,

144 + )

145

146 if not a.dtype.is_floating_point:

147 aiinfo = torch.iinfo(a.dtype)

@@ -133,21 +149,30 @@ def create_bench_tensors(shape: tuple[int, int, int], types: TypeConfig,

149

150 w_ref = w_ref.to(torch.float32)

151

136 - w_ch_s = None if types.channel_scale_type is None else\

137 - rand_data((n,), types.channel_scale_type)

138 - w_tok_s = None if types.token_scale_type is None else\

139 - rand_data((m,), types.token_scale_type)

152 + w_ch_s = (

153 + None

154 + if types.channel_scale_type is None

155 + else rand_data((n,), types.channel_scale_type)

156 + )

157 + w_tok_s = (

158 + None

159 + if types.token_scale_type is None

160 + else rand_data((m,), types.token_scale_type)

161 + )

162

163 benchmark_tensors.append(

142 - BenchmarkTensors(w_ref=w_ref,

143 - a=a,

144 - w_q=w_q_packed,

145 - wtype=types.weight_type,

146 - w_g_s=w_s,

147 - w_g_zp=w_zp,

148 - group_size=group_size,

149 - w_ch_s=w_ch_s,

150 - w_tok_s=w_tok_s))

164 + BenchmarkTensors(

165 + w_ref=w_ref,

166 + a=a,

167 + w_q=w_q_packed,

168 + wtype=types.weight_type,

169 + w_g_s=w_s,

170 + w_g_zp=w_zp,

171 + group_size=group_size,

172 + w_ch_s=w_ch_s,

173 + w_tok_s=w_tok_s,

174 + )

175 + )

176

177 return benchmark_tensors

178

@@ -170,50 +195,57 @@ def cutlass_scaled_mm_create_bench_fn(bt: BenchmarkTensors) -> Callable:

195 scale_b = torch.tensor(1.0, dtype=torch.float32, device=bt.a.device)

196 w_col_major = bt.w_ref.to(bt.a.dtype).t().contiguous().t()

197 return lambda: ops.cutlass_scaled_mm(

173 - bt.a, w_col_major, scale_a, scale_b, out_dtype=torch.float16)

198 + bt.a, w_col_major, scale_a, scale_b, out_dtype=torch.float16

199 + )

200

201

202 def marlin_create_bench_fn(bt: BenchmarkTensors) -> Callable:

203 device = bt.a.device

204

179 - workspace = MarlinWorkspace(bt.w_ref.shape[1], GPTQ_MARLIN_MIN_THREAD_N,

180 - GPTQ_MARLIN_MAX_PARALLEL)

205 + workspace = MarlinWorkspace(

206 + bt.w_ref.shape[1], GPTQ_MARLIN_MIN_THREAD_N, GPTQ_MARLIN_MAX_PARALLEL

207 + )

208

209 if bt.w_g_zp is None:

210 w_zp = torch.empty(0, dtype=torch.int, device=device)

211 else:

185 - w_zp = marlin_zero_points(bt.w_g_zp, bt.w_ref.shape[0],

186 - bt.w_ref.shape[1], bt.wtype.size_bits)

212 + w_zp = marlin_zero_points(

213 + bt.w_g_zp, bt.w_ref.shape[0], bt.w_ref.shape[1], bt.wtype.size_bits

214 + )

215

216 if bt.group_size is None:

217 w_s = torch.tensor([], device="cuda", dtype=torch.half)

218 else:

191 - w_s = marlin_permute_scales(bt.w_g_s, bt.w_ref.shape[0],

192 - bt.w_ref.shape[1], bt.group_size)

219 + w_s = marlin_permute_scales(

220 + bt.w_g_s, bt.w_ref.shape[0], bt.w_ref.shape[1], bt.group_size

221 + )

222

223 sort_indices = torch.empty(0, dtype=torch.int, device=device)

224 g_idx = torch.empty(0, dtype=torch.int, device=device)

196 - w_q = ops.gptq_marlin_repack(bt.w_q, sort_indices, bt.w_ref.shape[0],

197 - bt.w_ref.shape[1], bt.wtype.size_bits)

225 + w_q = ops.gptq_marlin_repack(

226 + bt.w_q, sort_indices, bt.w_ref.shape[0], bt.w_ref.shape[1], bt.wtype.size_bits

227 + )

228

229 if bt.a.dtype.is_floating_point:

230 assert bt.w_ch_s is None

231 assert bt.w_tok_s is None

232 assert bt.group_size is not None

233

204 - fn = lambda: ops.gptq_marlin_gemm(a=bt.a,

205 - b_q_weight=w_q,

206 - b_scales=w_s,

207 - b_zeros=w_zp,

208 - g_idx=g_idx,

209 - perm=sort_indices,

210 - workspace=workspace.scratch,

211 - b_q_type=bt.wtype,

212 - size_m=bt.a.shape[0],

213 - size_n=bt.w_ref.shape[1],

214 - size_k=bt.w_ref.shape[0],

215 - is_k_full=True,

216 - is_zp_float=False)

234 + fn = lambda: ops.gptq_marlin_gemm(

235 + a=bt.a,

236 + b_q_weight=w_q,

237 + b_scales=w_s,

238 + b_zeros=w_zp,

239 + g_idx=g_idx,

240 + perm=sort_indices,

241 + workspace=workspace.scratch,

242 + b_q_type=bt.wtype,

243 + size_m=bt.a.shape[0],

244 + size_n=bt.w_ref.shape[1],

245 + size_k=bt.w_ref.shape[0],

246 + is_k_full=True,

247 + is_zp_float=False,

248 + )

249 else:

250 assert bt.a.dtype == torch.int8

251 assert bt.wtype == scalar_types.uint4b8

@@ -221,36 +253,35 @@ def marlin_create_bench_fn(bt: BenchmarkTensors) -> Callable:

253 if bt.w_ch_s is not None:

254 s_ch = bt.w_ch_s.to(torch.float32)

255 else:

224 - s_ch = torch.ones(bt.w_ref.shape[1],

225 - dtype=torch.float32,

226 - device=device)

256 + s_ch = torch.ones(bt.w_ref.shape[1], dtype=torch.float32, device=device)

257

258 if bt.w_tok_s is not None:

259 s_tok = bt.w_tok_s.to(torch.float32)

260 else:

231 - s_tok = torch.ones(bt.a.shape[0],

232 - dtype=torch.float32,

233 - device=device)

234 -

235 - fn = lambda: ops.marlin_qqq_gemm(a=bt.a,

236 - b_q_weight=w_q,

237 - s_group=w_s,

238 - s_tok=s_tok,

239 - s_ch=s_ch,

240 - workspace=workspace.scratch,

241 - size_m=bt.a.shape[0],

242 - size_n=bt.w_ref.shape[1],

243 - size_k=bt.w_ref.shape[0])

261 + s_tok = torch.ones(bt.a.shape[0], dtype=torch.float32, device=device)

262 +

263 + fn = lambda: ops.marlin_qqq_gemm(

264 + a=bt.a,

265 + b_q_weight=w_q,

266 + s_group=w_s,

267 + s_tok=s_tok,

268 + s_ch=s_ch,

269 + workspace=workspace.scratch,

270 + size_m=bt.a.shape[0],

271 + size_n=bt.w_ref.shape[1],

272 + size_k=bt.w_ref.shape[0],

273 + )

274

275 return fn

276

277

248 -def machete_create_bench_fn(bt: BenchmarkTensors,

249 - out_type=torch.dtype,

250 - schedule=None) -> Callable:

278 +def machete_create_bench_fn(

279 + bt: BenchmarkTensors, out_type=torch.dtype, schedule=None

280 +) -> Callable:

281 w_q = bt.w_q.t().contiguous().t() # make col major

252 - w_q = ops.machete_prepack_B(w_q, bt.a.dtype, bt.wtype,

253 - None if bt.w_g_s is None else bt.w_g_s.dtype)

282 + w_q = ops.machete_prepack_B(

283 + w_q, bt.a.dtype, bt.wtype, None if bt.w_g_s is None else bt.w_g_s.dtype

284 + )

285

286 w_g_zp = bt.w_g_zp

287 if w_g_zp is not None:

@@ -275,26 +306,24 @@ def machete_create_bench_fn(bt: BenchmarkTensors,

306 # bench

307

308

278 -def bench_fns(label: str, sub_label: str, description: str,

279 - fns: list[Callable]):

280 -

309 +def bench_fns(label: str, sub_label: str, description: str, fns: list[Callable]):

310 min_run_time = 1 if not NVTX_PROFILE else 0.1

311 res = TBenchmark.Timer(

312 stmt="""

313 for fn in fns:

314 fn()

315 """,

287 - globals={

288 - "fns": fns

289 - },

316 + globals={"fns": fns},

317 label=label,

318 sub_label=sub_label,

319 description=description,

320 ).blocked_autorange(min_run_time=min_run_time)

321

322 if NVTX_PROFILE:

296 - with nvtx.annotate("mm-bench"), nvtx.annotate(

297 - f"{label}|{sub_label}|{description}"):

323 + with (

324 + nvtx.annotate("mm-bench"),

325 + nvtx.annotate(f"{label}|{sub_label}|{description}"),

326 + ):

327 fns[0]()

328

329 return res

@@ -304,19 +333,20 @@ _SWEEP_SCHEDULES_RESULTS: Optional[pd.DataFrame] = None

333 _SWEEP_SCHEDULES_RESULTS_CSV: Optional[str] = None

334

335

307 -def bench(types: TypeConfig,

308 - group_size: int,

309 - m: int,

310 - k: int,

311 - n: int,

312 - label: str,

313 - sub_label: str,

314 - sweep_schedules: bool = True) -> list[TMeasurement]:

336 +def bench(

337 + types: TypeConfig,

338 + group_size: int,

339 + m: int,

340 + k: int,

341 + n: int,

342 + label: str,

343 + sub_label: str,

344 + sweep_schedules: bool = True,

345 +) -> list[TMeasurement]:

346 benchmark_tensors = create_bench_tensors((m, n, k), types, group_size)

347 sub_label += f", L={len(benchmark_tensors)}"

348

318 - name_type_string = f"W{types.weight_type}"+\

319 - f"-A{terse_type_name(types.act_type)}"

349 + name_type_string = f"W{types.weight_type}" + f"-A{terse_type_name(types.act_type)}"

350 if types.group_scale_type is not None:

351 name_type_string += f"-GS{terse_type_name(types.group_scale_type)}"

352 if types.group_zero_type is not None:

@@ -332,31 +362,45 @@ def bench(types: TypeConfig,

362 # pytorch impl

363 timers.append(

364 bench_fns(

335 - label, sub_label, "torch.matmul (fp16)",

336 - [torch_matmul_f16_create_bench_fn(bt)

337 - for bt in benchmark_tensors]))

365 + label,

366 + sub_label,

367 + "torch.matmul (fp16)",

368 + [torch_matmul_f16_create_bench_fn(bt) for bt in benchmark_tensors],

369 + )

370 + )

371

372 if types.act_type == torch.int8 or types.act_type == torch.float8_e4m3fn:

373 timers.append(

374 bench_fns(

342 - label, sub_label,

343 - f"cutlass_scaled_mm ({terse_type_name(types.act_type)})", [

344 - cutlass_scaled_mm_create_bench_fn(bt)

345 - for bt in benchmark_tensors

346 - ]))

375 + label,

376 + sub_label,

377 + f"cutlass_scaled_mm ({terse_type_name(types.act_type)})",

378 + [cutlass_scaled_mm_create_bench_fn(bt) for bt in benchmark_tensors],

379 + )

380 + )

381

382 if types.act_type != torch.float8_e4m3fn:

383 timers.append(

350 - bench_fns(label, sub_label, f"marlin ({name_type_string})",

351 - [marlin_create_bench_fn(bt)

352 - for bt in benchmark_tensors]))

384 + bench_fns(

385 + label,

386 + sub_label,

387 + f"marlin ({name_type_string})",

388 + [marlin_create_bench_fn(bt) for bt in benchmark_tensors],

389 + )

390 + )

391

392 # machete

393 timers.append(

356 - bench_fns(label, sub_label, f"machete ({name_type_string})", [

357 - machete_create_bench_fn(bt, out_type=types.output_type)

358 - for bt in benchmark_tensors

359 - ]))

394 + bench_fns(

395 + label,

396 + sub_label,

397 + f"machete ({name_type_string})",

398 + [

399 + machete_create_bench_fn(bt, out_type=types.output_type)

400 + for bt in benchmark_tensors

401 + ],

402 + )

403 + )

404

405 if sweep_schedules:

406 global _SWEEP_SCHEDULES_RESULTS

@@ -371,7 +415,8 @@ def bench(types: TypeConfig,

415 group_zeros_type=types.group_zero_type,

416 token_scales_type=types.token_scale_type,

417 channel_scales_type=types.channel_scale_type,

374 - out_type=types.output_type)

418 + out_type=types.output_type,

419 + )

420

421 if schedules is None or len(schedules) == 0:

422 raise ValueError("No schedules found to sweep")

@@ -383,11 +428,17 @@ def bench(types: TypeConfig,

428 if schedule_M >= 2 * max(m, 16) or schedule_M < m // 4:

429 continue

430

386 - res = bench_fns(label, sub_label, "machete_best", [

387 - machete_create_bench_fn(

388 - bt, out_type=types.output_type, schedule=schedule)

389 - for bt in benchmark_tensors

390 - ])

431 + res = bench_fns(

432 + label,

433 + sub_label,

434 + "machete_best",

435 + [

436 + machete_create_bench_fn(

437 + bt, out_type=types.output_type, schedule=schedule

438 + )

439 + for bt in benchmark_tensors

440 + ],

441 + )

442

443 results_row = {

444 "M": m,

@@ -398,10 +449,8 @@ def bench(types: TypeConfig,

449 "median": res.median,

450 }

451 if _SWEEP_SCHEDULES_RESULTS is None:

401 - _SWEEP_SCHEDULES_RESULTS = pd.DataFrame(

402 - columns=results_row.keys())

403 - _SWEEP_SCHEDULES_RESULTS.\

404 - loc[len(_SWEEP_SCHEDULES_RESULTS)] = results_row

452 + _SWEEP_SCHEDULES_RESULTS = pd.DataFrame(columns=results_row.keys())

453 + _SWEEP_SCHEDULES_RESULTS.loc[len(_SWEEP_SCHEDULES_RESULTS)] = results_row

454

455 print(f" {res.median:5.5} ", schedule)

456 if not best or res.median < best.median:

@@ -422,8 +471,9 @@ def print_timers(timers: list[TMeasurement]):

471 def run(args, MKNs: Iterable[tuple[int, int, int]]) -> Iterable[TMeasurement]:

472 types = TypeConfig(

473 act_type=args.act_type,

425 - weight_type=scalar_types.uint4b8 if args.group_zero_type is None \

426 - else scalar_types.uint4,

474 + weight_type=scalar_types.uint4b8

475 + if args.group_zero_type is None

476 + else scalar_types.uint4,

477 output_type=args.out_type,

478 group_scale_type=args.group_scale_type,

479 group_zero_type=args.group_zero_type,

@@ -433,14 +483,16 @@ def run(args, MKNs: Iterable[tuple[int, int, int]]) -> Iterable[TMeasurement]:

483

484 results: list[TMeasurement] = []

485 for m, k, n in MKNs:

436 - timers = bench(types,

437 - args.group_size,

438 - m,

439 - k,

440 - n,

441 - f"{args.act_type}-gemm",

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

443 - sweep_schedules=args.sweep_schedules)

486 + timers = bench(

487 + types,

488 + args.group_size,

489 + m,

490 + k,

491 + n,

492 + f"{args.act_type}-gemm",

493 + f"MKN=({m}x{k}x{n})",

494 + sweep_schedules=args.sweep_schedules,

495 + )

496 print_timers(timers)

497 results.extend(timers)

498

@@ -454,7 +506,6 @@ def make_output(

506 base_description: str,

507 timestamp=None,

508 ):

457 -

509 print(f"== All Results {base_description} ====")

510 print_timers(data)

511

@@ -468,8 +519,7 @@ def make_output(

519

520

521 def run_square_bench(args):

471 - dim_sizes = list(

472 - range(args.dim_start, args.dim_end + 1, args.dim_increment))

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

523 MKNs = list(zip(dim_sizes, dim_sizes, dim_sizes))

524 data = run(args.dtype, args.sweep_schedules, MKNs)

525

@@ -479,8 +529,9 @@ def run_square_bench(args):

529 def run_range_bench(args):

530 m_start, k_start, n_start = (int(x) for x in args.dim_start.split(","))

531 m_end, k_end, n_end = (int(x) for x in args.dim_end.split(","))

482 - m_increment, k_increment, n_increment = \

483 - (int(x) for x in args.dim_increment.split(","))

532 + m_increment, k_increment, n_increment = (

533 + int(x) for x in args.dim_increment.split(",")

534 + )

535 Ms = list(range(m_start, m_end + 1, m_increment))

536 Ks = list(range(k_start, k_end + 1, k_increment))

537 Ns = list(range(n_start, n_end + 1, n_increment))

@@ -492,7 +543,6 @@ def run_range_bench(args):

543

544

545 def run_model_bench(args):

495 -

546 print("Benchmarking models:")

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

548 print(f"[{i}] {model}")

@@ -535,10 +585,13 @@ def run_model_bench(args):

585 with open(f"model_bench-{type_string}-{timestr}.pkl", "wb") as f:

586 args_dict = vars(args)

587 args_dict.pop("func")

538 - pkl.dump({

539 - "args": args_dict,

540 - "results": all_results,

541 - }, f)

588 + pkl.dump(

589 + {

590 + "args": args_dict,

591 + "results": all_results,

592 + },

593 + f,

594 + )

595

596

597 if __name__ == "__main__":

@@ -554,7 +607,6 @@ if __name__ == "__main__":

607 }[dt]

608

609 class ToTorchDtype(argparse.Action):

557 -

610 def __call__(self, parser, namespace, values, option_string=None):

611 setattr(namespace, self.dest, to_torch_dtype(values))

612

@@ -580,32 +632,32 @@ Benchmark Machete GEMM.

632 "--act-type",

633 action=ToTorchDtype,

634 required=True,

583 - choices=['bfloat16', 'float16', 'int8', 'float8_e4m3fn'],

635 + choices=["bfloat16", "float16", "int8", "float8_e4m3fn"],

636 )

637 parser.add_argument(

638 "--group-scale-type",

639 action=ToTorchDtype,

588 - choices=['bfloat16', 'float16'],

640 + choices=["bfloat16", "float16"],

641 )

642 parser.add_argument(

643 "--group-zero-type",

644 type=to_torch_dtype,

593 - choices=['bfloat16', 'float16'],

645 + choices=["bfloat16", "float16"],

646 )

647 parser.add_argument(

648 "--channel-scale-type",

649 action=ToTorchDtype,

598 - choices=['float'],

650 + choices=["float"],

651 )

652 parser.add_argument(

653 "--token-scale-type",

654 action=ToTorchDtype,

603 - choices=['float'],

655 + choices=["float"],

656 )

657 parser.add_argument(

658 "--out-type",

659 action=ToTorchDtype,

608 - choices=['bfloat16', 'float16'],

660 + choices=["bfloat16", "float16"],

661 )

662 parser.add_argument(

663 "--group-size",

@@ -618,9 +670,11 @@ Benchmark Machete GEMM.

670 action="store_true",

671 help="Run a sweep over all supported schedules",

672 )

621 - parser.add_argument("--sweep-csv-out",

622 - help="CSV to store sweep results",

623 - default="sch_sweep_results.csv")

673 + parser.add_argument(

674 + "--sweep-csv-out",

675 + help="CSV to store sweep results",

676 + default="sch_sweep_results.csv",

677 + )

678 subparsers = parser.add_subparsers(dest="cmd", required=True)

679

680 square_parser = subparsers.add_parser("square_bench")

@@ -634,17 +688,20 @@ Benchmark Machete GEMM.

688 "--dim-start",

689 type=str,

690 required=True,

637 - help="Start value for M,K,N as common separated list")

691 + help="Start value for M,K,N as common separated list",

692 + )

693 range_parser.add_argument(

694 "--dim-end",

695 type=str,

696 required=True,

642 - help="End value (inclusive) for M,K,N as common separated list")

697 + help="End value (inclusive) for M,K,N as common separated list",

698 + )

699 range_parser.add_argument(

700 "--dim-increment",

701 type=str,

702 required=True,

647 - help="Increment value for M,K,N as common separated list")

703 + help="Increment value for M,K,N as common separated list",

704 + )

705 range_parser.set_defaults(func=run_range_bench)

706

707 model_parser = subparsers.add_parser("model_bench")

@@ -655,14 +712,12 @@ Benchmark Machete GEMM.

712 default=DEFAULT_MODELS,

713 choices=WEIGHT_SHAPES.keys(),

714 )

658 - model_parser.add_argument("--tp-sizes",

659 - nargs="+",

660 - type=int,

661 - default=DEFAULT_TP_SIZES)

662 - model_parser.add_argument("--batch-sizes",

663 - nargs="+",

664 - type=int,

665 - default=DEFAULT_BATCH_SIZES)

715 + model_parser.add_argument(

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

717 + )

718 + model_parser.add_argument(

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

720 + )

721 model_parser.set_defaults(func=run_model_bench)

722

723 args = parser.parse_args()

benchmarks/kernels/benchmark_marlin.py

+117 -73

@@ -6,19 +6,34 @@ from benchmark_shapes import WEIGHT_SHAPES

7 from vllm import _custom_ops as ops

8 from vllm.model_executor.layers.quantization.gptq_marlin_24 import (

9 - GPTQ_MARLIN_24_MAX_PARALLEL, GPTQ_MARLIN_24_MIN_THREAD_N,

10 - GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES, GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES)

9 + GPTQ_MARLIN_24_MAX_PARALLEL,

10 + GPTQ_MARLIN_24_MIN_THREAD_N,

11 + GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES,

12 + GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES,

13 +)

14 from vllm.model_executor.layers.quantization.utils.allspark_utils import (

12 - ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD, ALLSPARK_SUPPORTED_QUANT_TYPES)

15 + ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,

16 + ALLSPARK_SUPPORTED_QUANT_TYPES,

17 +)

18 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

14 - GPTQ_MARLIN_MAX_PARALLEL, GPTQ_MARLIN_MIN_THREAD_N,

15 - MARLIN_SUPPORTED_GROUP_SIZES, query_marlin_supported_quant_types)

19 + GPTQ_MARLIN_MAX_PARALLEL,

20 + GPTQ_MARLIN_MIN_THREAD_N,

21 + MARLIN_SUPPORTED_GROUP_SIZES,

22 + query_marlin_supported_quant_types,

23 +)

24 from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (

17 - MarlinWorkspace, marlin_quantize)

25 + MarlinWorkspace,

26 + marlin_quantize,

27 +)

28 from vllm.model_executor.layers.quantization.utils.marlin_utils_test_24 import (

19 - marlin_24_quantize)

29 + marlin_24_quantize,

30 +)

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

21 - gptq_pack, gptq_quantize_weights, quantize_weights, sort_weights)

32 + gptq_pack,

33 + gptq_quantize_weights,

34 + quantize_weights,

35 + sort_weights,

36 +)

37 from vllm.scalar_type import ScalarType

38 from vllm.utils import FlexibleArgumentParser

@@ -29,22 +44,29 @@ ACT_ORDER_OPTS = [False, True]

44 K_FULL_OPTS = [False, True]

32 -def bench_run(results: list[benchmark.Measurement], model: str,

33 - act_order: bool, is_k_full: bool, quant_type: ScalarType,

34 - group_size: int, size_m: int, size_k: int, size_n: int):

47 +def bench_run(

48 + results: list[benchmark.Measurement],

49 + model: str,

50 + act_order: bool,

51 + is_k_full: bool,

52 + quant_type: ScalarType,

53 + group_size: int,

54 + size_m: int,

55 + size_k: int,

56 + size_n: int,

57 +):

58 label = "Quant Matmul"

37 - sub_label = ("{}, act={} k_full={}, q={}, g={}, "

38 - "MKN=({}x{}x{})".format(model, act_order, is_k_full,

39 - str(quant_type), group_size, size_m,

40 - size_k, size_n))

60 + sub_label = "{}, act={} k_full={}, q={}, g={}, MKN=({}x{}x{})".format(

61 + model, act_order, is_k_full, str(quant_type), group_size, size_m, size_k, size_n

62 + )

64 print(f"Testing: {sub_label}")

66 a = torch.randn(size_m, size_k).to(torch.half).cuda()

67 b = torch.rand(size_k, size_n).to(torch.half).cuda()

47 - a_tmp = (torch.zeros(size_m, size_k).to(torch.half).cuda())

69 + a_tmp = torch.zeros(size_m, size_k).to(torch.half).cuda()

71 # Marlin quant

72 (

@@ -57,14 +79,16 @@ def bench_run(results: list[benchmark.Measurement], model: str,

79 ) = marlin_quantize(b, quant_type, group_size, act_order)

81 # Marlin_24 quant

60 - (marlin_24_w_ref, marlin_24_q_w_comp, marlin_24_meta,

61 - marlin_24_s) = marlin_24_quantize(b, quant_type, group_size)

82 + (marlin_24_w_ref, marlin_24_q_w_comp, marlin_24_meta, marlin_24_s) = (

83 + marlin_24_quantize(b, quant_type, group_size)

84 + )

86 marlin_zp = torch.empty(0, dtype=torch.int, device=b.device)

88 # GPTQ quant

66 - (w_ref, q_w, s, g_idx,

67 - rand_perm) = gptq_quantize_weights(b, quant_type, group_size, act_order)

89 + (w_ref, q_w, s, g_idx, rand_perm) = gptq_quantize_weights(

90 + b, quant_type, group_size, act_order

91 + )

92 q_w_gptq = gptq_pack(q_w, quant_type.size_bits, size_k, size_n)

94 # For act_order, sort the "weights" and "g_idx"

@@ -74,32 +98,37 @@ def bench_run(results: list[benchmark.Measurement], model: str,

98 (q_w, g_idx, repack_sort_indices) = sort_weights(q_w, g_idx)

100 # Prepare

77 - marlin_workspace = MarlinWorkspace(size_n, GPTQ_MARLIN_MIN_THREAD_N,

78 - GPTQ_MARLIN_MAX_PARALLEL)

101 + marlin_workspace = MarlinWorkspace(

102 + size_n, GPTQ_MARLIN_MIN_THREAD_N, GPTQ_MARLIN_MAX_PARALLEL

103 + )

104

80 - marlin_24_workspace = MarlinWorkspace(size_n, GPTQ_MARLIN_24_MIN_THREAD_N,

81 - GPTQ_MARLIN_24_MAX_PARALLEL)

105 + marlin_24_workspace = MarlinWorkspace(

106 + size_n, GPTQ_MARLIN_24_MIN_THREAD_N, GPTQ_MARLIN_24_MAX_PARALLEL

107 + )

108 marlin_zp = torch.zeros_like(marlin_s, dtype=torch.int)

109

110 # AllSpark W8A16 quant

85 - as_supported_case = (quant_type in ALLSPARK_SUPPORTED_QUANT_TYPES

86 - and group_size == -1 and not act_order and is_k_full)

111 + as_supported_case = (

112 + quant_type in ALLSPARK_SUPPORTED_QUANT_TYPES

113 + and group_size == -1

114 + and not act_order

115 + and is_k_full

116 + )

117 if as_supported_case:

118 properties = torch.cuda.get_device_properties(b.device.index)

119 sm_count = properties.multi_processor_count

120 sm_version = properties.major * 10 + properties.minor

121

92 - supported_arch = (sm_version >= 80 and sm_version < 90)

122 + supported_arch = sm_version >= 80 and sm_version < 90

123 as_supported_case = as_supported_case and supported_arch

124 if supported_arch:

125 has_zp = False

96 - w_ref, qw, s, zp = quantize_weights(b, quant_type, group_size,

97 - has_zp)

126 + w_ref, qw, s, zp = quantize_weights(b, quant_type, group_size, has_zp)

127 qw = qw.to(torch.uint8)

128

100 - qw_reorder, s_reorder, zp_reorder = \

101 - ops.allspark_repack_weight(

102 - qw, s, zp, has_zp)

129 + qw_reorder, s_reorder, zp_reorder = ops.allspark_repack_weight(

130 + qw, s, zp, has_zp

131 + )

132 CUBLAS_M_THRESHOLD = ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD

133

134 globals = {

@@ -136,8 +165,7 @@ def bench_run(results: list[benchmark.Measurement], model: str,

165 "zp_reorder": zp_reorder if as_supported_case else None,

166 "sm_count": sm_count if as_supported_case else None,

167 "sm_version": sm_version if as_supported_case else None,

139 - "CUBLAS_M_THRESHOLD":

140 - CUBLAS_M_THRESHOLD if as_supported_case else None,

168 + "CUBLAS_M_THRESHOLD": CUBLAS_M_THRESHOLD if as_supported_case else None,

169 # Kernels

170 "gptq_marlin_gemm": ops.gptq_marlin_gemm,

171 "gptq_marlin_24_gemm": ops.gptq_marlin_24_gemm,

@@ -158,60 +186,63 @@ def bench_run(results: list[benchmark.Measurement], model: str,

186 label=label,

187 sub_label=sub_label,

188 description="pytorch_gemm",

161 - ).blocked_autorange(min_run_time=min_run_time))

189 + ).blocked_autorange(min_run_time=min_run_time)

190 + )

191

192 results.append(

193 benchmark.Timer(

165 - stmt=

166 - "output = gptq_marlin_gemm(a, marlin_q_w, marlin_s, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, False, False)", # noqa: E501

194 + stmt="output = gptq_marlin_gemm(a, marlin_q_w, marlin_s, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, False, False)", # noqa: E501

195 globals=globals,

196 label=label,

197 sub_label=sub_label,

198 description="gptq_marlin_gemm_fp16",

171 - ).blocked_autorange(min_run_time=min_run_time))

199 + ).blocked_autorange(min_run_time=min_run_time)

200 + )

201

202 results.append(

203 benchmark.Timer(

175 - stmt=

176 - "output = gptq_marlin_gemm(a, marlin_q_w, marlin_s, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, True, False)", # noqa: E501

204 + stmt="output = gptq_marlin_gemm(a, marlin_q_w, marlin_s, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, True, False)", # noqa: E501

205 globals=globals,

206 label=label,

207 sub_label=sub_label,

208 description="gptq_marlin_gemm_fp32",

181 - ).blocked_autorange(min_run_time=min_run_time))

209 + ).blocked_autorange(min_run_time=min_run_time)

210 + )

211

183 - if (quant_type in GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES

184 - and group_size in GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES):

212 + if (

213 + quant_type in GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES

214 + and group_size in GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES

215 + ):

216 results.append(

217 benchmark.Timer(

187 - stmt=

188 - "output = gptq_marlin_24_gemm(a, marlin_24_q_w_comp, marlin_24_meta, marlin_24_s, marlin_24_workspace.scratch, quant_type, size_m, size_n, size_k)", # noqa: E501

218 + stmt="output = gptq_marlin_24_gemm(a, marlin_24_q_w_comp, marlin_24_meta, marlin_24_s, marlin_24_workspace.scratch, quant_type, size_m, size_n, size_k)", # noqa: E501

219 globals=globals,

220 label=label,

221 sub_label=sub_label,

222 description="gptq_marlin_24_gemm",

193 - ).blocked_autorange(min_run_time=min_run_time))

223 + ).blocked_autorange(min_run_time=min_run_time)

224 + )

225

226 results.append(

227 benchmark.Timer(

197 - stmt=

198 - "q_res = gptq_marlin_repack(q_w_gptq, repack_sort_indices, size_k, size_n, quant_type.size_bits)", # noqa: E501

228 + stmt="q_res = gptq_marlin_repack(q_w_gptq, repack_sort_indices, size_k, size_n, quant_type.size_bits)", # noqa: E501

229 globals=globals,

230 label=label,

231 sub_label=sub_label,

232 description="gptq_marlin_repack",

203 - ).blocked_autorange(min_run_time=min_run_time))

233 + ).blocked_autorange(min_run_time=min_run_time)

234 + )

235

236 if as_supported_case:

237 results.append(

238 benchmark.Timer(

208 - stmt=

209 - "output = allspark_w8a16_gemm(a, qw_reorder, s_reorder, zp_reorder, size_n, group_size, sm_count, sm_version, CUBLAS_M_THRESHOLD, False, True)", # noqa: E501

239 + stmt="output = allspark_w8a16_gemm(a, qw_reorder, s_reorder, zp_reorder, size_n, group_size, sm_count, sm_version, CUBLAS_M_THRESHOLD, False, True)", # noqa: E501

240 globals=globals,

241 label=label,

242 sub_label=sub_label,

243 description="allspark_w8a16_gemm_fp32",

214 - ).blocked_autorange(min_run_time=min_run_time))

244 + ).blocked_autorange(min_run_time=min_run_time)

245 + )

246

247

248 def main(args):

@@ -233,37 +264,50 @@ def main(args):

264 continue

265

266 for act_order in ACT_ORDER_OPTS:

236 - if len(args.limit_act_order

237 - ) > 0 and act_order not in args.limit_act_order:

267 + if (

268 + len(args.limit_act_order) > 0

269 + and act_order not in args.limit_act_order

270 + ):

271 continue

272

273 for is_k_full in K_FULL_OPTS:

241 - if len(args.limit_k_full

242 - ) > 0 and is_k_full not in args.limit_k_full:

274 + if (

275 + len(args.limit_k_full) > 0

276 + and is_k_full not in args.limit_k_full

277 + ):

278 continue

279

245 - for quant_type in query_marlin_supported_quant_types(

246 - False):

247 - if len(args.limit_num_bits) > 0 and \

248 - quant_type.size_bits not in args.limit_num_bits:

280 + for quant_type in query_marlin_supported_quant_types(False):

281 + if (

282 + len(args.limit_num_bits) > 0

283 + and quant_type.size_bits not in args.limit_num_bits

284 + ):

285 continue

286

287 for group_size in MARLIN_SUPPORTED_GROUP_SIZES:

252 - if len(

253 - args.limit_group_size

254 - ) > 0 and group_size not in args.limit_group_size:

288 + if (

289 + len(args.limit_group_size) > 0

290 + and group_size not in args.limit_group_size

291 + ):

292 continue

293

294 # For act_order, the group_size must be less than

295 # size_k

259 - if act_order and (group_size == size_k

260 - or group_size == -1):

296 + if act_order and (group_size == size_k or group_size == -1):

297 continue

298

299 for size_m in args.batch_sizes:

264 - bench_run(results, model, act_order, is_k_full,

265 - quant_type, group_size, size_m,

266 - size_k, size_n)

300 + bench_run(

301 + results,

302 + model,

303 + act_order,

304 + is_k_full,

305 + quant_type,

306 + group_size,

307 + size_m,

308 + size_k,

309 + size_n,

310 + )

311

312 compare = benchmark.Compare(results)

313 compare.print()

@@ -274,7 +318,8 @@ def main(args):

318 #

319 if __name__ == "__main__":

320 parser = FlexibleArgumentParser(

277 - description="Benchmark Marlin across specified models/shapes/batches")

321 + description="Benchmark Marlin across specified models/shapes/batches"

322 + )

323 parser.add_argument(

324 "--models",

325 nargs="+",

@@ -282,10 +327,9 @@ if __name__ == "__main__":

327 default=DEFAULT_MODELS,

328 choices=WEIGHT_SHAPES.keys(),

329 )

285 - parser.add_argument("--batch-sizes",

286 - nargs="+",

287 - type=int,

288 - default=DEFAULT_BATCH_SIZES)

330 + parser.add_argument(

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

332 + )

333 parser.add_argument("--limit-k", nargs="+", type=int, default=[])

334 parser.add_argument("--limit-n", nargs="+", type=int, default=[])

335 parser.add_argument("--limit-group-size", nargs="+", type=int, default=[])

benchmarks/kernels/benchmark_moe.py

+237 -165

@@ -31,56 +31,60 @@ class BenchmarkConfig(TypedDict):

31 num_stages: int

34 -def benchmark_config(config: BenchmarkConfig,

35 - num_tokens: int,

36 - num_experts: int,

37 - shard_intermediate_size: int,

38 - hidden_size: int,

39 - topk: int,

40 - dtype: torch.dtype,

41 - use_fp8_w8a8: bool,

42 - use_int8_w8a16: bool,

43 - num_iters: int = 100,

44 - block_quant_shape: List[int] = None,

45 - use_deep_gemm: bool = False) -> float:

34 +def benchmark_config(

35 + config: BenchmarkConfig,

36 + num_tokens: int,

37 + num_experts: int,

38 + shard_intermediate_size: int,

39 + hidden_size: int,

40 + topk: int,

41 + dtype: torch.dtype,

42 + use_fp8_w8a8: bool,

43 + use_int8_w8a16: bool,

44 + num_iters: int = 100,

45 + block_quant_shape: List[int] = None,

46 + use_deep_gemm: bool = False,

47 +) -> float:

48 init_dtype = torch.float16 if use_fp8_w8a8 else dtype

49 x = torch.randn(num_tokens, hidden_size, dtype=dtype)

50 if use_int8_w8a16:

49 - w1 = torch.randint(-127,

50 - 127, (

51 - num_experts,

52 - shard_intermediate_size,

53 - hidden_size,

54 - ),

55 - dtype=torch.int8)

56 - w2 = torch.randint(-127,

57 - 127, (

58 - num_experts,

59 - hidden_size,

60 - shard_intermediate_size // 2,

61 - ),

62 - dtype=torch.int8)

51 + w1 = torch.randint(

52 + -127,

53 + 127,

54 + (

55 + num_experts,

56 + shard_intermediate_size,

57 + hidden_size,

58 + ),

59 + dtype=torch.int8,

60 + )

61 + w2 = torch.randint(

62 + -127,

63 + 127,

64 + (

65 + num_experts,

66 + hidden_size,

67 + shard_intermediate_size // 2,

68 + ),

69 + dtype=torch.int8,

70 + )

71 else:

64 - w1 = torch.randn(num_experts,

65 - shard_intermediate_size,

66 - hidden_size,

67 - dtype=init_dtype)

68 - w2 = torch.randn(num_experts,

69 - hidden_size,

70 - shard_intermediate_size // 2,

71 - dtype=init_dtype)

72 - gating_output = torch.randn(num_iters,

73 - num_tokens,

74 - num_experts,

75 - dtype=torch.float32)

72 + w1 = torch.randn(

73 + num_experts, shard_intermediate_size, hidden_size, dtype=init_dtype

74 + )

75 + w2 = torch.randn(

76 + num_experts, hidden_size, shard_intermediate_size // 2, dtype=init_dtype

77 + )

78 + gating_output = torch.randn(num_iters, num_tokens, num_experts, dtype=torch.float32)

80 w1_scale = None

81 w2_scale = None

82 a1_scale = None

83 a2_scale = None

84 if use_int8_w8a16:

82 - w1_scale = torch.randn((num_experts, 2 * shard_intermediate_size),

83 - dtype=torch.float32)

85 + w1_scale = torch.randn(

86 + (num_experts, 2 * shard_intermediate_size), dtype=torch.float32

87 + )

88 w2_scale = torch.randn((hidden_size, num_experts), dtype=torch.float32)

89 if use_fp8_w8a8:

90 if block_quant_shape:

@@ -93,10 +97,14 @@ def benchmark_config(config: BenchmarkConfig,

97 n_tiles_w2 = (K + block_n - 1) // block_n

98 k_tiles_w1 = (K + block_k - 1) // block_k

99 k_tiles_w2 = (N + block_k - 1) // block_k

96 - w1_scale = torch.rand((E, n_tiles_w1, k_tiles_w1),

97 - dtype=torch.float32) * factor_for_scale

98 - w2_scale = torch.rand((E, n_tiles_w2, k_tiles_w2),

99 - dtype=torch.float32) * factor_for_scale

100 + w1_scale = (

101 + torch.rand((E, n_tiles_w1, k_tiles_w1), dtype=torch.float32)

102 + * factor_for_scale

103 + )

104 + w2_scale = (

105 + torch.rand((E, n_tiles_w2, k_tiles_w2), dtype=torch.float32)

106 + * factor_for_scale

107 + )

108 else:

109 w1_scale = torch.randn(num_experts, dtype=torch.float32)

110 w2_scale = torch.randn(num_experts, dtype=torch.float32)

@@ -114,10 +122,12 @@ def benchmark_config(config: BenchmarkConfig,

122

123 def run():

124 from vllm.model_executor.layers.fused_moe import override_config

125 +

126 with override_config(config):

127 if use_deep_gemm:

128 topk_weights, topk_ids, token_expert_indices = fused_topk(

120 - x, input_gating, topk, False)

129 + x, input_gating, topk, False

130 + )

131 return fused_experts(

132 x,

133 w1,

@@ -213,8 +223,7 @@ def get_rocm_tuning_space(use_fp16):

223 return param_ranges

224

225

216 -def get_configs_compute_bound(use_fp16,

217 - block_quant_shape) -> list[dict[str, int]]:

226 +def get_configs_compute_bound(use_fp16, block_quant_shape) -> list[dict[str, int]]:

227 configs: list[BenchmarkConfig] = []

228

229 if current_platform.is_rocm():

@@ -250,20 +259,25 @@ def get_configs_compute_bound(use_fp16,

259 if block_quant_shape is not None and not use_fp16:

260 block_n, block_k = block_quant_shape[0], block_quant_shape[1]

261 for config in configs[:]:

253 - if config["BLOCK_SIZE_K"] % block_k != 0 or config[

254 - "BLOCK_SIZE_N"] % block_n != 0:

262 + if (

263 + config["BLOCK_SIZE_K"] % block_k != 0

264 + or config["BLOCK_SIZE_N"] % block_n != 0

265 + ):

266 configs.remove(config)

267 return configs

268

269

259 -def prune_rocm_search_space(num_tokens, shard_intermediate_size, hidden_size,

260 - search_space, is_fp16, topk):

270 +def prune_rocm_search_space(

271 + num_tokens, shard_intermediate_size, hidden_size, search_space, is_fp16, topk

272 +):

273 N1, K1 = shard_intermediate_size, hidden_size

274 N2, K2 = hidden_size, shard_intermediate_size // 2

263 - pruned_space_1 = prune_rocm_configs(num_tokens * topk, N1, K1,

264 - search_space, is_fp16)

265 - pruned_space_2 = prune_rocm_configs(num_tokens * topk, N2, K2,

266 - search_space, is_fp16)

275 + pruned_space_1 = prune_rocm_configs(

276 + num_tokens * topk, N1, K1, search_space, is_fp16

277 + )

278 + pruned_space_2 = prune_rocm_configs(

279 + num_tokens * topk, N2, K2, search_space, is_fp16

280 + )

281 search_space = merge_unique_dicts(pruned_space_1, pruned_space_2)

282 return search_space

283

@@ -301,14 +315,14 @@ def prune_rocm_configs(M, N, K, configs, is_fp16=True):

315 SPLIT_K = config.get("SPLIT_K", 1)

316 GROUP_M = config.get("GROUP_SIZE_M")

317 if is_fp16:

304 - if (matrix_instr_nonkdim > BLOCK_SIZE_M

305 - or matrix_instr_nonkdim > BLOCK_SIZE_N):

318 + if (

319 + matrix_instr_nonkdim > BLOCK_SIZE_M

320 + or matrix_instr_nonkdim > BLOCK_SIZE_N

321 + ):

322 continue

307 - if (matrix_instr_nonkdim >= M

308 - and matrix_instr_nonkdim != BLOCK_SIZE_M):

323 + if matrix_instr_nonkdim >= M and matrix_instr_nonkdim != BLOCK_SIZE_M:

324 continue

310 - if (matrix_instr_nonkdim >= N

311 - and matrix_instr_nonkdim != BLOCK_SIZE_N):

325 + if matrix_instr_nonkdim >= N and matrix_instr_nonkdim != BLOCK_SIZE_N:

326 continue

327 # Skip BLOCK_SIZE that is too large compare to M/N

328 # unless BLOCK_SIZE is already small enough

@@ -329,8 +343,10 @@ def prune_rocm_configs(M, N, K, configs, is_fp16=True):

343 continue

344 # out of shared memory resource

345 # TODO (zhanglx): This does not consider the LDS usage in the epilogue

332 - LDS = (BLOCK_SIZE_K * BLOCK_SIZE_M * elemBytes_a +

333 - BLOCK_SIZE_K * BLOCK_SIZE_N * elemBytes_b)

346 + LDS = (

347 + BLOCK_SIZE_K * BLOCK_SIZE_M * elemBytes_a

348 + + BLOCK_SIZE_K * BLOCK_SIZE_N * elemBytes_b

349 + )

350 if LDS > 65536:

351 continue

352 # Skip small block sizes and num_warps for large gemm

@@ -364,7 +380,6 @@ def merge_unique_dicts(list1, list2):

380

381 @ray.remote(num_gpus=1)

382 class BenchmarkWorker:

367 -

383 def __init__(self, seed: int) -> None:

384 torch.set_default_device("cuda")

385 current_platform.seed_everything(seed)

@@ -388,36 +403,40 @@ class BenchmarkWorker:

403 use_deep_gemm: bool = False,

404 ) -> tuple[dict[str, int], float]:

405 current_platform.seed_everything(self.seed)

391 - dtype_str = get_config_dtype_str(dtype,

392 - use_int8_w8a16=use_int8_w8a16,

393 - use_fp8_w8a8=use_fp8_w8a8)

406 + dtype_str = get_config_dtype_str(

407 + dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8

408 + )

409 # NOTE(woosuk): The current naming convention uses w2.shape[2], which

410 # is the intermediate size after silu_and_mul.

396 - op_config = get_moe_configs(num_experts, shard_intermediate_size // 2,

397 - dtype_str)

411 + op_config = get_moe_configs(

412 + num_experts, shard_intermediate_size // 2, dtype_str

413 + )

414 if op_config is None:

399 - config = get_default_config(num_tokens,

400 - num_experts,

401 - shard_intermediate_size,

402 - hidden_size,

403 - topk,

404 - dtype_str,

405 - is_marlin=False)

415 + config = get_default_config(

416 + num_tokens,

417 + num_experts,

418 + shard_intermediate_size,

419 + hidden_size,

420 + topk,

421 + dtype_str,

422 + is_marlin=False,

423 + )

424 else:

407 - config = op_config[min(op_config.keys(),

408 - key=lambda x: abs(x - num_tokens))]

409 - kernel_time = benchmark_config(config,

410 - num_tokens,

411 - num_experts,

412 - shard_intermediate_size,

413 - hidden_size,

414 - topk,

415 - dtype,

416 - use_fp8_w8a8,

417 - use_int8_w8a16,

418 - num_iters=100,

419 - block_quant_shape=block_quant_shape,

420 - use_deep_gemm=use_deep_gemm)

425 + config = op_config[min(op_config.keys(), key=lambda x: abs(x - num_tokens))]

426 + kernel_time = benchmark_config(

427 + config,

428 + num_tokens,

429 + num_experts,

430 + shard_intermediate_size,

431 + hidden_size,

432 + topk,

433 + dtype,

434 + use_fp8_w8a8,

435 + use_int8_w8a16,

436 + num_iters=100,

437 + block_quant_shape=block_quant_shape,

438 + use_deep_gemm=use_deep_gemm,

439 + )

440 return config, kernel_time

441

442 def tune(

@@ -438,10 +457,14 @@ class BenchmarkWorker:

457 best_time = float("inf")

458 if current_platform.is_rocm():

459 is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16)

441 - search_space = prune_rocm_search_space(num_tokens,

442 - shard_intermediate_size,

443 - hidden_size, search_space,

444 - is_fp16, topk)

460 + search_space = prune_rocm_search_space(

461 + num_tokens,

462 + shard_intermediate_size,

463 + hidden_size,

464 + search_space,

465 + is_fp16,

466 + topk,

467 + )

468

469 need_device_guard = False

470 if current_platform.is_rocm():

@@ -449,8 +472,7 @@ class BenchmarkWorker:

472 if visible_device != f"{self.device_id}":

473 need_device_guard = True

474

452 - with torch.cuda.device(

453 - self.device_id) if need_device_guard else nullcontext():

475 + with torch.cuda.device(self.device_id) if need_device_guard else nullcontext():

476 for config in tqdm(search_space):

477 try:

478 kernel_time = benchmark_config(

@@ -465,7 +487,8 @@ class BenchmarkWorker:

487 use_int8_w8a16,

488 num_iters=20,

489 block_quant_shape=block_quant_shape,

468 - use_deep_gemm=use_deep_gemm)

490 + use_deep_gemm=use_deep_gemm,

491 + )

492 except triton.runtime.autotuner.OutOfResources:

493 # Some configurations may be invalid and fail to compile.

494 continue

@@ -481,42 +504,44 @@ class BenchmarkWorker:

504

505 def sort_config(config: BenchmarkConfig) -> BenchmarkConfig:

506 return {

484 - "BLOCK_SIZE_M":

485 - config["BLOCK_SIZE_M"],

486 - "BLOCK_SIZE_N":

487 - config["BLOCK_SIZE_N"],

488 - "BLOCK_SIZE_K":

489 - config["BLOCK_SIZE_K"],

490 - "GROUP_SIZE_M":

491 - config["GROUP_SIZE_M"],

492 - "num_warps":

493 - config["num_warps"],

494 - "num_stages":

495 - config["num_stages"],

496 - **({

497 - "waves_per_eu": config["waves_per_eu"]

498 - } if "waves_per_eu" in config else {}),

499 - **({

500 - "matrix_instr_nonkdim": config["matrix_instr_nonkdim"]

501 - } if "matrix_instr_nonkdim" in config else {}),

502 - **({

503 - "kpack": config["kpack"]

504 - } if "kpack" in config else {}),

507 + "BLOCK_SIZE_M": config["BLOCK_SIZE_M"],

508 + "BLOCK_SIZE_N": config["BLOCK_SIZE_N"],

509 + "BLOCK_SIZE_K": config["BLOCK_SIZE_K"],

510 + "GROUP_SIZE_M": config["GROUP_SIZE_M"],

511 + "num_warps": config["num_warps"],

512 + "num_stages": config["num_stages"],

513 + **(

514 + {"waves_per_eu": config["waves_per_eu"]} if "waves_per_eu" in config else {}

515 + ),

516 + **(

517 + {"matrix_instr_nonkdim": config["matrix_instr_nonkdim"]}

518 + if "matrix_instr_nonkdim" in config

519 + else {}

520 + ),

521 + **({"kpack": config["kpack"]} if "kpack" in config else {}),

522 }

523

524

508 -def save_configs(configs: dict[int, BenchmarkConfig], num_experts: int,

509 - shard_intermediate_size: int, hidden_size: int, topk: int,

510 - dtype: torch.dtype, use_fp8_w8a8: bool, use_int8_w8a16: bool,

511 - block_quant_shape: List[int]) -> None:

512 - dtype_str = get_config_dtype_str(dtype,

513 - use_int8_w8a16=use_int8_w8a16,

514 - use_fp8_w8a8=use_fp8_w8a8)

525 +def save_configs(

526 + configs: dict[int, BenchmarkConfig],

527 + num_experts: int,

528 + shard_intermediate_size: int,

529 + hidden_size: int,

530 + topk: int,

531 + dtype: torch.dtype,

532 + use_fp8_w8a8: bool,

533 + use_int8_w8a16: bool,

534 + block_quant_shape: List[int],

535 +) -> None:

536 + dtype_str = get_config_dtype_str(

537 + dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8

538 + )

539

540 # NOTE(woosuk): The current naming convention uses w2.shape[2], which

541 # is the intermediate size after silu_and_mul.

518 - filename = get_config_file_name(num_experts, shard_intermediate_size // 2,

519 - dtype_str, block_quant_shape)

542 + filename = get_config_file_name(

543 + num_experts, shard_intermediate_size // 2, dtype_str, block_quant_shape

544 + )

545

546 print(f"Writing best config to {filename}...")

547 with open(filename, "w") as f:

@@ -525,18 +550,16 @@ def save_configs(configs: dict[int, BenchmarkConfig], num_experts: int,

550

551

552 def get_weight_block_size_safety(config, default_value=None):

528 -

529 - quantization_config = getattr(config, 'quantization_config', {})

553 + quantization_config = getattr(config, "quantization_config", {})

554 if isinstance(quantization_config, dict):

531 - return quantization_config.get('weight_block_size', default_value)

555 + return quantization_config.get("weight_block_size", default_value)

556 return default_value

557

558

559 def main(args: argparse.Namespace):

560 print(args)

561

538 - config = get_config(model=args.model,

539 - trust_remote_code=args.trust_remote_code)

562 + config = get_config(model=args.model, trust_remote_code=args.trust_remote_code)

563 if args.model_prefix:

564 config = getattr(config, args.model_prefix)

565 config = SimpleNamespace(**config)

@@ -551,14 +574,12 @@ def main(args: argparse.Namespace):

574 topk = config.num_experts_per_tok

575 intermediate_size = config.intermediate_size

576 shard_intermediate_size = 2 * intermediate_size // args.tp_size

554 - elif (config.architectures[0]

555 - in ("DeepseekV3ForCausalLM", "DeepseekV2ForCausalLM")):

577 + elif config.architectures[0] in ("DeepseekV3ForCausalLM", "DeepseekV2ForCausalLM"):

578 E = config.n_routed_experts

579 topk = config.num_experts_per_tok

580 intermediate_size = config.moe_intermediate_size

581 shard_intermediate_size = 2 * intermediate_size // args.tp_size

560 - elif config.architectures[0] in ("Qwen2MoeForCausalLM",

561 - "Qwen3MoeForCausalLM"):

582 + elif config.architectures[0] in ("Qwen2MoeForCausalLM", "Qwen3MoeForCausalLM"):

583 E = config.num_experts

584 topk = config.num_experts_per_tok

585 intermediate_size = config.moe_intermediate_size

@@ -573,16 +594,35 @@ def main(args: argparse.Namespace):

594 shard_intermediate_size = 2 * intermediate_size // args.tp_size

595

596 hidden_size = config.hidden_size

576 - dtype = torch.float16 if current_platform.is_rocm() else getattr(

577 - torch, config.torch_dtype)

597 + dtype = (

598 + torch.float16

599 + if current_platform.is_rocm()

600 + else getattr(torch, config.torch_dtype)

601 + )

602 use_fp8_w8a8 = args.dtype == "fp8_w8a8"

603 use_int8_w8a16 = args.dtype == "int8_w8a16"

604 block_quant_shape = get_weight_block_size_safety(config)

605

606 if args.batch_size is None:

607 batch_sizes = [

584 - 1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 256, 512, 1024, 1536,

585 - 2048, 3072, 4096

608 + 1,

609 + 2,

610 + 4,

611 + 8,

612 + 16,

613 + 24,

614 + 32,

615 + 48,

616 + 64,

617 + 96,

618 + 128,

619 + 256,

620 + 512,

621 + 1024,

622 + 1536,

623 + 2048,

624 + 3072,

625 + 4096,

626 ]

627 else:

628 batch_sizes = [args.batch_size]

@@ -593,7 +633,8 @@ def main(args: argparse.Namespace):

633 # Ray will set ROCR_VISIBLE_DEVICES for device visibility

634 logger.warning(

635 "Ray uses ROCR_VISIBLE_DEVICES to control device accessibility."

596 - "Replacing HIP_VISIBLE_DEVICES with ROCR_VISIBLE_DEVICES.")

636 + "Replacing HIP_VISIBLE_DEVICES with ROCR_VISIBLE_DEVICES."

637 + )

638 val = os.environ["HIP_VISIBLE_DEVICES"]

639 os.environ["ROCR_VISIBLE_DEVICES"] = val

640 del os.environ["HIP_VISIBLE_DEVICES"]

@@ -620,25 +661,59 @@ def main(args: argparse.Namespace):

661

662 start = time.time()

663 configs = _distribute(

623 - "tune", [(batch_size, E, shard_intermediate_size, hidden_size,

624 - topk, dtype, use_fp8_w8a8, use_int8_w8a16, search_space,

625 - block_quant_shape, use_deep_gemm)

626 - for batch_size in batch_sizes])

664 + "tune",

665 + [

666 + (

667 + batch_size,

668 + E,

669 + shard_intermediate_size,

670 + hidden_size,

671 + topk,

672 + dtype,

673 + use_fp8_w8a8,

674 + use_int8_w8a16,

675 + search_space,

676 + block_quant_shape,

677 + use_deep_gemm,

678 + )

679 + for batch_size in batch_sizes

680 + ],

681 + )

682 best_configs = {

628 - M: sort_config(config)

629 - for M, config in zip(batch_sizes, configs)

683 + M: sort_config(config) for M, config in zip(batch_sizes, configs)

684 }

631 - save_configs(best_configs, E, shard_intermediate_size, hidden_size,

632 - topk, dtype, use_fp8_w8a8, use_int8_w8a16,

633 - block_quant_shape)

685 + save_configs(

686 + best_configs,

687 + E,

688 + shard_intermediate_size,

689 + hidden_size,

690 + topk,

691 + dtype,

692 + use_fp8_w8a8,

693 + use_int8_w8a16,

694 + block_quant_shape,

695 + )

696 end = time.time()

697 print(f"Tuning took {end - start:.2f} seconds")

698 else:

699 outputs = _distribute(

700 "benchmark",

639 - [(batch_size, E, shard_intermediate_size, hidden_size, topk, dtype,

640 - use_fp8_w8a8, use_int8_w8a16, block_quant_shape, use_deep_gemm)

641 - for batch_size in batch_sizes])

701 + [

702 + (

703 + batch_size,

704 + E,

705 + shard_intermediate_size,

706 + hidden_size,

707 + topk,

708 + dtype,

709 + use_fp8_w8a8,

710 + use_int8_w8a16,

711 + block_quant_shape,

712 + use_deep_gemm,

713 + )

714 + for batch_size in batch_sizes

715 + ],

716 + )

717

718 for batch_size, (config, kernel_time) in zip(batch_sizes, outputs):

719 print(f"Batch size: {batch_size}, config: {config}")

@@ -647,18 +722,15 @@ def main(args: argparse.Namespace):

722

723 if __name__ == "__main__":

724 parser = FlexibleArgumentParser()

650 - parser.add_argument("--model",

651 - type=str,

652 - default="mistralai/Mixtral-8x7B-Instruct-v0.1")

653 - parser.add_argument("--tp-size",

654 - "-tp",

655 - "--tensor-parallel-size",

656 - type=int,

657 - default=2)

658 - parser.add_argument("--dtype",

659 - type=str,

660 - choices=["auto", "fp8_w8a8", "int8_w8a16"],

661 - default="auto")

725 + parser.add_argument(

726 + "--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1"

727 + )

728 + parser.add_argument(

729 + "--tp-size", "-tp", "--tensor-parallel-size", type=int, default=2

730 + )

731 + parser.add_argument(

732 + "--dtype", type=str, choices=["auto", "fp8_w8a8", "int8_w8a16"], default="auto"

733 + )

734 parser.add_argument("--use-deep-gemm", action="store_true")

735 parser.add_argument("--seed", type=int, default=0)

736 parser.add_argument("--batch-size", type=int, required=False)

benchmarks/kernels/benchmark_moe_permute_unpermute.py

+160 -93

@@ -8,7 +8,9 @@ import torch

8 from transformers import AutoConfig

10 from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (

11 - _moe_permute, _moe_unpermute_and_reduce)

11 + _moe_permute,

12 + _moe_unpermute_and_reduce,

13 +)

14 from vllm.model_executor.layers.fused_moe.fused_moe import *

15 from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import *

16 from vllm.model_executor.layers.fused_moe.utils import _fp8_quantize

@@ -27,15 +29,17 @@ class BenchmarkConfig(TypedDict):

29 num_stages: int

30 -def benchmark_permute(num_tokens: int,

31 - num_experts: int,

32 - hidden_size: int,

33 - topk: int,

34 - dtype: torch.dtype,

35 - use_fp8_w8a8: bool,

36 - use_int8_w8a16: bool,

37 - num_iters: int = 100,

38 - use_customized_permute: bool = False) -> float:

32 +def benchmark_permute(

33 + num_tokens: int,

34 + num_experts: int,

35 + hidden_size: int,

36 + topk: int,

37 + dtype: torch.dtype,

38 + use_fp8_w8a8: bool,

39 + use_int8_w8a16: bool,

40 + num_iters: int = 100,

41 + use_customized_permute: bool = False,

42 +) -> float:

43 # init_dtype = torch.float16 if use_fp8_w8a8 else dtype

44 hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype)

45 # output_hidden_states = torch.empty_like(hidden_states)

@@ -46,36 +50,41 @@ def benchmark_permute(num_tokens: int,

50 align_block_size = None

51 qhidden_states = hidden_states

49 - gating_output = torch.randn(num_iters,

50 - num_tokens,

51 - num_experts,

52 - dtype=torch.float32)

53 + gating_output = torch.randn(num_iters, num_tokens, num_experts, dtype=torch.float32)

55 input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32)

56 topk_weights, topk_ids, token_expert_indices = fused_topk(

56 - qhidden_states, input_gating, topk, False)

57 + qhidden_states, input_gating, topk, False

58 + )

60 def prepare(i: int):

61 input_gating.copy_(gating_output[i])

63 def run():

64 if use_customized_permute:

63 - (permuted_hidden_states, first_token_off, inv_perm_idx,

64 - m_indices) = moe_permute(

65 - qhidden_states,

66 - topk_weights=topk_weights,

67 - topk_ids=topk_ids,

68 - token_expert_indices=token_expert_indices,

69 - topk=topk,

70 - n_expert=num_experts,

71 - n_local_expert=num_experts,

72 - expert_map=None,

73 - align_block_size=align_block_size,

74 - )

65 + (permuted_hidden_states, first_token_off, inv_perm_idx, m_indices) = (

66 + moe_permute(

67 + qhidden_states,

68 + topk_weights=topk_weights,

69 + topk_ids=topk_ids,

70 + token_expert_indices=token_expert_indices,

71 + topk=topk,

72 + n_expert=num_experts,

73 + n_local_expert=num_experts,

74 + expert_map=None,

75 + align_block_size=align_block_size,

76 + )

77 + )

78 else:

76 - (permuted_hidden_states, a1q_scale, sorted_token_ids, expert_ids,

77 - inv_perm) = _moe_permute(qhidden_states, None, topk_ids,

78 - num_experts, None, align_block_size)

79 + (

80 + permuted_hidden_states,

81 + a1q_scale,

82 + sorted_token_ids,

83 + expert_ids,

84 + inv_perm,

85 + ) = _moe_permute(

86 + qhidden_states, None, topk_ids, num_experts, None, align_block_size

87 + )

89 # JIT compilation & warmup

90 run()

@@ -111,15 +120,17 @@ def benchmark_permute(num_tokens: int,

120 return avg

121

122

114 -def benchmark_unpermute(num_tokens: int,

115 - num_experts: int,

116 - hidden_size: int,

117 - topk: int,

118 - dtype: torch.dtype,

119 - use_fp8_w8a8: bool,

120 - use_int8_w8a16: bool,

121 - num_iters: int = 100,

122 - use_customized_permute: bool = False) -> float:

123 +def benchmark_unpermute(

124 + num_tokens: int,

125 + num_experts: int,

126 + hidden_size: int,

127 + topk: int,

128 + dtype: torch.dtype,

129 + use_fp8_w8a8: bool,

130 + use_int8_w8a16: bool,

131 + num_iters: int = 100,

132 + use_customized_permute: bool = False,

133 +) -> float:

134 # init_dtype = torch.float16 if use_fp8_w8a8 else dtype

135 hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype)

136 output_hidden_states = torch.empty_like(hidden_states)

@@ -133,46 +144,74 @@ def benchmark_unpermute(num_tokens: int,

144 input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32)

145

146 topk_weights, topk_ids, token_expert_indices = fused_topk(

136 - qhidden_states, input_gating, topk, False)

147 + qhidden_states, input_gating, topk, False

148 + )

149

150 def prepare():

151 if use_customized_permute:

140 - (permuted_hidden_states, first_token_off, inv_perm_idx,

141 - m_indices) = moe_permute(

142 - qhidden_states,

143 - topk_weights=topk_weights,

144 - topk_ids=topk_ids,

145 - token_expert_indices=token_expert_indices,

146 - topk=topk,

147 - n_expert=num_experts,

148 - n_local_expert=num_experts,

149 - expert_map=None,

150 - align_block_size=align_block_size,

151 - )

152 + (permuted_hidden_states, first_token_off, inv_perm_idx, m_indices) = (

153 + moe_permute(

154 + qhidden_states,

155 + topk_weights=topk_weights,

156 + topk_ids=topk_ids,

157 + token_expert_indices=token_expert_indices,

158 + topk=topk,

159 + n_expert=num_experts,

160 + n_local_expert=num_experts,

161 + expert_map=None,

162 + align_block_size=align_block_size,

163 + )

164 + )

165 # convert to fp16/bf16 as gemm output

153 - return (permuted_hidden_states.to(dtype), first_token_off,

154 - inv_perm_idx, m_indices)

166 + return (

167 + permuted_hidden_states.to(dtype),

168 + first_token_off,

169 + inv_perm_idx,

170 + m_indices,

171 + )

172 else:

156 - (permuted_qhidden_states, a1q_scale, sorted_token_ids, expert_ids,

157 - inv_perm) = _moe_permute(qhidden_states, None, topk_ids,

158 - num_experts, None, align_block_size)

173 + (

174 + permuted_qhidden_states,

175 + a1q_scale,

176 + sorted_token_ids,

177 + expert_ids,

178 + inv_perm,

179 + ) = _moe_permute(

180 + qhidden_states, None, topk_ids, num_experts, None, align_block_size

181 + )

182 # convert to fp16/bf16 as gemm output

160 - return (permuted_qhidden_states.to(dtype), a1q_scale,

161 - sorted_token_ids, expert_ids, inv_perm)

183 + return (

184 + permuted_qhidden_states.to(dtype),

185 + a1q_scale,

186 + sorted_token_ids,

187 + expert_ids,

188 + inv_perm,

189 + )

190

191 def run(input: tuple):

192 if use_customized_permute:

165 - (permuted_hidden_states, first_token_off, inv_perm_idx,

166 - m_indices) = input

167 - moe_unpermute(permuted_hidden_states, topk_weights, topk_ids,

168 - inv_perm_idx, first_token_off, topk, num_experts,

169 - num_experts)

193 + (permuted_hidden_states, first_token_off, inv_perm_idx, m_indices) = input

194 + moe_unpermute(

195 + permuted_hidden_states,

196 + topk_weights,

197 + topk_ids,

198 + inv_perm_idx,

199 + first_token_off,

200 + topk,

201 + num_experts,

202 + num_experts,

203 + )

204 else:

171 - (permuted_hidden_states, a1q_scale, sorted_token_ids, expert_ids,

172 - inv_perm) = input

173 - _moe_unpermute_and_reduce(output_hidden_states,

174 - permuted_hidden_states, inv_perm,

175 - topk_weights)

205 + (

206 + permuted_hidden_states,

207 + a1q_scale,

208 + sorted_token_ids,

209 + expert_ids,

210 + inv_perm,

211 + ) = input

212 + _moe_unpermute_and_reduce(

213 + output_hidden_states, permuted_hidden_states, inv_perm, topk_weights

214 + )

215

216 # JIT compilation & warmup

217 input = prepare()

@@ -209,7 +248,6 @@ def benchmark_unpermute(num_tokens: int,

248

249 @ray.remote(num_gpus=1)

250 class BenchmarkWorker:

212 -

251 def __init__(self, seed: int) -> None:

252 torch.set_default_device("cuda")

253 current_platform.seed_everything(seed)

@@ -241,7 +279,8 @@ class BenchmarkWorker:

279 use_fp8_w8a8,

280 use_int8_w8a16,

281 num_iters=100,

244 - use_customized_permute=use_customized_permute)

282 + use_customized_permute=use_customized_permute,

283 + )

284 unpermute_time = benchmark_unpermute(

285 num_tokens,

286 num_experts,

@@ -251,15 +290,15 @@ class BenchmarkWorker:

290 use_fp8_w8a8,

291 use_int8_w8a16,

292 num_iters=100,

254 - use_customized_permute=use_customized_permute)

293 + use_customized_permute=use_customized_permute,

294 + )

295 return permute_time, unpermute_time

296

297

298 def get_weight_block_size_safety(config, default_value=None):

259 -

260 - quantization_config = getattr(config, 'quantization_config', {})

299 + quantization_config = getattr(config, "quantization_config", {})

300 if isinstance(quantization_config, dict):

262 - return quantization_config.get('weight_block_size', default_value)

301 + return quantization_config.get("weight_block_size", default_value)

302 return default_value

303

304

@@ -267,20 +306,21 @@ def main(args: argparse.Namespace):

306 print(args)

307

308 config = AutoConfig.from_pretrained(

270 - args.model, trust_remote_code=args.trust_remote_code)

309 + args.model, trust_remote_code=args.trust_remote_code

310 + )

311 if config.architectures[0] == "DbrxForCausalLM":

312 E = config.ffn_config.moe_num_experts

313 topk = config.ffn_config.moe_top_k

314 elif config.architectures[0] == "JambaForCausalLM":

315 E = config.num_experts

316 topk = config.num_experts_per_tok

277 - elif (config.architectures[0] == "DeepseekV3ForCausalLM"

278 - or config.architectures[0] == "DeepseekV2ForCausalLM"):

317 + elif (

318 + config.architectures[0] == "DeepseekV3ForCausalLM"

319 + or config.architectures[0] == "DeepseekV2ForCausalLM"

320 + ):

321 E = config.n_routed_experts

322 topk = config.num_experts_per_tok

281 - elif config.architectures[0] in [

282 - "Qwen2MoeForCausalLM", "Qwen3MoeForCausalLM"

283 - ]:

323 + elif config.architectures[0] in ["Qwen2MoeForCausalLM", "Qwen3MoeForCausalLM"]:

324 E = config.num_experts

325 topk = config.num_experts_per_tok

326

@@ -299,8 +339,24 @@ def main(args: argparse.Namespace):

339

340 if args.batch_size is None:

341 batch_sizes = [

302 - 1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 256, 512, 1024, 1536,

303 - 2048, 3072, 4096

342 + 1,

343 + 2,

344 + 4,

345 + 8,

346 + 16,

347 + 24,

348 + 32,

349 + 48,

350 + 64,

351 + 96,

352 + 128,

353 + 256,

354 + 512,

355 + 1024,

356 + 1536,

357 + 2048,

358 + 3072,

359 + 4096,

360 ]

361 else:

362 batch_sizes = [args.batch_size]

@@ -321,9 +377,21 @@ def main(args: argparse.Namespace):

377 return ray.get(outputs)

378

379 outputs = _distribute(

324 - "benchmark", [(batch_size, E, hidden_size, topk, dtype, use_fp8_w8a8,

325 - use_int8_w8a16, use_customized_permute)

326 - for batch_size in batch_sizes])

380 + "benchmark",

381 + [

382 + (

383 + batch_size,

384 + E,

385 + hidden_size,

386 + topk,

387 + dtype,

388 + use_fp8_w8a8,

389 + use_int8_w8a16,

390 + use_customized_permute,

391 + )

392 + for batch_size in batch_sizes

393 + ],

394 + )

395

396 for batch_size, (permute, unpermute) in zip(batch_sizes, outputs):

397 print(f"Batch size: {batch_size}")

@@ -333,13 +401,12 @@ def main(args: argparse.Namespace):

401

402 if __name__ == "__main__":

403 parser = FlexibleArgumentParser()

336 - parser.add_argument("--model",

337 - type=str,

338 - default="mistralai/Mixtral-8x7B-Instruct-v0.1")

339 - parser.add_argument("--dtype",

340 - type=str,

341 - choices=["auto", "fp8_w8a8", "int8_w8a16"],

342 - default="auto")

404 + parser.add_argument(

405 + "--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1"

406 + )

407 + parser.add_argument(

408 + "--dtype", type=str, choices=["auto", "fp8_w8a8", "int8_w8a16"], default="auto"

409 + )

410 parser.add_argument("--use-customized-permute", action="store_true")

411 parser.add_argument("--seed", type=int, default=0)

412 parser.add_argument("--batch-size", type=int, required=False)

benchmarks/kernels/benchmark_paged_attention.py

+45 -52

@@ -9,8 +9,11 @@ import torch

9 from vllm import _custom_ops as ops

10 from vllm.logger import init_logger

11 from vllm.platforms import current_platform

12 -from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, FlexibleArgumentParser,

13 - create_kv_caches_with_random)

12 +from vllm.utils import (

13 + STR_DTYPE_TO_TORCH_DTYPE,

14 + FlexibleArgumentParser,

15 + create_kv_caches_with_random,

16 +)

18 logger = init_logger(__name__)

@@ -38,19 +41,15 @@ def main(

41 current_platform.seed_everything(seed)

43 scale = float(1.0 / (head_size**0.5))

41 - query = torch.empty(num_seqs,

42 - num_query_heads,

43 - head_size,

44 - dtype=dtype,

45 - device=device)

44 + query = torch.empty(

45 + num_seqs, num_query_heads, head_size, dtype=dtype, device=device

46 + )

47 query.uniform_(-scale, scale)

49 assert num_query_heads % num_kv_heads == 0

50 alibi_slopes = None

51 if use_alibi:

51 - alibi_slopes = torch.randn(num_query_heads,

52 - dtype=torch.float,

53 - device=device)

52 + alibi_slopes = torch.randn(num_query_heads, dtype=torch.float, device=device)

54 seq_lens = [seq_len for _ in range(num_seqs)]

55 max_seq_len = max(seq_lens)

@@ -61,24 +60,23 @@ def main(

60 block_tables_lst: list[list[int]] = []

61 for _ in range(num_seqs):

62 block_table = [

64 - random.randint(0, NUM_BLOCKS - 1)

65 - for _ in range(max_num_blocks_per_seq)

63 + random.randint(0, NUM_BLOCKS - 1) for _ in range(max_num_blocks_per_seq)

64 ]

65 block_tables_lst.append(block_table)

69 - block_tables = torch.tensor(block_tables_lst,

70 - dtype=torch.int,

71 - device=device)

67 + block_tables = torch.tensor(block_tables_lst, dtype=torch.int, device=device)

69 # Create the KV cache.

74 - key_caches, value_caches = create_kv_caches_with_random(NUM_BLOCKS,

75 - block_size,

76 - 1,

77 - num_kv_heads,

78 - head_size,

79 - kv_cache_dtype,

80 - dtype,

81 - device=device)

70 + key_caches, value_caches = create_kv_caches_with_random(

71 + NUM_BLOCKS,

72 + block_size,

73 + 1,

74 + num_kv_heads,

75 + head_size,

76 + kv_cache_dtype,

77 + dtype,

78 + device=device,

79 + )

80 key_cache, value_cache = key_caches[0], value_caches[0]

82 # Prepare for the paged attention kernel.

@@ -86,11 +84,8 @@ def main(

84 if version == "v2":

85 if current_platform.is_rocm():

86 global PARTITION_SIZE

89 - if not args.custom_paged_attn:

90 - PARTITION_SIZE = 1024

91 - else:

92 - PARTITION_SIZE = PARTITION_SIZE_ROCM

93 - num_partitions = ((max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE)

87 + PARTITION_SIZE = 1024 if not args.custom_paged_attn else PARTITION_SIZE_ROCM

88 + num_partitions = (max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE

89 tmp_output = torch.empty(

90 size=(num_seqs, num_query_heads, num_partitions, head_size),

91 dtype=output.dtype,

@@ -110,9 +105,7 @@ def main(

105 start_time = time.perf_counter()

106

107 # Using default kv_scale

113 - k_scale = v_scale = torch.tensor(1.0,

114 - dtype=torch.float32,

115 - device=device)

108 + k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)

109

110 for _ in range(num_iters):

111 if version == "v1":

@@ -195,30 +188,29 @@ def main(

188 print(f"Kernel running time: {latency * 1000000:.3f} us")

189

190

198 -if __name__ == '__main__':

199 - logger.warning("This script benchmarks the paged attention kernel. "

200 - "By default this is no longer used in vLLM inference.")

191 +if __name__ == "__main__":

192 + logger.warning(

193 + "This script benchmarks the paged attention kernel. "

194 + "By default this is no longer used in vLLM inference."

195 + )

196

202 - parser = FlexibleArgumentParser(

203 - description="Benchmark the paged attention kernel.")

204 - parser.add_argument("--version",

205 - type=str,

206 - choices=["v1", "v2"],

207 - default="v2")

197 + parser = FlexibleArgumentParser(description="Benchmark the paged attention kernel.")

198 + parser.add_argument("--version", type=str, choices=["v1", "v2"], default="v2")

199 parser.add_argument("--batch-size", type=int, default=8)

200 parser.add_argument("--seq-len", type=int, default=4096)

201 parser.add_argument("--num-query-heads", type=int, default=64)

202 parser.add_argument("--num-kv-heads", type=int, default=8)

212 - parser.add_argument("--head-size",

213 - type=int,

214 - choices=[64, 80, 96, 112, 120, 128, 192, 256],

215 - default=128)

203 + parser.add_argument(

204 + "--head-size",

205 + type=int,

206 + choices=[64, 80, 96, 112, 120, 128, 192, 256],

207 + default=128,

208 + )

209 parser.add_argument("--block-size", type=int, choices=[16, 32], default=16)

210 parser.add_argument("--use-alibi", action="store_true")

218 - parser.add_argument("--dtype",

219 - type=str,

220 - choices=["half", "bfloat16", "float"],

221 - default="half")

211 + parser.add_argument(

212 + "--dtype", type=str, choices=["half", "bfloat16", "float"], default="half"

213 + )

214 parser.add_argument("--seed", type=int, default=0)

215 parser.add_argument("--profile", action="store_true")

216 parser.add_argument(

@@ -228,10 +220,11 @@ if __name__ == '__main__':

220 default="auto",

221 help="Data type for kv cache storage. If 'auto', will use model "

222 "data type. CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. "

231 - "ROCm (AMD GPU) supports fp8 (=fp8_e4m3)")

232 - parser.add_argument("--custom-paged-attn",

233 - action="store_true",

234 - help="Use custom paged attention")

223 + "ROCm (AMD GPU) supports fp8 (=fp8_e4m3)",

224 + )

225 + parser.add_argument(

226 + "--custom-paged-attn", action="store_true", help="Use custom paged attention"

227 + )

228 args = parser.parse_args()

229 print(args)

230

benchmarks/kernels/benchmark_quant.py

+38 -33

@@ -10,15 +10,17 @@ from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, FlexibleArgumentParser

12 @torch.inference_mode()

13 -def main(num_tokens: int,

14 - hidden_size: int,

15 - static_scale: bool,

16 - quant_dtype: torch.dtype,

17 - dtype: torch.dtype,

18 - seed: int = 0,

19 - do_profile: bool = False,

20 - num_warmup_iters: int = 5,

21 - num_iters: int = 100) -> None:

13 +def main(

14 + num_tokens: int,

15 + hidden_size: int,

16 + static_scale: bool,

17 + quant_dtype: torch.dtype,

18 + dtype: torch.dtype,

19 + seed: int = 0,

20 + do_profile: bool = False,

21 + num_warmup_iters: int = 5,

22 + num_iters: int = 100,

23 +) -> None:

24 current_platform.seed_everything(seed)

25 torch.set_default_device("cuda")

@@ -56,7 +58,7 @@ def main(num_tokens: int,

58 print(f"Kernel running time: {latency * 1000000:.3f} us")

59 -if __name__ == '__main__':

61 +if __name__ == "__main__":

63 def to_torch_dtype(dt):

64 if dt == "int8":

@@ -66,37 +68,40 @@ if __name__ == '__main__':

68 raise ValueError(f"Unsupported dtype: {dt}")

70 parser = FlexibleArgumentParser(

69 - description="Benchmark the quantization (fp8 or int8) kernel.")

71 + description="Benchmark the quantization (fp8 or int8) kernel."

72 + )

73 parser.add_argument("--num-tokens", type=int, default=4096)

74 parser.add_argument("--hidden-size", type=int, default=8192)

75 parser.add_argument("--static-scale", action="store_true")

73 - parser.add_argument("--quant-dtype",

74 - type=str,

75 - choices=["fp8", "int8"],

76 - default="int8")

77 - parser.add_argument("--dtype",

78 - type=str,

79 - choices=["half", "bfloat16", "float"],

80 - default="half")

76 + parser.add_argument(

77 + "--quant-dtype", type=str, choices=["fp8", "int8"], default="int8"

78 + )

79 + parser.add_argument(

80 + "--dtype", type=str, choices=["half", "bfloat16", "float"], default="half"

81 + )

83 parser.add_argument("--seed", type=int, default=0)

84 parser.add_argument("--profile", action="store_true")

85 parser.add_argument("--num-warmup-iters", type=int, default=5)

85 - parser.add_argument("--num-iters",

86 - type=int,

87 - default=100,

88 - help="Number of benchmark iterations. "

89 - "If --profile is set, this number is ignored")

86 + parser.add_argument(

87 + "--num-iters",

88 + type=int,

89 + default=100,

90 + help="Number of benchmark iterations. "

91 + "If --profile is set, this number is ignored",

92 + )

94 args = parser.parse_args()

95 print(args)

94 - main(num_tokens=args.num_tokens,

95 - hidden_size=args.hidden_size,

96 - static_scale=args.static_scale,

97 - quant_dtype=to_torch_dtype(args.quant_dtype),

98 - dtype=STR_DTYPE_TO_TORCH_DTYPE[args.dtype],

99 - seed=args.seed,

100 - do_profile=args.profile,

101 - num_warmup_iters=args.num_warmup_iters,

102 - num_iters=args.num_iters)

97 + main(

98 + num_tokens=args.num_tokens,

99 + hidden_size=args.hidden_size,

100 + static_scale=args.static_scale,

101 + quant_dtype=to_torch_dtype(args.quant_dtype),

102 + dtype=STR_DTYPE_TO_TORCH_DTYPE[args.dtype],

103 + seed=args.seed,

104 + do_profile=args.profile,

105 + num_warmup_iters=args.num_warmup_iters,

106 + num_iters=args.num_iters,

107 + )

benchmarks/kernels/benchmark_rmsnorm.py

+24 -33

@@ -12,7 +12,6 @@ from vllm.triton_utils import triton

14 class HuggingFaceRMSNorm(nn.Module):

15 -

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

16 super().__init__()

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

@@ -114,23 +113,19 @@ def rmsnorm_vllm(

113

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

115 dtype = torch.bfloat16

117 - x = torch.randn(batch_size,

118 - seq_len,

119 - hidden_size,

120 - dtype=dtype,

121 - device="cuda")

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(

126 - x.clone(), weight,

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

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

122 + )

123 output_flashinfer = rmsnorm_flashinfer(

129 - x.clone(), weight,

130 - residual.clone() if residual is not None else None)

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

125 + )

126 output_vllm = rmsnorm_vllm(

132 - x.clone(), weight,

133 - residual.clone() if residual is not None else None)

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]

@@ -141,9 +136,9 @@ def calculate_diff(batch_size, seq_len, hidden_size, use_residual=True):

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

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

138

144 - if torch.allclose(output_naive, output_flashinfer, atol=1e-2,

145 - rtol=1e-2) and torch.allclose(

146 - output_naive, output_vllm, atol=1e-2, rtol=1e-2):

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")

@@ -152,12 +147,10 @@ def calculate_diff(batch_size, seq_len, hidden_size, use_residual=True):

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]

155 -configs = list(

156 - itertools.product(head_num_range, batch_size_range, seq_length_range))

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

151

152

153 def get_benchmark(use_residual):

160 -

154 @triton.testing.perf_report(

155 triton.testing.Benchmark(

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

@@ -167,19 +160,15 @@ def get_benchmark(use_residual):

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

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

162 ylabel="us",

170 - plot_name=

171 - f"rmsnorm-perf-{'with' if use_residual else 'without'}-residual",

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

164 args={},

173 - ))

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

178 - x = torch.randn(batch_size,

179 - seq_len,

180 - hidden_size,

181 - dtype=dtype,

182 - device="cuda")

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

@@ -240,9 +229,9 @@ if __name__ == "__main__":

229 default=4096,

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

231 )

243 - parser.add_argument("--use-residual",

244 - action="store_true",

245 - help="Whether to use residual connection")

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,

@@ -253,10 +242,12 @@ if __name__ == "__main__":

242 args = parser.parse_args()

243

244 # Run correctness test

256 - calculate_diff(batch_size=args.batch_size,

257 - seq_len=args.seq_len,

258 - hidden_size=args.hidden_size,

259 - use_residual=args.use_residual)

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)

benchmarks/kernels/benchmark_rope.py

+46 -37

@@ -6,8 +6,7 @@ from typing import Optional

6 import nvtx

7 import torch

9 -from vllm.model_executor.layers.rotary_embedding import (RotaryEmbedding,

10 - get_rope)

9 +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope

10 from vllm.platforms import current_platform

11 from vllm.utils import FlexibleArgumentParser

@@ -32,40 +31,49 @@ def benchmark_rope_kernels_multi_lora(

31 # silulating serving 4 LoRAs

32 scaling_factors = [1, 2, 4, 8]

33 # batched RoPE can take multiple scaling factors

35 - batched_rope = get_rope(head_size, rotary_dim, max_position, base,

36 - is_neox_style, {

37 - "rope_type": "linear",

38 - "factor": tuple(scaling_factors)

39 - })

34 + batched_rope = get_rope(

35 + head_size,

36 + rotary_dim,

37 + max_position,

38 + base,

39 + is_neox_style,

40 + {"rope_type": "linear", "factor": tuple(scaling_factors)},

41 + )

42 # non-batched RoPE takes only one scaling factor, we create multiple

43 # instances to simulate the same behavior

44 non_batched_ropes: list[RotaryEmbedding] = []

45 for scaling_factor in scaling_factors:

46 non_batched_ropes.append(

45 - get_rope(head_size, rotary_dim, max_position, base, is_neox_style,

46 - {

47 - "rope_type": "linear",

48 - "factor": (scaling_factor, )

49 - }))

47 + get_rope(

48 + head_size,

49 + rotary_dim,

50 + max_position,

51 + base,

52 + is_neox_style,

53 + {"rope_type": "linear", "factor": (scaling_factor,)},

54 + )

55 + )

57 positions = torch.randint(0, max_position, (batch_size, seq_len))

52 - query = torch.randn(batch_size,

53 - seq_len,

54 - num_heads * head_size,

55 - dtype=dtype)

58 + query = torch.randn(batch_size, seq_len, num_heads * head_size, dtype=dtype)

59 key = torch.randn_like(query)

61 # create query offsets for batched RoPE, we concat multiple kv cache

62 # together and each query needs to find the right kv cache of its type

63 offset_map = torch.tensor(

64 list(

62 - accumulate([0] + [

63 - max_position * scaling_factor * 2

64 - for scaling_factor in scaling_factors[:-1]

65 - ])))

66 - query_types = torch.randint(0,

67 - len(scaling_factors), (batch_size, seq_len),

68 - device=device)

65 + accumulate(

66 + [0]

67 + + [

68 + max_position * scaling_factor * 2

69 + for scaling_factor in scaling_factors[:-1]

70 + ]

71 + )

72 + )

73 + )

74 + query_types = torch.randint(

75 + 0, len(scaling_factors), (batch_size, seq_len), device=device

76 + )

77 # map query types to offsets

78 query_offsets = offset_map[query_types]

79 # the kernel takes flattened offsets

@@ -86,27 +94,28 @@ def benchmark_rope_kernels_multi_lora(

94 torch.cuda.synchronize()

89 -if __name__ == '__main__':

97 +if __name__ == "__main__":

98 parser = FlexibleArgumentParser(

91 - description="Benchmark the rotary embedding kernels.")

99 + description="Benchmark the rotary embedding kernels."

100 + )

101 parser.add_argument("--is-neox-style", type=bool, default=True)

102 parser.add_argument("--batch-size", type=int, default=16)

103 parser.add_argument("--seq-len", type=int, default=512)

104 parser.add_argument("--num-heads", type=int, default=8)

96 - parser.add_argument("--head-size",

97 - type=int,

98 - choices=[64, 80, 96, 112, 120, 128, 192, 256],

99 - default=128)

105 + parser.add_argument(

106 + "--head-size",

107 + type=int,

108 + choices=[64, 80, 96, 112, 120, 128, 192, 256],

109 + default=128,

110 + )

111 parser.add_argument("--rotary-dim", type=int, choices=[16, 32], default=32)

101 - parser.add_argument("--dtype",

102 - type=str,

103 - choices=["bfloat16", "float"],

104 - default="float")

112 + parser.add_argument(

113 + "--dtype", type=str, choices=["bfloat16", "float"], default="float"

114 + )

115 parser.add_argument("--seed", type=int, default=0)

106 - parser.add_argument("--device",

107 - type=str,

108 - choices=["cuda:0", "cuda:1"],

109 - default="cuda:0")

116 + parser.add_argument(

117 + "--device", type=str, choices=["cuda:0", "cuda:1"], default="cuda:0"

118 + )

119 args = parser.parse_args()

120 print(args)

121

benchmarks/kernels/benchmark_w8a8_block_fp8.py

+53 -60

@@ -14,14 +14,16 @@ import tqdm

14 import triton

16 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

17 - _w8a8_block_fp8_matmul)

17 + _w8a8_block_fp8_matmul,

18 +)

19 from vllm.platforms import current_platform

20 from vllm.utils import FlexibleArgumentParser

22 mp.set_start_method("spawn", force=True)

23 -assert current_platform.is_cuda(

24 -), "Only support tune w8a8 block fp8 kernel on CUDA device."

24 +assert current_platform.is_cuda(), (

25 + "Only support tune w8a8 block fp8 kernel on CUDA device."

26 +)

28 DTYPE_MAP = {

29 "float32": torch.float32,

@@ -40,7 +42,7 @@ def w8a8_block_matmul(

42 config: dict[str, Any],

43 output_dtype: torch.dtype = torch.float16,

44 ) -> torch.Tensor:

43 - """This function performs matrix multiplication with

45 + """This function performs matrix multiplication with

46 block-wise quantization.

48 It takes two input tensors `A` and `B` with scales `As` and `Bs`.

@@ -51,7 +53,7 @@ def w8a8_block_matmul(

53 B: The input tensor, e.g., weight.

54 As: The per-token-group quantization scale for `A`.

55 Bs: The per-block quantization scale for `B`.

54 - block_size: The block size for per-block quantization.

56 + block_size: The block size for per-block quantization.

57 It should be 2-dim, e.g., [128, 128].

58 output_dytpe: The dtype of the returned tensor.

@@ -71,18 +73,18 @@ def w8a8_block_matmul(

73 assert triton.cdiv(N, block_n) == Bs.shape[0]

74 assert triton.cdiv(K, block_k) == Bs.shape[1]

74 - C_shape = A.shape[:-1] + (N, )

76 + C_shape = A.shape[:-1] + (N,)

77 C = A.new_empty(C_shape, dtype=output_dtype)

79 def grid(META):

78 - return (triton.cdiv(M, META["BLOCK_SIZE_M"]) *

79 - triton.cdiv(N, META["BLOCK_SIZE_N"]), )

80 + return (

81 + triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),

82 + )

84 if A.dtype == torch.float8_e4m3fn:

85 kernel = _w8a8_block_fp8_matmul

86 else:

84 - raise RuntimeError(

85 - "Currently, only support tune w8a8 block fp8 kernel.")

87 + raise RuntimeError("Currently, only support tune w8a8 block fp8 kernel.")

89 kernel[grid](

90 A,

@@ -119,14 +121,16 @@ def get_configs_compute_bound():

121 for block_n in [32, 64, 128, 256]:

122 for num_warps in [4, 8]:

123 for group_size in [1, 16, 32, 64]:

122 - configs.append({

123 - "BLOCK_SIZE_M": block_m,

124 - "BLOCK_SIZE_N": block_n,

125 - "BLOCK_SIZE_K": block_k,

126 - "GROUP_SIZE_M": group_size,

127 - "num_warps": num_warps,

128 - "num_stages": num_stages,

129 - })

124 + configs.append(

125 + {

126 + "BLOCK_SIZE_M": block_m,

127 + "BLOCK_SIZE_N": block_n,

128 + "BLOCK_SIZE_K": block_k,

129 + "GROUP_SIZE_M": group_size,

130 + "num_warps": num_warps,

131 + "num_stages": num_stages,

132 + }

133 + )

134 return configs

135

136

@@ -165,15 +169,9 @@ def get_weight_shapes(tp_size):

169 return weight_shapes

170

171

168 -def benchmark_config(A,

169 - B,

170 - As,

171 - Bs,

172 - block_size,

173 - config,

174 - out_dtype=torch.float16,

175 - num_iters=10):

176 -

172 +def benchmark_config(

173 + A, B, As, Bs, block_size, config, out_dtype=torch.float16, num_iters=10

174 +):

175 def run():

176 w8a8_block_matmul(A, B, As, Bs, block_size, config, out_dtype)

177

@@ -206,26 +204,26 @@ def tune(M, N, K, block_size, out_dtype, search_space, input_type):

204 fp8_max, fp8_min = fp8_info.max, fp8_info.min

205

206 A_fp32 = (

209 - (torch.rand(M, K, dtype=torch.float32, device="cuda") - 0.5) * 2 *

210 - fp8_max)

207 + (torch.rand(M, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max

208 + )

209 A = A_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)

210

211 B_fp32 = (

214 - (torch.rand(N, K, dtype=torch.float32, device="cuda") - 0.5) * 2 *

215 - fp8_max)

212 + (torch.rand(N, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max

213 + )

214 B = B_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)

215 else:

218 - raise RuntimeError(

219 - "Currently, only support tune w8a8 block fp8 kernel.")

216 + raise RuntimeError("Currently, only support tune w8a8 block fp8 kernel.")

217

218 block_n, block_k = block_size[0], block_size[1]

219 n_tiles = (N + block_n - 1) // block_n

220 k_tiles = (K + block_k - 1) // block_k

221

225 - As = torch.rand(M, k_tiles, dtype=torch.float32,

226 - device="cuda") * factor_for_scale

227 - Bs = (torch.rand(n_tiles, k_tiles, dtype=torch.float32, device="cuda") *

228 - factor_for_scale)

222 + As = torch.rand(M, k_tiles, dtype=torch.float32, device="cuda") * factor_for_scale

223 + Bs = (

224 + torch.rand(n_tiles, k_tiles, dtype=torch.float32, device="cuda")

225 + * factor_for_scale

226 + )

227

228 best_config = None

229 best_time = float("inf")

@@ -267,7 +265,8 @@ def save_configs(

265 device_name = current_platform.get_device_name().replace(" ", "_")

266 json_file_name = (

267 f"N={N},K={K},device_name={device_name},dtype={input_type}_w8a8,"

270 - f"block_shape=[{block_n},{block_k}].json")

268 + f"block_shape=[{block_n},{block_k}].json"

269 + )

270

271 config_file_path = os.path.join(save_path, json_file_name)

272 print(f"Writing best config to {config_file_path}...")

@@ -295,8 +294,7 @@ def tune_on_gpu(args_dict):

294

295 search_space = get_configs_compute_bound()

296 search_space = [

298 - config for config in search_space

299 - if block_k % config["BLOCK_SIZE_K"] == 0

297 + config for config in search_space if block_k % config["BLOCK_SIZE_K"] == 0

298 ]

299

300 start = time.time()

@@ -312,15 +310,11 @@ def tune_on_gpu(args_dict):

310 out_dtype,

311 search_space,

312 input_type,

315 - ) for batch_size in tqdm(batch_sizes,

316 - desc=f"GPU {gpu_id} - Batch sizes")

313 + )

314 + for batch_size in tqdm(batch_sizes, desc=f"GPU {gpu_id} - Batch sizes")

315 ]

318 - best_configs = {

319 - M: config

320 - for M, config in zip(batch_sizes, benchmark_results)

321 - }

322 - save_configs(N, K, block_n, block_k, best_configs, save_path,

323 - input_type)

316 + best_configs = {M: config for M, config in zip(batch_sizes, benchmark_results)}

317 + save_configs(N, K, block_n, block_k, best_configs, save_path, input_type)

318

319 end = time.time()

320 print(f"Tuning on GPU {gpu_id} took {end - start:.2f} seconds")

@@ -376,13 +370,14 @@ def main(args):

370

371 process_args = []

372 for gpu_id in range(num_gpus):

379 - process_args.append({

380 - "gpu_id": gpu_id,

381 - "batch_sizes": batches_per_gpu[gpu_id],

382 - "weight_shapes":

383 - weight_shapes, # Each GPU processes all weight shapes

384 - "args": args,

385 - })

373 + process_args.append(

374 + {

375 + "gpu_id": gpu_id,

376 + "batch_sizes": batches_per_gpu[gpu_id],

377 + "weight_shapes": weight_shapes, # Each GPU processes all weight shapes

378 + "args": args,

379 + }

380 + )

381

382 ctx = mp.get_context("spawn")

383 with ctx.Pool(num_gpus) as pool:

@@ -398,13 +393,11 @@ Tune triton w8a8 block fp8 for DeepSeek-V3/DeepSeek-R1:

393 python3 benchmark_w8a8_block_fp8.py --tp-size 8 --input-type fp8

394 Then copy to model_executor/layers/quantization/utils/configs

395 """,

401 - formatter_class=argparse.RawTextHelpFormatter)

396 + formatter_class=argparse.RawTextHelpFormatter,

397 + )

398

399 parser.add_argument("--tp-size", "-tp", type=int, default=8)

404 - parser.add_argument("--input-type",

405 - type=str,

406 - choices=["fp8"],

407 - default="fp8")

400 + parser.add_argument("--input-type", type=str, choices=["fp8"], default="fp8")

401 parser.add_argument(

402 "--out-dtype",

403 type=str,

benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py

+3 -1

@@ -11,7 +11,9 @@ from deep_gemm import calc_diff, ceil_div, get_col_major_tma_aligned_tensor

11 # Import vLLM functions

12 from vllm import _custom_ops as ops

13 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

14 - per_token_group_quant_fp8, w8a8_block_fp8_matmul)

14 + per_token_group_quant_fp8,

15 + w8a8_block_fp8_matmul,

16 +)

17 from vllm.triton_utils import triton

benchmarks/kernels/graph_machete_bench.py

+16 -17

@@ -14,13 +14,14 @@ from vllm.utils import FlexibleArgumentParser

15 if __name__ == "__main__":

16 parser = FlexibleArgumentParser(

17 - description='Benchmark the latency of processing a single batch of '

18 - 'requests till completion.')

19 - parser.add_argument('filename', type=str)

17 + description="Benchmark the latency of processing a single batch of "

18 + "requests till completion."

19 + )

20 + parser.add_argument("filename", type=str)

22 args = parser.parse_args()

23 - with open(args.filename, 'rb') as f:

24 + with open(args.filename, "rb") as f:

25 data = pickle.load(f)

26 raw_results: list[TMeasurement] = data["results"]

@@ -38,11 +39,7 @@ if __name__ == "__main__":

39 raise Exception("MKN not found")

41 kernel = v.task_spec.description

41 - results[KN].append({

42 - "kernel": kernel,

43 - "batch_size": M,

44 - "median": v.median

45 - })

42 + results[KN].append({"kernel": kernel, "batch_size": M, "median": v.median})

44 rows = int(math.ceil(len(results) / 2))

45 fig, axs = plt.subplots(rows, 2, figsize=(12, 5 * rows))

@@ -50,14 +47,16 @@ if __name__ == "__main__":

47 for axs_idx, (shape, data) in enumerate(results.items()):

48 plt.sca(axs[axs_idx])

49 df = pd.DataFrame(data)

53 - sns.lineplot(data=df,

54 - x="batch_size",

55 - y="median",

56 - hue="kernel",

57 - style="kernel",

58 - markers=True,

59 - dashes=False,

60 - palette="Dark2")

50 + sns.lineplot(

51 + data=df,

52 + x="batch_size",

53 + y="median",

54 + hue="kernel",

55 + style="kernel",

56 + markers=True,

57 + dashes=False,

58 + palette="Dark2",

59 + )

60 plt.title(f"Shape: {shape}")

61 plt.ylabel("time (median, s)")

62 plt.tight_layout()

benchmarks/kernels/utils.py

+26 -26

@@ -23,6 +23,7 @@ class ArgPool:

23 For every invocation during a benchmarking run, it will choose a

24 different value from the list.

25 """

26 +

27 values: Iterable[Any]

29 def __getitem__(self, index):

@@ -30,9 +31,7 @@ class ArgPool:

33 class Bench:

33 -

34 class ArgsIterator:

35 -

35 def __init__(self, args_list, kwargs_list):

36 assert len(args_list) == len(kwargs_list)

37 self.args_list = args_list

@@ -53,10 +52,16 @@ class Bench:

52 def n_args(self):

53 return self.n

56 - def __init__(self, cuda_graph_params: Optional[CudaGraphBenchParams],

57 - label: str, sub_label: str, description: str, fn: Callable,

58 - *args, **kwargs):

59 -

55 + def __init__(

56 + self,

57 + cuda_graph_params: Optional[CudaGraphBenchParams],

58 + label: str,

59 + sub_label: str,

60 + description: str,

61 + fn: Callable,

62 + *args,

63 + **kwargs,

64 + ):

65 self.cuda_graph_params = cuda_graph_params

66 self.use_cuda_graph = self.cuda_graph_params is not None

67 self.label = label

@@ -67,10 +72,8 @@ class Bench:

72 # Process args

73 self._args = args

74 self._kwargs = kwargs

70 - self.args_list, self.kwargs_list = self.collapse_argpool(

71 - *args, **kwargs)

72 - self.args_iterator = self.ArgsIterator(self.args_list,

73 - self.kwargs_list)

75 + self.args_list, self.kwargs_list = self.collapse_argpool(*args, **kwargs)

76 + self.args_iterator = self.ArgsIterator(self.args_list, self.kwargs_list)

78 # Cudagraph runner

79 self.g = None

@@ -100,16 +103,13 @@ class Bench:

103

104 for i in range(argpool_size):

105 # collapse args; Just pick the ith value

103 - args_list[i] = tuple([

104 - arg[i] if isinstance(arg, ArgPool) else arg

105 - for arg in args_list[i]

106 - ])

106 + args_list[i] = tuple(

107 + [arg[i] if isinstance(arg, ArgPool) else arg for arg in args_list[i]]

108 + )

109

110 # collapse kwargs

111 kwargs_i = kwargs_list[i]

110 - arg_pool_keys = [

111 - k for k, v in kwargs_i.items() if isinstance(v, ArgPool)

112 - ]

112 + arg_pool_keys = [k for k, v in kwargs_i.items() if isinstance(v, ArgPool)]

113 for k in arg_pool_keys:

114 # again just pick the ith value

115 kwargs_i[k] = kwargs_i[k][i]

@@ -142,7 +142,7 @@ class Bench:

142

143 def run_cudagrah(self) -> TMeasurement:

144 assert self.use_cuda_graph

145 - globals = {'g': self.g}

145 + globals = {"g": self.g}

146

147 return TBenchmark.Timer(

148 stmt="g.replay()",

@@ -162,15 +162,15 @@ class Bench:

162

163 has_arg_pool = self.args_iterator.n_args > 1

164 if has_arg_pool:

165 - setup = '''

165 + setup = """

166 args_iterator.reset()

167 args_it = args_iterator.__next__()

168 - '''

169 - stmt = '''

168 + """

169 + stmt = """

170 args, kwargs = next(args_it)

171 fn(*args, **kwargs)

172 - '''

173 - globals = {'fn': self.fn, 'args_iterator': self.args_iterator}

172 + """

173 + globals = {"fn": self.fn, "args_iterator": self.args_iterator}

174 else:

175 # no arg pool. Just use the args and kwargs directly

176 self.args_iterator.reset()

@@ -178,10 +178,10 @@ class Bench:

178 args, kwargs = next(args_it)

179

180 setup = ""

181 - stmt = '''

181 + stmt = """

182 fn(*args, **kwargs)

183 - '''

184 - globals = {'fn': self.fn, 'args': args, 'kwargs': kwargs}

183 + """

184 + globals = {"fn": self.fn, "args": args, "kwargs": kwargs}

185

186 return TBenchmark.Timer(

187 stmt=stmt,

benchmarks/overheads/benchmark_hashing.py

+19 -17

@@ -7,9 +7,8 @@ from vllm import LLM, SamplingParams

7 from vllm.utils import FlexibleArgumentParser

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

10 -LONG_PROMPT = ["You are an expert in large language models, aren't you?"

11 - ] * 1000

12 -LONG_PROMPT = ' '.join(LONG_PROMPT)

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

11 +LONG_PROMPT = " ".join(LONG_PROMPT)

14 def main(args):

@@ -30,32 +29,35 @@ def main(args):

30 print("------start generating------")

31 for i in range(3):

33 - profiler.runctx('llm.generate(LONG_PROMPT, sampling_params)',

34 - globals(), locals())

32 + profiler.runctx(

33 + "llm.generate(LONG_PROMPT, sampling_params)", globals(), locals()

34 + )

36 # analyze the runtime of hashing function

37 stats = pstats.Stats(profiler)

38 - stats.sort_stats('cumulative')

38 + stats.sort_stats("cumulative")

39 total_time = 0

40 total_calls = 0

41 for func in stats.stats:

42 - if 'hash_of_block' in func[2]:

42 + if "hash_of_block" in func[2]:

43 total_time = stats.stats[func][3]

44 total_calls = stats.stats[func][0]

45 percentage = (total_time / stats.total_tt) * 100

46 - print(f"Hashing took {total_time:.2f} seconds,"

47 - f"{percentage:.2f}% of the total runtime.")

46 + print(

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

48 + )

51 if __name__ == "__main__":

52 parser = FlexibleArgumentParser(

52 - description='Benchmark the performance of hashing function in'

53 - 'automatic prefix caching.')

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

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

56 - parser.add_argument('--output-len', type=int, default=10)

57 - parser.add_argument('--enable-prefix-caching',

58 - action='store_true',

59 - help='enable prefix caching')

53 + description="Benchmark the performance of hashing function in"

54 + "automatic prefix caching."

55 + )

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

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

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

59 + parser.add_argument(

60 + "--enable-prefix-caching", action="store_true", help="enable prefix caching"

61 + )

62 args = parser.parse_args()

63 main(args)

benchmarks/pyproject.toml new

+54

@@ -0,0 +1,54 @@

1 +# This local pyproject file is part of the migration from yapf to ruff format.

2 +# It uses the same core rules as the main pyproject.toml file, but with the

3 +# following differences:

4 +# - ruff line length is overridden to 88

5 +# - deprecated typing ignores (UP006, UP035) have been removed

6 +

7 +[tool.ruff]

8 +line-length = 88

9 +exclude = [

10 + # External file, leaving license intact

11 + "examples/other/fp8/quantizer/quantize.py",

12 + "vllm/vllm_flash_attn/flash_attn_interface.pyi"

13 +]

14 +

15 +[tool.ruff.lint.per-file-ignores]

16 +"vllm/third_party/**" = ["ALL"]

17 +"vllm/version.py" = ["F401"]

18 +"vllm/_version.py" = ["ALL"]

19 +

20 +[tool.ruff.lint]

21 +select = [

22 + # pycodestyle

23 + "E",

24 + # Pyflakes

25 + "F",

26 + # pyupgrade

27 + "UP",

28 + # flake8-bugbear

29 + "B",

30 + # flake8-simplify

31 + "SIM",

32 + # isort

33 + "I",

34 + # flake8-logging-format

35 + "G",

36 +]

37 +ignore = [

38 + # star imports

39 + "F405", "F403",

40 + # lambda expression assignment

41 + "E731",

42 + # Loop control variable not used within loop body

43 + "B007",

44 + # f-string format

45 + "UP032",

46 + # Can remove once 3.10+ is the minimum Python version

47 + "UP007",

48 +]

49 +

50 +[tool.ruff.lint.isort]

51 +known-first-party = ["vllm"]

52 +

53 +[tool.ruff.format]

54 +docstring-code-format = true

\ No newline at end of file

pyproject.toml

@@ -54,6 +54,7 @@ include = ["vllm*"]

54 [tool.yapfignore]

55 ignore_patterns = [

56 ".buildkite/**",

57 + "benchmarks/**",

58 "build/**",

59 ]

@@ -155,6 +156,10 @@ ignore-words-list = "dout, te, indicies, subtile, ElementE"

156 skip = "tests/models/fixtures/*,tests/prompts/*,benchmarks/sonnet.txt,tests/lora/data/*,build/*,vllm/third_party/*"

157

158 [tool.isort]

159 +skip_glob = [

160 + ".buildkite/*",

161 + "benchmarks/*",

162 +]

163 use_parentheses = true

164 skip_gitignore = true

165