wehub-resource-sync
252 lines
9.3 KiB
Python
252 lines
9.3 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
tvm = pytest.importorskip("tvm")
|
|
from tvm import relax # noqa: E402
|
|
from tvm.relax.frontend import nn # noqa: E402
|
|
from tvm.relax.frontend.nn import spec # noqa: E402
|
|
from tvm.runtime import tensor as tvm_tensor # noqa: E402
|
|
|
|
from mlc_llm.op import ( # noqa: E402
|
|
MultimodalRotaryEmbedding,
|
|
VisionPositionMetadata,
|
|
apply_multimodal_rotary_pos_emb,
|
|
get_mrope_position_ids,
|
|
)
|
|
|
|
|
|
def _numpy_rotate_half(x: np.ndarray) -> np.ndarray:
|
|
x1, x2 = np.split(x, 2, axis=-1)
|
|
return np.concatenate([-x2, x1], axis=-1)
|
|
|
|
|
|
def _numpy_apply_mrope(
|
|
q: np.ndarray,
|
|
k: np.ndarray,
|
|
position_ids: np.ndarray,
|
|
theta: float,
|
|
mrope_section: tuple[int, ...],
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
if position_ids.ndim != 3:
|
|
raise ValueError(f"position_ids must be rank-3, got shape {position_ids.shape}")
|
|
if position_ids.shape[0] == 3:
|
|
position_ids = np.transpose(position_ids, (1, 2, 0))
|
|
elif position_ids.shape[-1] != 3:
|
|
raise ValueError(
|
|
"position_ids must have shape (batch, seq, 3) or (3, batch, seq), "
|
|
f"got {position_ids.shape}"
|
|
)
|
|
|
|
head_dim = q.shape[-1]
|
|
inv_freq = 1.0 / (theta ** (np.arange(0, head_dim, 2, dtype=np.float32) / float(head_dim)))
|
|
pos = np.transpose(position_ids, (2, 0, 1))
|
|
inv = inv_freq.reshape(1, 1, -1, 1).astype(np.float32)
|
|
inv = np.broadcast_to(inv, (3, pos.shape[1], inv_freq.size, 1))
|
|
pos = pos.reshape(3, pos.shape[1], 1, pos.shape[2]).astype(np.float32)
|
|
freqs = np.matmul(inv, pos)
|
|
freqs = np.transpose(freqs, (0, 1, 3, 2))
|
|
emb = np.concatenate([freqs, freqs], axis=-1)
|
|
cos = np.cos(emb)
|
|
sin = np.sin(emb)
|
|
split_sizes = list(mrope_section) * 2
|
|
split_points = np.cumsum(split_sizes)[:-1]
|
|
cos_chunks = np.split(cos, split_points, axis=-1)
|
|
sin_chunks = np.split(sin, split_points, axis=-1)
|
|
cos = np.concatenate([chunk[idx % 3] for idx, chunk in enumerate(cos_chunks)], axis=-1)
|
|
sin = np.concatenate([chunk[idx % 3] for idx, chunk in enumerate(sin_chunks)], axis=-1)
|
|
cos = np.expand_dims(cos, axis=2)
|
|
sin = np.expand_dims(sin, axis=2)
|
|
q_out = q * cos + _numpy_rotate_half(q) * sin
|
|
k_out = k * cos + _numpy_rotate_half(k) * sin
|
|
return q_out, k_out
|
|
|
|
|
|
def _evaluate_tensor(expr):
|
|
mod = tvm.IRModule.from_expr(expr)
|
|
target = tvm.target.Target("llvm")
|
|
ex = tvm.relax.build(mod, target)
|
|
vm = tvm.relax.VirtualMachine(ex, tvm.cpu())
|
|
return vm["main"]().numpy()
|
|
|
|
|
|
def _run_mlc_mrope(
|
|
q_np: np.ndarray,
|
|
k_np: np.ndarray,
|
|
position_ids_np: np.ndarray,
|
|
theta: float,
|
|
mrope_section: tuple[int, ...],
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
class RopeModule(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.rotary = MultimodalRotaryEmbedding(q_np.shape[-1], theta, mrope_section)
|
|
|
|
def forward(
|
|
self,
|
|
q: nn.Tensor,
|
|
k: nn.Tensor,
|
|
pos: nn.Tensor,
|
|
):
|
|
"""Run MRoPE on test tensors and return rotated query/key outputs."""
|
|
cos, sin = self.rotary(q, pos)
|
|
return apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section)
|
|
|
|
module = RopeModule()
|
|
mod, _, _ = module.export_tvm(
|
|
spec={
|
|
"forward": {
|
|
"q": spec.Tensor(q_np.shape, "float32"),
|
|
"k": spec.Tensor(k_np.shape, "float32"),
|
|
"pos": spec.Tensor(position_ids_np.shape, "int64"),
|
|
}
|
|
},
|
|
allow_extern=True,
|
|
)
|
|
target = tvm.target.Target("llvm")
|
|
exec_mod = relax.build(mod, target=target)
|
|
vm = relax.VirtualMachine(exec_mod, tvm.cpu())
|
|
device = tvm.cpu()
|
|
q_nd = tvm_tensor(q_np.astype("float32"), device=device)
|
|
k_nd = tvm_tensor(k_np.astype("float32"), device=device)
|
|
pos_nd = tvm_tensor(position_ids_np.astype("int64"), device=device)
|
|
out_q, out_k = vm["forward"](q_nd, k_nd, pos_nd)
|
|
return out_q.numpy(), out_k.numpy()
|
|
|
|
|
|
def test_apply_mrope_matches_numpy_reference():
|
|
theta = 10000.0
|
|
mrope_section = (2, 2, 2)
|
|
batch, seq_len, heads, head_dim = 1, 4, 2, 12
|
|
rng = np.random.default_rng(0)
|
|
q_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
|
|
k_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
|
|
position_ids = np.zeros((batch, seq_len, 3), dtype=np.int64)
|
|
position_ids[0, :, 0] = np.arange(seq_len)
|
|
position_ids[0, :, 1] = np.arange(seq_len) * 2
|
|
position_ids[0, :, 2] = np.arange(seq_len) * 3
|
|
|
|
mlc_q, mlc_k = _run_mlc_mrope(q_np, k_np, position_ids, theta, mrope_section)
|
|
ref_q, ref_k = _numpy_apply_mrope(q_np, k_np, position_ids, theta, mrope_section)
|
|
|
|
np.testing.assert_allclose(mlc_q, ref_q, rtol=1e-5, atol=1e-5)
|
|
np.testing.assert_allclose(mlc_k, ref_k, rtol=1e-5, atol=1e-5)
|
|
|
|
|
|
def test_get_mrope_position_ids_text_only():
|
|
input_ids = np.array([[1, 2, 3, 0, 0]], dtype=np.int64)
|
|
attention_mask = np.array([[1, 1, 1, 0, 0]], dtype=np.int64)
|
|
meta = VisionPositionMetadata(
|
|
vision_start_token_id=1000,
|
|
image_token_id=1001,
|
|
video_token_id=1002,
|
|
spatial_merge_size=2,
|
|
tokens_per_second=4.0,
|
|
)
|
|
position_ids, deltas = get_mrope_position_ids(
|
|
input_ids,
|
|
meta,
|
|
attention_mask=attention_mask,
|
|
image_grid_thw=None,
|
|
video_grid_thw=None,
|
|
second_per_grid_ts=None,
|
|
)
|
|
expected = attention_mask.cumsum(axis=-1) - 1
|
|
expected = np.where(attention_mask == 0, 1, expected)
|
|
expected = np.expand_dims(expected, axis=0).repeat(3, axis=0)
|
|
np.testing.assert_array_equal(position_ids, expected)
|
|
np.testing.assert_array_equal(deltas, np.array([[-2]], dtype=np.int64))
|
|
|
|
|
|
def test_get_mrope_position_ids_single_image_block():
|
|
meta = VisionPositionMetadata(
|
|
vision_start_token_id=5000,
|
|
image_token_id=5001,
|
|
video_token_id=6000,
|
|
spatial_merge_size=2,
|
|
tokens_per_second=4.0,
|
|
)
|
|
input_ids = np.array(
|
|
[[11, 12, 5000, 5001, 21, 22, 23, 24, 31, 32]],
|
|
dtype=np.int64,
|
|
)
|
|
attention_mask = np.ones_like(input_ids, dtype=np.int64)
|
|
image_grid_thw = np.array([[1, 4, 4]], dtype=np.int64)
|
|
position_ids, deltas = get_mrope_position_ids(
|
|
input_ids,
|
|
meta,
|
|
attention_mask=attention_mask,
|
|
image_grid_thw=image_grid_thw,
|
|
video_grid_thw=None,
|
|
second_per_grid_ts=None,
|
|
)
|
|
expected = np.array(
|
|
[
|
|
[0, 1, 2, 3, 3, 3, 3, 5, 6, 7],
|
|
[0, 1, 2, 3, 3, 4, 4, 5, 6, 7],
|
|
[0, 1, 2, 3, 4, 3, 4, 5, 6, 7],
|
|
],
|
|
dtype=np.int64,
|
|
).reshape(3, 1, -1)
|
|
np.testing.assert_array_equal(position_ids, expected)
|
|
np.testing.assert_array_equal(deltas, np.array([[-2]], dtype=np.int64))
|
|
|
|
|
|
def test_apply_mrope_accepts_3_batch_seq_layout():
|
|
theta = 10000.0
|
|
mrope_section = (2, 2, 2)
|
|
batch, seq_len, heads, head_dim = 1, 4, 2, 12
|
|
rng = np.random.default_rng(1)
|
|
q_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
|
|
k_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
|
|
|
|
position_ids_bsc = np.zeros((batch, seq_len, 3), dtype=np.int64)
|
|
position_ids_bsc[0, :, 0] = np.arange(seq_len)
|
|
position_ids_bsc[0, :, 1] = np.arange(seq_len) * 2
|
|
position_ids_bsc[0, :, 2] = np.arange(seq_len) * 3
|
|
position_ids_3bs = np.transpose(position_ids_bsc, (2, 0, 1))
|
|
|
|
mlc_q_bsc, mlc_k_bsc = _run_mlc_mrope(q_np, k_np, position_ids_bsc, theta, mrope_section)
|
|
mlc_q_3bs, mlc_k_3bs = _run_mlc_mrope(q_np, k_np, position_ids_3bs, theta, mrope_section)
|
|
ref_q, ref_k = _numpy_apply_mrope(q_np, k_np, position_ids_bsc, theta, mrope_section)
|
|
|
|
np.testing.assert_allclose(mlc_q_bsc, ref_q, rtol=1e-5, atol=1e-5)
|
|
np.testing.assert_allclose(mlc_k_bsc, ref_k, rtol=1e-5, atol=1e-5)
|
|
np.testing.assert_allclose(mlc_q_3bs, ref_q, rtol=1e-5, atol=1e-5)
|
|
np.testing.assert_allclose(mlc_k_3bs, ref_k, rtol=1e-5, atol=1e-5)
|
|
|
|
|
|
def test_get_mrope_position_ids_output_is_directly_usable():
|
|
theta = 10000.0
|
|
mrope_section = (2, 2, 2)
|
|
meta = VisionPositionMetadata(
|
|
vision_start_token_id=7000,
|
|
image_token_id=7001,
|
|
video_token_id=7002,
|
|
spatial_merge_size=2,
|
|
tokens_per_second=4.0,
|
|
)
|
|
input_ids = np.array([[11, 12, 7000, 7001, 21, 22, 23, 24, 31, 32]], dtype=np.int64)
|
|
attention_mask = np.ones_like(input_ids, dtype=np.int64)
|
|
image_grid_thw = np.array([[1, 4, 4]], dtype=np.int64)
|
|
position_ids_3bs, _ = get_mrope_position_ids(
|
|
input_ids,
|
|
meta,
|
|
attention_mask=attention_mask,
|
|
image_grid_thw=image_grid_thw,
|
|
video_grid_thw=None,
|
|
second_per_grid_ts=None,
|
|
)
|
|
position_ids_bsc = np.transpose(position_ids_3bs, (1, 2, 0))
|
|
|
|
batch, seq_len = input_ids.shape
|
|
heads, head_dim = 2, 12
|
|
rng = np.random.default_rng(2)
|
|
q_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
|
|
k_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
|
|
|
|
mlc_q_3bs, mlc_k_3bs = _run_mlc_mrope(q_np, k_np, position_ids_3bs, theta, mrope_section)
|
|
mlc_q_bsc, mlc_k_bsc = _run_mlc_mrope(q_np, k_np, position_ids_bsc, theta, mrope_section)
|
|
|
|
np.testing.assert_allclose(mlc_q_3bs, mlc_q_bsc, rtol=1e-5, atol=1e-5)
|
|
np.testing.assert_allclose(mlc_k_3bs, mlc_k_bsc, rtol=1e-5, atol=1e-5)
|