[VLM] Optimize get_rope_index for GLM4v (#17420)
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
425
benchmark/bench_rope/benchmark_rope_index.py
Normal file
425
benchmark/bench_rope/benchmark_rope_index.py
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,
|
||||
)
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user