wehub-resource-sync
87 lines
3.3 KiB
Python
87 lines
3.3 KiB
Python
import numpy as np
|
|
import pytest
|
|
import tvm
|
|
import tvm.testing
|
|
|
|
from mlc_llm.op.top_p_pivot import top_p_pivot, top_p_renorm
|
|
|
|
# mypy: disable-error-code="var-annotated"
|
|
|
|
# test category "op_correctness"
|
|
pytestmark = [pytest.mark.op_correctness]
|
|
|
|
|
|
@pytest.mark.parametrize("batch_size", [32, 64])
|
|
@pytest.mark.parametrize("vocab", [3, 32, 64, 128])
|
|
def test_top_p_renorm(batch_size, vocab):
|
|
top_p = 0.95
|
|
init_pivots_np = np.array([1 - top_p, 0.02, 0.01]).astype(np.float32)
|
|
top_p_np = np.array([top_p]).astype(np.float32)
|
|
|
|
p_np = np.random.exponential(3, size=(batch_size, vocab)).astype(np.float32)
|
|
p_np /= np.sum(p_np, axis=-1, keepdims=True)
|
|
final_pivot_np = np.zeros(batch_size).astype(np.float32)
|
|
final_lsum_np = np.zeros(batch_size).astype(np.float32)
|
|
|
|
dev = tvm.cuda(0)
|
|
var_prob = tvm.runtime.tensor(p_np, dev)
|
|
var_init_pivots = tvm.runtime.tensor(init_pivots_np, dev)
|
|
top_p_global = tvm.runtime.tensor(top_p_np, dev)
|
|
var_final_pivot = tvm.runtime.tensor(final_pivot_np, dev)
|
|
var_final_lsum = tvm.runtime.tensor(final_lsum_np, dev)
|
|
|
|
kernel = top_p_pivot(init_pivots_np.shape[0])
|
|
mod = tvm.build(kernel, target="cuda")
|
|
mod(var_prob, top_p_global, var_init_pivots, var_final_pivot, var_final_lsum)
|
|
|
|
final_pivot = var_final_pivot.asnumpy()
|
|
final_lsum = var_final_lsum.asnumpy()
|
|
|
|
renorm_np = p_np.copy()
|
|
var_renorm = tvm.runtime.tensor(renorm_np, dev)
|
|
|
|
kernel_renorm = top_p_renorm()
|
|
mod_renorm = tvm.build(kernel_renorm, target="cuda")
|
|
mod_renorm(var_prob, var_final_pivot, var_final_lsum, var_renorm)
|
|
|
|
renorm = var_renorm.asnumpy()
|
|
|
|
def verify_pivot(probs: np.ndarray, pivot: float, lsum: float, renorm: np.ndarray):
|
|
sorted_probs = np.sort(probs, axis=-1)[::-1]
|
|
num_larger_than_pivot = np.sum(sorted_probs >= pivot)
|
|
filtered_sorted_probs = sorted_probs[:num_larger_than_pivot]
|
|
min_larger_than_pivot = min(filtered_sorted_probs)
|
|
|
|
sum_larger_than_pivot = np.sum(np.where(sorted_probs >= pivot, sorted_probs, 0))
|
|
sum_larger_than_pivot_exclude_min = np.sum(
|
|
np.where(filtered_sorted_probs != min_larger_than_pivot, filtered_sorted_probs, 0)
|
|
)
|
|
|
|
probs[probs < pivot] = 0
|
|
renorm_prob = probs / np.sum(probs, axis=-1, keepdims=True)
|
|
try:
|
|
assert sum_larger_than_pivot >= top_p
|
|
assert sum_larger_than_pivot_exclude_min < top_p
|
|
assert abs(lsum - sum_larger_than_pivot) < 1e-6
|
|
assert np.allclose(renorm, renorm_prob, atol=1e-6, rtol=1e-6)
|
|
except AssertionError:
|
|
print("Failed")
|
|
print("probs:", repr(probs))
|
|
print("pivot:", pivot)
|
|
print("sorted_probs:", sorted_probs)
|
|
print("num_larger_than_pivot:", num_larger_than_pivot)
|
|
print("filtered_sorted_probs:", filtered_sorted_probs)
|
|
print("min_larger_than_pivot:", min_larger_than_pivot)
|
|
print("sum_larger_than_pivot:", sum_larger_than_pivot)
|
|
print("sum_larger_than_pivot_exclude_min:", sum_larger_than_pivot_exclude_min)
|
|
print("renom_prob:", renorm_prob)
|
|
print("renorm:", renorm)
|
|
raise
|
|
|
|
for i in range(batch_size):
|
|
verify_pivot(p_np[i], final_pivot[i], final_lsum[i], renorm[i])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|