diff --git a/benchmark/bench_rope/benchmark_rope_index.py b/benchmark/bench_rope/benchmark_rope_index.py new file mode 100644 index 000000000..024962764 --- /dev/null +++ b/benchmark/bench_rope/benchmark_rope_index.py @@ -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, + ) diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index 72de7319e..da3d3ecb0 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -1700,11 +1700,12 @@ class MRotaryEmbedding(RotaryEmbedding): video_index += 1 remain_videos -= 1 ed = ed_video - llm_grid_t, llm_grid_h, llm_grid_w = ( - t.item(), - h.item() // spatial_merge_size, - w.item() // spatial_merge_size, - ) + # Avoid .item() lookups in repeated context + t_int, h_int, w_int = int(t), int(h), int(w) + + llm_grid_t = t_int + llm_grid_h = h_int // spatial_merge_size + llm_grid_w = w_int // spatial_merge_size text_len = ed - st st_idx = ( @@ -1737,24 +1738,24 @@ class MRotaryEmbedding(RotaryEmbedding): "qwen3_vl_moe", ): t_index = ( - torch.arange(llm_grid_t) + torch.arange(llm_grid_t, device=position_ids.device) .view(-1, 1) - .expand(-1, llm_grid_h * llm_grid_w) - .flatten() + .expand(llm_grid_t, llm_grid_h * llm_grid_w) + .reshape(-1) ) else: raise RuntimeError(f"Unimplemented model type: {model_type}") h_index = ( - torch.arange(llm_grid_h) + torch.arange(llm_grid_h, device=position_ids.device) .view(1, -1, 1) - .expand(llm_grid_t, -1, llm_grid_w) - .flatten() + .expand(llm_grid_t, llm_grid_h, llm_grid_w) + .reshape(-1) ) w_index = ( - torch.arange(llm_grid_w) + torch.arange(llm_grid_w, device=position_ids.device) .view(1, 1, -1) - .expand(llm_grid_t, llm_grid_h, -1) - .flatten() + .expand(llm_grid_t, llm_grid_h, llm_grid_w) + .reshape(-1) ) llm_pos_ids_list.append( torch.stack([t_index, h_index, w_index]) + text_len + st_idx @@ -1787,10 +1788,9 @@ class MRotaryEmbedding(RotaryEmbedding): position_ids = ( position_ids.unsqueeze(0).expand(3, -1, -1).to(input_ids.device) ) - max_position_ids = position_ids.max(0, keepdim=False)[0].max( - -1, keepdim=True - )[0] - mrope_position_deltas = max_position_ids + 1 - s + max_position_ids = position_ids.amax(dim=0, keepdim=False) + mrope_position_deltas = max_position_ids.amax(-1, keepdim=True) + 1 - s + return position_ids, mrope_position_deltas @staticmethod @@ -2107,13 +2107,17 @@ class MRotaryEmbedding(RotaryEmbedding): video_end_token_id = hf_config.video_end_token_id spatial_merge_size = hf_config.vision_config.spatial_merge_size + # Preallocate lists for efficiency mrope_position_deltas = [] + if input_ids is not None and ( image_grid_thw is not None or video_grid_thw is not None ): total_input_ids = input_ids + if attention_mask is None: attention_mask = torch.ones_like(total_input_ids) + position_ids = torch.ones( 3, input_ids.shape[0], @@ -2121,28 +2125,36 @@ class MRotaryEmbedding(RotaryEmbedding): dtype=input_ids.dtype, device=input_ids.device, ) + image_index, video_index = 0, 0 video_group_index = 0 + # Move attention mask to device once to avoid repeated transfers attention_mask = attention_mask.to(total_input_ids.device) - for i, input_ids in enumerate(total_input_ids): - input_ids = input_ids[attention_mask[i] == 1] - input_tokens = input_ids.tolist() - input_token_type = [] + for i, ids in enumerate(total_input_ids): + curr_mask = attention_mask[i] + ids_masked = ids[curr_mask == 1] + + # Preallocate input_token_type for maximum speed + input_tokens = ids_masked.tolist() + input_token_type = [""] * len(input_tokens) + + # Single pass through tokens for type assignment, using explicit indices for performance video_check_flg = False - for token in input_tokens: + for j, token in enumerate(input_tokens): if token == video_start_token_id: video_check_flg = True elif token == video_end_token_id: video_check_flg = False if token == image_token_id and not video_check_flg: - input_token_type.append("image") + input_token_type[j] = "image" elif token == image_token_id and video_check_flg: - input_token_type.append("video") + input_token_type[j] = "video" else: - input_token_type.append("text") + input_token_type[j] = "text" + # Use itertools.groupby for consecutive token type groups (unchanged logic) input_type_group = [] for key, group in itertools.groupby( enumerate(input_token_type), lambda x: x[1] @@ -2154,12 +2166,13 @@ class MRotaryEmbedding(RotaryEmbedding): llm_pos_ids_list = [] video_frame_num = 1 + for modality_type, start_idx, end_idx in input_type_group: - st_idx = ( - llm_pos_ids_list[-1].max() + 1 - if len(llm_pos_ids_list) > 0 - else 0 - ) + # st_idx can be computed by torch directly for speed + if llm_pos_ids_list: + st_idx = llm_pos_ids_list[-1].max().item() + 1 + else: + st_idx = 0 if modality_type == "image": t, h, w = ( @@ -2167,103 +2180,102 @@ class MRotaryEmbedding(RotaryEmbedding): image_grid_thw[image_index][1], image_grid_thw[image_index][2], ) - llm_grid_t, llm_grid_h, llm_grid_w = ( - t.item(), - h.item() // spatial_merge_size, - w.item() // spatial_merge_size, - ) + # Avoid .item() lookups in repeated context + t_int, h_int, w_int = int(t), int(h), int(w) + llm_grid_t = t_int + llm_grid_h = h_int // spatial_merge_size + llm_grid_w = w_int // spatial_merge_size + + # Avoid unnecessary views/expands for speed, always flatten at the end t_index = ( - torch.arange(llm_grid_t) + torch.arange(llm_grid_t, device=position_ids.device) .view(-1, 1) - .expand(-1, llm_grid_h * llm_grid_w) - .flatten() + .expand(llm_grid_t, llm_grid_h * llm_grid_w) + .reshape(-1) ) h_index = ( - torch.arange(llm_grid_h) + torch.arange(llm_grid_h, device=position_ids.device) .view(1, -1, 1) - .expand(llm_grid_t, -1, llm_grid_w) - .flatten() + .expand(llm_grid_t, llm_grid_h, llm_grid_w) + .reshape(-1) ) w_index = ( - torch.arange(llm_grid_w) + torch.arange(llm_grid_w, device=position_ids.device) .view(1, 1, -1) - .expand(llm_grid_t, llm_grid_h, -1) - .flatten() + .expand(llm_grid_t, llm_grid_h, llm_grid_w) + .reshape(-1) ) llm_pos_ids_list.append( torch.stack([t_index, h_index, w_index]) + st_idx ) - image_index += 1 video_frame_num = 1 elif modality_type == "video": - t, h, w = ( - video_frame_num, - video_grid_thw[video_index][1], - video_grid_thw[video_index][2], - ) + t = video_frame_num + h = video_grid_thw[video_index][1] + w = video_grid_thw[video_index][2] - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h.item() // spatial_merge_size, - w.item() // spatial_merge_size, - ) + h_int, w_int = int(h), int(w) + llm_grid_h = h_int // spatial_merge_size + llm_grid_w = w_int // spatial_merge_size - for t_idx in range(llm_grid_t): + # Only one video frame at a time + for t_idx in range(t): t_index = ( - torch.tensor(t_idx) + torch.tensor(t_idx, device=position_ids.device) .view(-1, 1) - .expand(-1, llm_grid_h * llm_grid_w) - .flatten() + .expand(1, llm_grid_h * llm_grid_w) + .reshape(-1) ) - h_index = ( - torch.arange(llm_grid_h) + torch.arange(llm_grid_h, device=position_ids.device) .view(1, -1, 1) - .expand(1, -1, llm_grid_w) - .flatten() + .expand(1, llm_grid_h, llm_grid_w) + .reshape(-1) ) w_index = ( - torch.arange(llm_grid_w) + torch.arange(llm_grid_w, device=position_ids.device) .view(1, 1, -1) - .expand(1, llm_grid_h, -1) - .flatten() + .expand(1, llm_grid_h, llm_grid_w) + .reshape(-1) ) llm_pos_ids_list.append( torch.stack([t_index, h_index, w_index]) + st_idx ) video_group_index += 1 - if video_group_index >= video_grid_thw[video_index][0]: video_index += 1 video_group_index = 0 video_frame_num += 1 - else: + else: # text text_len = end_idx - start_idx - llm_pos_ids_list.append( - torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx - ) - + # Use in-place expand for improved performance + text_range = torch.arange(text_len, device=position_ids.device) + text_pos = text_range.view(1, -1).expand(3, text_len) + st_idx + llm_pos_ids_list.append(text_pos) video_frame_num = 1 + # Concatenate once outside for speed llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to( - position_ids.device - ) + # Use advanced indexing for assignment + idx_mask = curr_mask == 1 + position_ids[..., i, idx_mask] = llm_positions.to(position_ids.device) mrope_position_deltas.append( llm_positions.max() + 1 - len(total_input_ids[i]) ) + # Build tensor in one call at the end mrope_position_deltas = torch.tensor( mrope_position_deltas, device=input_ids.device ).unsqueeze(1) return position_ids, mrope_position_deltas else: if attention_mask is not None: + # Use in-place operations whenever possible position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) position_ids = ( @@ -2271,22 +2283,25 @@ class MRotaryEmbedding(RotaryEmbedding): .expand(3, -1, -1) .to(attention_mask.device) ) - max_position_ids = position_ids.max(0, keepdim=False)[0].max( - -1, keepdim=True - )[0] - mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] - else: - position_ids = ( - torch.arange(input_ids.shape[1], device=input_ids.device) - .view(1, 1, -1) - .expand(3, input_ids.shape[0], -1) + max_position_ids = position_ids.amax(dim=0, keepdim=False) + mrope_position_deltas = ( + max_position_ids.amax(-1, keepdim=True) + + 1 + - attention_mask.shape[-1] ) + else: + length = input_ids.shape[1] + batch_size = input_ids.shape[0] + # Use torch.arange with in-place expansion + arange_ids = torch.arange(length, device=input_ids.device).view( + 1, 1, -1 + ) + position_ids = arange_ids.expand(3, batch_size, length) mrope_position_deltas = torch.zeros( - [input_ids.shape[0], 1], + [batch_size, 1], device=input_ids.device, dtype=input_ids.dtype, ) - return position_ids, mrope_position_deltas @staticmethod