wehub-resource-sync
336 lines
15 KiB
Python
336 lines
15 KiB
Python
"""Operators for choosing the pivot to cut-off top-p percentile"""
|
|
|
|
import tvm
|
|
from tvm.script import tirx as T
|
|
|
|
from mlc_llm.support.max_thread_check import get_max_num_threads_per_block
|
|
|
|
# mypy: disable-error-code="attr-defined,valid-type,name-defined"
|
|
|
|
|
|
def top_p_pivot(pN, target: tvm.target.Target):
|
|
"""Top-p pivot function. This function finds the pivot to cut-off top-p percentile.
|
|
|
|
A valide pivot should satisfy the following conditions:
|
|
- lsum >= top_p
|
|
- top_p > lsum - cmin * lmin
|
|
where lsum is the sum of elements that are larger or equal to the pivot,
|
|
lmin is the minimum elements that is larger or equal to the pivot,
|
|
cmin is the count of elements that are equal to lmin,
|
|
|
|
Parameters
|
|
----------
|
|
prob:
|
|
The probability vector
|
|
|
|
top_p_arr:
|
|
The top-p threshold
|
|
|
|
init_pivots:
|
|
The initial pivot candidates
|
|
|
|
final_pivot:
|
|
The final pivot to cut-off top-p percentile
|
|
|
|
final_lsum:
|
|
The final sum of the values after top-p filtering.
|
|
"""
|
|
TX = 1024
|
|
K = 32
|
|
eps_LR = 1e-7
|
|
|
|
max_num_threads_per_block = get_max_num_threads_per_block(target)
|
|
TX = min(TX, max_num_threads_per_block)
|
|
|
|
def _var(dtype="int32"):
|
|
return T.sblock_alloc_buffer((1,), dtype, scope="local")
|
|
|
|
def valid(lsum, lmin, cmin, top_p):
|
|
return tvm.tirx.all(lsum >= top_p, top_p > lsum - cmin * lmin)
|
|
|
|
# fmt: off
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def _func(
|
|
var_prob: T.handle,
|
|
var_top_p_arr: T.handle,
|
|
var_init_pivots: T.handle,
|
|
var_final_pivot: T.handle,
|
|
var_final_lsum: T.handle,
|
|
):
|
|
T.func_attr({"tirx.is_scheduled": 1, "tirx.noalias": True})
|
|
B = T.int32()
|
|
N = T.int32()
|
|
prob = T.match_buffer(var_prob, (B, N,), "float32")
|
|
top_p_arr = T.match_buffer(var_top_p_arr, (B,), dtype="float32")
|
|
init_pivots = T.match_buffer(var_init_pivots, (B, pN), "float32")
|
|
final_pivot = T.match_buffer(var_final_pivot, (B,), "float32")
|
|
final_lsum = T.match_buffer(var_final_lsum, (B,), "float32")
|
|
|
|
with T.sblock("kernel"):
|
|
pivot = T.sblock_alloc_buffer((pN,), "float32", scope="local")
|
|
top_p = _var("float32")
|
|
|
|
L = T.sblock_alloc_buffer((1,), "float32", scope="shared")
|
|
R = T.sblock_alloc_buffer((1,), "float32", scope="shared")
|
|
L_local = _var("float32")
|
|
R_local = _var("float32")
|
|
|
|
q = _var("float32")
|
|
lsum = T.sblock_alloc_buffer((pN,), "float32", scope="local")
|
|
lmin_broadcast = T.sblock_alloc_buffer((1), "float32", scope="shared")
|
|
lmin_broadcast_local = _var("float32")
|
|
lmin = T.sblock_alloc_buffer((pN,), "float32", scope="local")
|
|
cmin = T.sblock_alloc_buffer((pN,), "int32", scope="local")
|
|
total_sum = _var("float32")
|
|
|
|
it = _var("int32")
|
|
es_local = _var("bool")
|
|
es = T.sblock_alloc_buffer((1,), "bool", scope="shared")
|
|
find_pivot_local = _var("bool")
|
|
find_pivot = T.sblock_alloc_buffer((1,), "bool", scope="shared")
|
|
|
|
total_sum_reduce = _var("float32")
|
|
lsum_reduce = _var("float32")
|
|
lmin_reduce = _var("float32")
|
|
cmin_reduce = _var("int32")
|
|
|
|
for _bx in T.thread_binding(0, B, thread="blockIdx.x"):
|
|
for _tx in T.thread_binding(0, TX, thread="threadIdx.x"):
|
|
with T.sblock("CTA"):
|
|
b, tx = T.axis.remap("SS", [_bx, _tx])
|
|
|
|
top_p[0] = top_p_arr[b]
|
|
|
|
if tx == 0:
|
|
# leader thread initializes L, R
|
|
L[0] = 1.0 - top_p[0]
|
|
R[0] = eps_LR
|
|
find_pivot[0] = False
|
|
T.tvm_storage_sync("shared")
|
|
|
|
L_local[0] = L[0]
|
|
R_local[0] = R[0]
|
|
for i in T.unroll(0, pN):
|
|
# pivots are in descending order
|
|
pivot[i] = init_pivots[b, i]
|
|
find_pivot_local[0] = False
|
|
if L_local[0] - R_local[0] <= eps_LR:
|
|
# When the initial value is too small, set the result directly.
|
|
if tx == 0:
|
|
final_lsum[b] = 1.0
|
|
final_pivot[b] = 0.0
|
|
find_pivot_local[0] = True
|
|
|
|
while T.tvm_thread_invariant(
|
|
L_local[0] - R_local[0] > eps_LR
|
|
and T.Not(find_pivot_local[0])
|
|
):
|
|
# sync before each iteration
|
|
T.tvm_storage_sync("shared")
|
|
|
|
### get lsum, lmin, total_sum
|
|
for pidx in T.unroll(0, pN):
|
|
lsum[pidx] = 0.0
|
|
lmin[pidx] = T.max_value("float32")
|
|
cmin[pidx] = 0
|
|
total_sum[0] = 0.0
|
|
it[0] = 0
|
|
es_local[0] = False
|
|
while it[0] < T.ceildiv(N, TX) and T.Not(es_local[0]):
|
|
idx = T.meta_var(it[0] * TX + tx)
|
|
q[0] = T.if_then_else(idx < N, prob[b, idx], 0.0)
|
|
total_sum[0] += q[0]
|
|
for pidx in T.unroll(0, pN):
|
|
if q[0] >= pivot[pidx]:
|
|
lsum[pidx] += q[0]
|
|
if lmin[pidx] > q[0]:
|
|
lmin[pidx] = q[0]
|
|
cmin[pidx] = 1
|
|
elif lmin[pidx] == q[0]:
|
|
cmin[pidx] += 1
|
|
it[0] += 1
|
|
|
|
# early stop every K iterations
|
|
if it[0] % K == 0:
|
|
# reduce total_sum over tx
|
|
# T.tvm_storage_sync("shared")
|
|
with T.sblock("block_cross_thread"):
|
|
T.reads(total_sum[0])
|
|
T.writes(total_sum_reduce[0])
|
|
T.attr(
|
|
T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
|
|
"reduce_scope",
|
|
T.int32(0),
|
|
)
|
|
T.tvm_thread_allreduce(T.uint32(1), total_sum[0], True, total_sum_reduce[0], tx, dtype="void") # noqa: E501
|
|
# T.tvm_storage_sync("shared")
|
|
|
|
if tx == 0:
|
|
# leader thread checks if we can stop early
|
|
es[0] = 1 - total_sum_reduce[0] < pivot[pN - 1]
|
|
T.tvm_storage_sync("shared")
|
|
es_local[0] = es[0]
|
|
|
|
T.tvm_storage_sync("shared")
|
|
|
|
# reduce lsum, lmin, cmin, over tx
|
|
for pidx in T.serial(0, pN):
|
|
# reduce lsum over tx for pivot[j]
|
|
with T.sblock("block_cross_thread"):
|
|
T.reads(lsum[pidx])
|
|
T.writes(lsum_reduce[0])
|
|
T.attr(
|
|
T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
|
|
"reduce_scope",
|
|
T.int32(0),
|
|
)
|
|
T.tvm_thread_allreduce(T.uint32(1), lsum[pidx], True, lsum_reduce[0], tx, dtype="void") # noqa: E501
|
|
|
|
# reduce lmin over tx for pivot[j]
|
|
with T.sblock("block_cross_thread"):
|
|
T.reads(lmin[pidx])
|
|
T.writes(lmin_reduce[0])
|
|
T.attr(
|
|
T.comm_reducer(lambda x0, y0: T.min(x0, y0), [T.float32(0)]), # noqa: E501
|
|
"reduce_scope",
|
|
T.int32(0),
|
|
)
|
|
T.tvm_thread_allreduce(T.uint32(1), lmin[pidx], True, lmin_reduce[0], tx, dtype="void") # noqa: E501
|
|
|
|
if tx == 0:
|
|
# broadcast lmin to all threads
|
|
lmin_broadcast[0] = lmin_reduce[0]
|
|
T.tvm_storage_sync("shared")
|
|
lmin_broadcast_local[0] = lmin_broadcast[0]
|
|
if lmin[pidx] > lmin_broadcast_local[0]:
|
|
cmin[pidx] = 0
|
|
if tx == 0:
|
|
# only the leader thread updates lsum, lmin
|
|
lsum[pidx] = lsum_reduce[0]
|
|
lmin[pidx] = lmin_reduce[0]
|
|
|
|
# reduce cmin over tx for pivot[j]
|
|
with T.sblock("block_cross_thread"):
|
|
T.reads(cmin[pidx])
|
|
T.writes(cmin_reduce[0])
|
|
T.attr(
|
|
T.comm_reducer(lambda x0, y0: x0 + y0, [T.int32(0)]),
|
|
"reduce_scope",
|
|
T.int32(0),
|
|
)
|
|
T.tvm_thread_allreduce(T.uint32(1), cmin[pidx], True, cmin_reduce[0], tx, dtype="void") # noqa: E501
|
|
|
|
if tx == 0:
|
|
# only the leader thread updates cmin
|
|
cmin[pidx] = cmin_reduce[0]
|
|
|
|
T.tvm_storage_sync("shared")
|
|
|
|
if tx == 0:
|
|
# leader thread checks if we have found the pivot, or updates L, R
|
|
it[0] = 0
|
|
while it[0] < pN and T.Not(find_pivot_local[0]):
|
|
pidx = T.meta_var(it[0])
|
|
if valid(lsum[pidx], lmin[pidx], cmin[pidx], top_p[0]):
|
|
find_pivot[0] = True
|
|
find_pivot_local[0] = True
|
|
# write back the pivot and lsum
|
|
final_pivot[b] = pivot[pidx]
|
|
final_lsum[b] = lsum[pidx]
|
|
elif lsum[pidx] - lmin[pidx] * cmin[pidx] >= top_p[0]:
|
|
R[0] = pivot[pidx]
|
|
final_lsum[b] = lsum[pidx]
|
|
elif lsum[pidx] < top_p[0]:
|
|
L[0] = pivot[pidx]
|
|
it[0] += 1
|
|
|
|
T.tvm_storage_sync("shared")
|
|
|
|
L_local[0] = L[0]
|
|
R_local[0] = R[0]
|
|
find_pivot_local[0] = find_pivot[0]
|
|
# new pivots for next iteration
|
|
# uniform spacing between L and R
|
|
for pidx in T.unroll(0, pN):
|
|
pivot[pidx] = L[0] - (pidx + 1) * (L_local[0] - R_local[0]) / (pN + 1) # noqa: E501
|
|
|
|
if tx == 0:
|
|
# leader thread writes back the pivot
|
|
if T.Not(find_pivot_local[0]):
|
|
final_pivot[b] = R_local[0]
|
|
if R_local[0] == eps_LR:
|
|
final_lsum[b] = lsum[pN - 1]
|
|
# fmt: on
|
|
|
|
return _func
|
|
|
|
|
|
def top_p_renorm(target: tvm.target.Target = None):
|
|
"""Top-p renormalization function. This function renormalizes the probability vector.
|
|
|
|
Given the pivot, the probability vector is renormalized as follows:
|
|
- if prob >= pivot, renorm_prob = prob / lsum
|
|
- otherwise, renorm_prob = 0
|
|
|
|
Parameters
|
|
----------
|
|
prob:
|
|
The probability vector
|
|
|
|
final_pivot:
|
|
The final pivot to cut-off top-p percentile
|
|
|
|
final_lsum:
|
|
The sum of elements that are larger or equal to the pivot
|
|
|
|
renorm_prob:
|
|
The renormalized probability vector
|
|
"""
|
|
TX = 1024
|
|
CTA_COUNT = 512
|
|
|
|
if target:
|
|
max_num_threads_per_block = get_max_num_threads_per_block(target)
|
|
TX = min(TX, max_num_threads_per_block)
|
|
|
|
def _var(dtype="int32"):
|
|
return T.sblock_alloc_buffer((1,), dtype, scope="local")
|
|
|
|
# fmt: off
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def _func(
|
|
var_prob: T.handle,
|
|
var_final_pivot: T.handle,
|
|
var_final_lsum: T.handle,
|
|
var_renorm_prob: T.handle,
|
|
):
|
|
T.func_attr({"tirx.is_scheduled": 1, "tirx.noalias": True})
|
|
B = T.int32()
|
|
N = T.int32()
|
|
prob = T.match_buffer(var_prob, (B, N,), "float32")
|
|
final_pivot = T.match_buffer(var_final_pivot, (B,), "float32")
|
|
final_lsum = T.match_buffer(var_final_lsum, (B,), "float32")
|
|
renorm_prob = T.match_buffer(var_renorm_prob, (B, N,), "float32")
|
|
|
|
with T.sblock("kernel"):
|
|
pivot = _var("float32")
|
|
lsum = _var("float32")
|
|
BX = T.meta_var(T.ceildiv(CTA_COUNT, B))
|
|
|
|
for _by in T.thread_binding(0, B, thread="blockIdx.y"):
|
|
for _bx in T.thread_binding(0, BX, thread="blockIdx.x"):
|
|
for _tx in T.thread_binding(0, TX, thread="threadIdx.x"):
|
|
with T.sblock("CTA"):
|
|
by, bx, tx = T.axis.remap("SSS", [_by, _bx, _tx])
|
|
|
|
pivot[0] = final_pivot[by]
|
|
lsum[0] = final_lsum[by]
|
|
|
|
for i in T.serial(T.ceildiv(N, BX * TX)):
|
|
idx = T.meta_var(i * BX * TX + bx * TX + tx)
|
|
if idx < N:
|
|
renorm_prob[by, idx] = T.if_then_else(prob[by, idx] >= pivot[0], prob[by, idx] / lsum[0], 0.0) # noqa: E501
|
|
# fmt: on
|
|
|
|
return _func
|