chore: vendor sglang v0.5.10 snapshot
This commit is contained in:
425
third_party/sglang/benchmark/bench_rope/benchmark_rope_index.py
vendored
Normal file
425
third_party/sglang/benchmark/bench_rope/benchmark_rope_index.py
vendored
Normal file
@@ -0,0 +1,425 @@
|
||||
# This script benchmarks MRotaryEmbedding.get_rope_index_glm4v (GLM4V mrope index builder).
|
||||
# It generates synthetic multimodal input_ids + attention_mask (+ optional image/video grids),
|
||||
# runs benchmarks.
|
||||
#
|
||||
# == Usage Examples ==
|
||||
#
|
||||
# python3 benchmark_rope_index.py --device cuda --num-tokens 1024 2048 --benchmark-iter 200
|
||||
|
||||
import argparse
|
||||
import math
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.rotary_embedding import MRotaryEmbedding
|
||||
|
||||
|
||||
# -----------------------------
|
||||
# Minimal config objects
|
||||
# -----------------------------
|
||||
@dataclass
|
||||
class DummyVisionConfig:
|
||||
spatial_merge_size: int = 2
|
||||
|
||||
|
||||
@dataclass
|
||||
class DummyHFConfig:
|
||||
image_token_id: int = 32000
|
||||
video_start_token_id: int = 32001
|
||||
video_end_token_id: int = 32002
|
||||
vision_config: DummyVisionConfig = field(
|
||||
default_factory=lambda: DummyVisionConfig(spatial_merge_size=2)
|
||||
)
|
||||
|
||||
|
||||
# -----------------------------
|
||||
# Helpers
|
||||
# -----------------------------
|
||||
def calculate_stats(times: list[float]) -> dict[str, float]:
|
||||
"""Calculate statistics from a list of times."""
|
||||
times_array = np.array(times, dtype=np.float64)
|
||||
return {
|
||||
"mean": float(np.mean(times_array)),
|
||||
"median": float(np.median(times_array)),
|
||||
"p99": float(np.percentile(times_array, 99)),
|
||||
"min": float(np.min(times_array)),
|
||||
"max": float(np.max(times_array)),
|
||||
}
|
||||
|
||||
|
||||
def _sync(device: torch.device):
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize()
|
||||
|
||||
|
||||
def _approx_hw(patches: int, merge: int) -> tuple[int, int]:
|
||||
# want (h/merge)*(w/merge) ~= patches
|
||||
gh = int(math.sqrt(max(1, patches)))
|
||||
gw = max(1, patches // max(1, gh))
|
||||
return gh * merge, gw * merge
|
||||
|
||||
|
||||
def generate_test_data(
|
||||
num_tokens: int,
|
||||
batch_size: int,
|
||||
hf_config: DummyHFConfig,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
pad_ratio: float,
|
||||
num_images_per_sample: int,
|
||||
image_patch_tokens: int,
|
||||
num_videos_per_sample: int,
|
||||
video_patch_tokens: int,
|
||||
seed: int,
|
||||
):
|
||||
"""
|
||||
Generate synthetic (input_ids, attention_mask, image_grid_thw, video_grid_thw).
|
||||
|
||||
NOTE:
|
||||
- image_grid_thw / video_grid_thw are global lists across the entire batch in encounter order,
|
||||
matching the function's image_index/video_index behavior.
|
||||
- image patches are represented by repeated image_token_id.
|
||||
- video patches are represented by image_token_id wrapped with start/end tokens.
|
||||
"""
|
||||
torch.manual_seed(seed)
|
||||
|
||||
forbidden = {
|
||||
0,
|
||||
hf_config.image_token_id,
|
||||
hf_config.video_start_token_id,
|
||||
hf_config.video_end_token_id,
|
||||
}
|
||||
vocab_size = 50000
|
||||
|
||||
def rand_text(n: int) -> torch.Tensor:
|
||||
# generate random ids not in forbidden
|
||||
out = torch.randint(1, vocab_size, (n,), device=device, dtype=torch.long)
|
||||
# fix forbidden by +1 until ok (cheap, deterministic enough for benchmark data)
|
||||
for bad in forbidden:
|
||||
out = torch.where(out == bad, out + 1, out)
|
||||
return out
|
||||
|
||||
image_grids: list[list[int]] = []
|
||||
video_grids: list[list[int]] = []
|
||||
|
||||
input_ids = torch.zeros((batch_size, num_tokens), device=device, dtype=torch.long)
|
||||
attention_mask = torch.zeros(
|
||||
(batch_size, num_tokens), device=device, dtype=torch.long
|
||||
)
|
||||
|
||||
eff_len = int(round(num_tokens * (1.0 - pad_ratio)))
|
||||
eff_len = max(1, min(num_tokens, eff_len))
|
||||
|
||||
min_needed = 1
|
||||
min_needed += num_images_per_sample * image_patch_tokens
|
||||
min_needed += num_videos_per_sample * (2 + video_patch_tokens)
|
||||
if eff_len < min_needed:
|
||||
num_images_per_sample = 0
|
||||
num_videos_per_sample = 0
|
||||
|
||||
for b in range(batch_size):
|
||||
blocks: list[torch.Tensor] = []
|
||||
|
||||
reserved = (
|
||||
num_images_per_sample * image_patch_tokens
|
||||
+ num_videos_per_sample * (2 + video_patch_tokens)
|
||||
)
|
||||
reserved = min(reserved, max(0, eff_len - 1))
|
||||
text_budget = max(1, eff_len - reserved)
|
||||
|
||||
n_text_chunks = num_images_per_sample + num_videos_per_sample + 1
|
||||
base = text_budget // n_text_chunks
|
||||
rem = text_budget % n_text_chunks
|
||||
text_chunks = [base + (1 if i < rem else 0) for i in range(n_text_chunks)]
|
||||
|
||||
tci = 0
|
||||
for _ in range(num_images_per_sample):
|
||||
blocks.append(rand_text(text_chunks[tci]))
|
||||
tci += 1
|
||||
blocks.append(
|
||||
torch.full(
|
||||
(image_patch_tokens,),
|
||||
hf_config.image_token_id,
|
||||
device=device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
)
|
||||
|
||||
h, w = _approx_hw(
|
||||
image_patch_tokens, hf_config.vision_config.spatial_merge_size
|
||||
)
|
||||
image_grids.append([1, h, w])
|
||||
|
||||
for _ in range(num_videos_per_sample):
|
||||
blocks.append(rand_text(text_chunks[tci]))
|
||||
tci += 1
|
||||
blocks.append(
|
||||
torch.tensor(
|
||||
[hf_config.video_start_token_id], device=device, dtype=torch.long
|
||||
)
|
||||
)
|
||||
blocks.append(
|
||||
torch.full(
|
||||
(video_patch_tokens,),
|
||||
hf_config.image_token_id,
|
||||
device=device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
)
|
||||
blocks.append(
|
||||
torch.tensor(
|
||||
[hf_config.video_end_token_id], device=device, dtype=torch.long
|
||||
)
|
||||
)
|
||||
|
||||
h, w = _approx_hw(
|
||||
video_patch_tokens, hf_config.vision_config.spatial_merge_size
|
||||
)
|
||||
# first field = group count used by code; set to 1
|
||||
video_grids.append([1, h, w])
|
||||
|
||||
blocks.append(rand_text(text_chunks[tci]))
|
||||
|
||||
tokens = torch.cat(blocks, dim=0)[:eff_len]
|
||||
pad = torch.zeros(
|
||||
(num_tokens - tokens.numel(),), device=device, dtype=torch.long
|
||||
)
|
||||
ids = torch.cat([tokens, pad], dim=0)
|
||||
|
||||
mask = torch.cat(
|
||||
[
|
||||
torch.ones((tokens.numel(),), device=device, dtype=torch.long),
|
||||
torch.zeros(
|
||||
(num_tokens - tokens.numel(),), device=device, dtype=torch.long
|
||||
),
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
|
||||
input_ids[b] = ids
|
||||
attention_mask[b] = mask
|
||||
|
||||
image_grid_thw = (
|
||||
torch.tensor(image_grids, device=device, dtype=torch.long)
|
||||
if len(image_grids)
|
||||
else None
|
||||
)
|
||||
video_grid_thw = (
|
||||
torch.tensor(video_grids, device=device, dtype=torch.long)
|
||||
if len(video_grids)
|
||||
else None
|
||||
)
|
||||
return (
|
||||
input_ids.to(dtype=torch.long),
|
||||
attention_mask.to(dtype=torch.long),
|
||||
image_grid_thw,
|
||||
video_grid_thw,
|
||||
)
|
||||
|
||||
|
||||
def benchmark_rope_index(
|
||||
model_name: str,
|
||||
tp_size: int,
|
||||
num_tokens: int,
|
||||
batch_size: int,
|
||||
pad_ratio: float,
|
||||
spatial_merge_size: int,
|
||||
num_images: int,
|
||||
image_patch_tokens: int,
|
||||
num_videos: int,
|
||||
video_patch_tokens: int,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
warmup_iter: int,
|
||||
benchmark_iter: int,
|
||||
device: torch.device,
|
||||
):
|
||||
torch.manual_seed(seed)
|
||||
hf_config = DummyHFConfig(
|
||||
image_token_id=32000,
|
||||
video_start_token_id=32001,
|
||||
video_end_token_id=32002,
|
||||
vision_config=DummyVisionConfig(spatial_merge_size=spatial_merge_size),
|
||||
)
|
||||
|
||||
print(80 * "=")
|
||||
print(
|
||||
f"Evaluating: {model_name} tp_size={tp_size} "
|
||||
f"num_tokens={num_tokens} batch={batch_size} pad_ratio={pad_ratio} "
|
||||
f"images/sample={num_images} image_patch_tokens={image_patch_tokens} "
|
||||
f"videos/sample={num_videos} video_patch_tokens={video_patch_tokens} "
|
||||
f"dtype={dtype} device={device}"
|
||||
)
|
||||
|
||||
input_ids, attention_mask, image_grid_thw, video_grid_thw = generate_test_data(
|
||||
num_tokens=num_tokens,
|
||||
batch_size=batch_size,
|
||||
hf_config=hf_config,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
pad_ratio=pad_ratio,
|
||||
num_images_per_sample=num_images,
|
||||
image_patch_tokens=image_patch_tokens,
|
||||
num_videos_per_sample=num_videos,
|
||||
video_patch_tokens=video_patch_tokens,
|
||||
seed=seed,
|
||||
)
|
||||
|
||||
# Smoke test
|
||||
has_mm = (image_grid_thw is not None) or (video_grid_thw is not None)
|
||||
if has_mm:
|
||||
pos, delta = MRotaryEmbedding.get_rope_index_glm4v(
|
||||
input_ids=input_ids,
|
||||
hf_config=hf_config,
|
||||
image_grid_thw=image_grid_thw,
|
||||
video_grid_thw=video_grid_thw,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
assert pos.shape == (3, batch_size, num_tokens)
|
||||
assert delta.shape == (batch_size, 1)
|
||||
|
||||
# Warm up
|
||||
for _ in range(warmup_iter):
|
||||
if has_mm:
|
||||
MRotaryEmbedding.get_rope_index_glm4v(
|
||||
input_ids=input_ids,
|
||||
hf_config=hf_config,
|
||||
image_grid_thw=image_grid_thw,
|
||||
video_grid_thw=video_grid_thw,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
MRotaryEmbedding.get_rope_index_glm4v(
|
||||
input_ids=input_ids,
|
||||
hf_config=hf_config,
|
||||
image_grid_thw=None,
|
||||
video_grid_thw=None,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
|
||||
_sync(device)
|
||||
|
||||
# Time multimodal branch
|
||||
multimodal_times = []
|
||||
for _ in range(benchmark_iter):
|
||||
_sync(device)
|
||||
start = time.time()
|
||||
MRotaryEmbedding.get_rope_index_glm4v(
|
||||
input_ids=input_ids,
|
||||
hf_config=hf_config,
|
||||
image_grid_thw=image_grid_thw,
|
||||
video_grid_thw=video_grid_thw,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
_sync(device)
|
||||
multimodal_times.append(time.time() - start)
|
||||
|
||||
# Time fallback branch
|
||||
fallback_times = []
|
||||
for _ in range(benchmark_iter):
|
||||
_sync(device)
|
||||
start = time.time()
|
||||
MRotaryEmbedding.get_rope_index_glm4v(
|
||||
input_ids=input_ids,
|
||||
hf_config=hf_config,
|
||||
image_grid_thw=None,
|
||||
video_grid_thw=None,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
_sync(device)
|
||||
fallback_times.append(time.time() - start)
|
||||
|
||||
multimodal_stats = calculate_stats(multimodal_times)
|
||||
fallback_stats = calculate_stats(fallback_times)
|
||||
|
||||
print(f"\nPerformance for config (B={batch_size}, T={num_tokens}):")
|
||||
print(
|
||||
f"Multimodal: mean={multimodal_stats['mean']:.8f}s, "
|
||||
f"median={multimodal_stats['median']:.8f}s, "
|
||||
f"p99={multimodal_stats['p99']:.8f}s"
|
||||
)
|
||||
print(
|
||||
f"Fallback: mean={fallback_stats['mean']:.8f}s, "
|
||||
f"median={fallback_stats['median']:.8f}s, "
|
||||
f"p99={fallback_stats['p99']:.8f}s"
|
||||
)
|
||||
|
||||
if has_mm:
|
||||
speedup = (
|
||||
multimodal_stats["mean"] / fallback_stats["mean"]
|
||||
if fallback_stats["mean"] > 0
|
||||
else float("inf")
|
||||
)
|
||||
print(f"Fallback Speedup over Multimodal: {speedup:.8f}x")
|
||||
else:
|
||||
speedup = float("nan")
|
||||
print(
|
||||
"[INFO] num_tokens too small for multimodal segments; skip multimodal benchmark."
|
||||
)
|
||||
|
||||
print(f"Fallback Speedup over Multimodal: {speedup:.8f}x")
|
||||
|
||||
return multimodal_stats, fallback_stats, speedup
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Benchmark GLM4V get_rope_index_glm4v."
|
||||
)
|
||||
parser.add_argument("--model-name", type=str, default="GLM4V")
|
||||
parser.add_argument("--tp-size", type=int, default=1)
|
||||
parser.add_argument(
|
||||
"--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu"
|
||||
)
|
||||
parser.add_argument("--warmup-iter", type=int, default=10)
|
||||
parser.add_argument("--benchmark-iter", type=int, default=100)
|
||||
parser.add_argument("--dtype", type=str, choices=["int64"], default="int64")
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
|
||||
# token length sweep
|
||||
parser.add_argument("--num-tokens", type=int, nargs="+", required=False)
|
||||
|
||||
# data shape knobs
|
||||
parser.add_argument("--batch-size", type=int, default=1)
|
||||
parser.add_argument("--pad-ratio", type=float, default=0.0)
|
||||
parser.add_argument("--spatial-merge-size", type=int, default=2)
|
||||
parser.add_argument("--num-images", type=int, default=1)
|
||||
parser.add_argument("--image-patch-tokens", type=int, default=256)
|
||||
parser.add_argument("--num-videos", type=int, default=1)
|
||||
parser.add_argument("--video-patch-tokens", type=int, default=256)
|
||||
|
||||
# output
|
||||
parser.add_argument("--out-dir", type=str, default=".")
|
||||
args = parser.parse_args()
|
||||
print(args)
|
||||
|
||||
device = torch.device(args.device)
|
||||
|
||||
if args.num_tokens is None:
|
||||
num_tokens_list = [2**i for i in range(0, 18)]
|
||||
else:
|
||||
num_tokens_list = args.num_tokens
|
||||
|
||||
rows: list[dict[str, Any]] = []
|
||||
|
||||
for num_tokens in num_tokens_list:
|
||||
multimodal_stats, fallback_stats, speedup = benchmark_rope_index(
|
||||
model_name=args.model_name,
|
||||
tp_size=args.tp_size,
|
||||
num_tokens=num_tokens,
|
||||
batch_size=args.batch_size,
|
||||
pad_ratio=args.pad_ratio,
|
||||
spatial_merge_size=args.spatial_merge_size,
|
||||
num_images=args.num_images,
|
||||
image_patch_tokens=args.image_patch_tokens,
|
||||
num_videos=args.num_videos,
|
||||
video_patch_tokens=args.video_patch_tokens,
|
||||
dtype=getattr(torch, args.dtype),
|
||||
seed=args.seed,
|
||||
warmup_iter=args.warmup_iter,
|
||||
benchmark_iter=args.benchmark_iter,
|
||||
device=device,
|
||||
)
|
||||
Reference in New Issue
Block a user