wehub-resource-sync
391 lines
16 KiB
Python
391 lines
16 KiB
Python
"""The pass that attaches GPU sampler functions to the IRModule."""
|
|
|
|
from typing import Dict # noqa: UP035
|
|
|
|
import tvm
|
|
from tvm import IRModule, relax, te, tirx
|
|
from tvm.relax.frontend import nn
|
|
from tvm.script import tirx as T
|
|
|
|
from mlc_llm.op.batch_spec_verify import batch_spec_verify
|
|
from mlc_llm.op.top_p_pivot import top_p_pivot, top_p_renorm
|
|
|
|
|
|
@tvm.transform.module_pass(opt_level=0, name="AttachGPUSamplingFunc")
|
|
class AttachGPUSamplingFunc:
|
|
"""Attach GPU sampling functions to IRModule."""
|
|
|
|
def __init__(self, target: tvm.target.Target, variable_bounds: Dict[str, int]): # noqa: UP006
|
|
# Specifically for RWKV workloads, which contains -1 max_seq_len
|
|
max_batch_size = variable_bounds["batch_size"]
|
|
self.variable_bounds = {
|
|
"batch_size": max_batch_size,
|
|
"num_samples": max_batch_size,
|
|
"num_positions": 6 * max_batch_size,
|
|
}
|
|
self.non_negative_var = ["vocab_size"]
|
|
self.target = target
|
|
|
|
def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
|
|
"""Entrypoint"""
|
|
target_kind = self.target.kind.name
|
|
if target_kind not in ["cuda", "vulkan", "metal", "webgpu"]:
|
|
# Only enable GPU sampling for CUDA, Vulkan, Metal, and WebGPU.
|
|
return mod
|
|
|
|
bb = relax.BlockBuilder(mod)
|
|
if target_kind == "webgpu":
|
|
# Only attach functions that do not contain i8s for WebGPU
|
|
gv_names = [
|
|
gv.name_hint
|
|
for gv in [
|
|
_attach_argsort_func(bb),
|
|
_attach_sample_with_top_p(bb),
|
|
]
|
|
]
|
|
else:
|
|
gv_names = [
|
|
gv.name_hint
|
|
for gv in [
|
|
_attach_multinomial_sampling_func(bb),
|
|
_attach_argsort_func(bb),
|
|
_attach_sample_with_top_p(bb),
|
|
_attach_take_probs_func(bb),
|
|
_attach_batch_verifier(bb),
|
|
_attach_renormalize_by_top_p(bb, self.target),
|
|
]
|
|
]
|
|
|
|
mod = bb.finalize()
|
|
for gv_name in gv_names:
|
|
mod[gv_name] = (
|
|
mod[gv_name]
|
|
.with_attr("tir_var_upper_bound", self.variable_bounds)
|
|
.with_attr("tir_non_negative_var", self.non_negative_var)
|
|
)
|
|
return mod
|
|
|
|
|
|
def _attach_multinomial_sampling_func(bb: relax.BlockBuilder):
|
|
batch_size = tirx.Var("batch_size", "int64")
|
|
num_samples = tirx.Var("num_samples", "int64")
|
|
vocab_size = tirx.Var("vocab_size", "int64")
|
|
probs = relax.Var("probs", relax.TensorType((batch_size, vocab_size), "float32"))
|
|
uniform_samples = relax.Var("uniform_samples", relax.TensorType((num_samples,), "float32"))
|
|
sample_indices = relax.Var("sample_indices", relax.TensorType((num_samples,), "int32"))
|
|
with bb.function("multinomial_from_uniform", [probs, uniform_samples, sample_indices]):
|
|
with bb.dataflow():
|
|
sample_shape = relax.ShapeExpr([num_samples, 1])
|
|
probs_tensor = nn.wrap_nested(probs, name="probs")
|
|
uniform_samples_tensor = nn.wrap_nested(
|
|
relax.call_pure_packed(
|
|
"vm.builtin.reshape",
|
|
uniform_samples,
|
|
sample_shape,
|
|
ty_args=relax.TensorType(sample_shape, "float32"),
|
|
),
|
|
name="uniform_samples",
|
|
)
|
|
sample_indices_tensor = nn.wrap_nested(
|
|
relax.call_pure_packed(
|
|
"vm.builtin.reshape",
|
|
sample_indices,
|
|
sample_shape,
|
|
ty_args=relax.TensorType(sample_shape, "int32"),
|
|
),
|
|
name="sample_indices",
|
|
)
|
|
result_tensor = nn.multinomial_from_uniform(
|
|
probs_tensor,
|
|
uniform_samples_tensor,
|
|
sample_indices_tensor,
|
|
"int32",
|
|
name="nn_multinomial_from_uniform",
|
|
)
|
|
result = bb.emit(
|
|
relax.call_pure_packed(
|
|
"vm.builtin.reshape",
|
|
result_tensor._expr,
|
|
sample_indices.ty.shape,
|
|
ty_args=sample_indices.ty,
|
|
)
|
|
)
|
|
output = bb.emit_output(result)
|
|
gv = bb.emit_func_output(output)
|
|
return gv
|
|
|
|
|
|
def _attach_argsort_func(bb: relax.BlockBuilder):
|
|
batch_size = tirx.Var("batch_size", "int64")
|
|
vocab_size = tirx.Var("vocab_size", "int64")
|
|
probs = relax.Var("probs", relax.TensorType((batch_size, vocab_size), "float32"))
|
|
with bb.function("argsort_probs", [probs]):
|
|
with bb.dataflow():
|
|
sorted_indices = bb.emit(relax.op.argsort(probs, descending=True, dtype="int32"))
|
|
sorted_values = bb.emit_te(
|
|
lambda unsorted_probs, sorted_indices: te.compute(
|
|
(batch_size, vocab_size),
|
|
lambda i, j: unsorted_probs[i, sorted_indices[i, j]],
|
|
name="take_sorted_probs",
|
|
),
|
|
probs,
|
|
sorted_indices,
|
|
primfunc_name_hint="take_sorted_probs",
|
|
)
|
|
output = bb.emit_output((sorted_values, sorted_indices))
|
|
gv = bb.emit_func_output(output)
|
|
return gv
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def full(var_result: T.handle, value: T.int32):
|
|
"""The filling function for top k."""
|
|
batch_size = T.int32()
|
|
result = T.match_buffer(var_result, (batch_size, 1), "int32")
|
|
for i in T.serial(batch_size):
|
|
with T.sblock("block"):
|
|
vi = T.axis.spatial(batch_size, i)
|
|
result[vi, 0] = value
|
|
|
|
|
|
def _attach_sample_with_top_p(bb: relax.BlockBuilder):
|
|
batch_size = tirx.Var("batch_size", "int64")
|
|
num_samples = tirx.Var("num_samples", "int64")
|
|
vocab_size = tirx.Var("vocab_size", "int64")
|
|
sorted_probs = relax.Var("sorted_probs", relax.TensorType((batch_size, vocab_size), "float32"))
|
|
sorted_indices = relax.Var(
|
|
"sorted_indices", relax.TensorType((batch_size, vocab_size), "int32")
|
|
)
|
|
uniform_samples = relax.Var("uniform_samples", relax.TensorType((num_samples,), "float32"))
|
|
sample_indices = relax.Var("sample_indices", relax.TensorType((num_samples,), "int32"))
|
|
top_p = relax.Var("top_p", relax.TensorType((batch_size,), "float32"))
|
|
|
|
with bb.function(
|
|
"sample_with_top_p",
|
|
[sorted_probs, sorted_indices, uniform_samples, sample_indices, top_p],
|
|
):
|
|
with bb.dataflow():
|
|
sample_shape = relax.ShapeExpr([num_samples, 1])
|
|
top_p_shape = relax.ShapeExpr([batch_size, 1])
|
|
sorted_probs_tensor = nn.wrap_nested(sorted_probs, name="sorted_probs")
|
|
sorted_indices_tensor = nn.wrap_nested(sorted_indices, name="sorted_indices")
|
|
uniform_samples_tensor = nn.wrap_nested(
|
|
relax.call_pure_packed(
|
|
"vm.builtin.reshape",
|
|
uniform_samples,
|
|
sample_shape,
|
|
ty_args=relax.TensorType(sample_shape, "float32"),
|
|
),
|
|
name="uniform_samples",
|
|
)
|
|
sample_indices_tensor = nn.wrap_nested(
|
|
relax.call_pure_packed(
|
|
"vm.builtin.reshape",
|
|
sample_indices,
|
|
sample_shape,
|
|
ty_args=relax.TensorType(sample_shape, "int32"),
|
|
),
|
|
name="sample_indices",
|
|
)
|
|
top_p_tensor = nn.wrap_nested(
|
|
relax.call_pure_packed(
|
|
"vm.builtin.reshape",
|
|
top_p,
|
|
top_p_shape,
|
|
ty_args=relax.TensorType(top_p_shape, "float32"),
|
|
),
|
|
name="sample_indices",
|
|
)
|
|
top_k_tensor = nn.tensor_ir_op(
|
|
full,
|
|
name_hint="full",
|
|
args=[vocab_size],
|
|
out=nn.Tensor.placeholder(
|
|
[batch_size, 1],
|
|
"int32",
|
|
),
|
|
)
|
|
|
|
result_tensor = nn.sample_top_p_top_k_from_sorted_prob(
|
|
sorted_probs_tensor,
|
|
sorted_indices_tensor,
|
|
top_p_tensor,
|
|
top_k_tensor,
|
|
uniform_samples_tensor,
|
|
sample_indices_tensor,
|
|
)
|
|
result = bb.emit_output(
|
|
relax.call_pure_packed(
|
|
"vm.builtin.reshape",
|
|
result_tensor._expr,
|
|
sample_indices.ty.shape,
|
|
ty_args=sample_indices.ty,
|
|
)
|
|
)
|
|
gv = bb.emit_func_output(result)
|
|
return gv
|
|
|
|
|
|
def _attach_renormalize_by_top_p(bb: relax.BlockBuilder, target: tvm.target.Target):
|
|
batch_size = tirx.Var("batch_size", "int64")
|
|
vocab_size = tirx.Var("vocab_size", "int64")
|
|
num_pivots = 3
|
|
probs = relax.Var("probs", relax.TensorType((batch_size, vocab_size), "float32"))
|
|
top_p = relax.Var("top_p", relax.TensorType((batch_size,), "float32"))
|
|
init_pivots = relax.Var("init_pivots", relax.TensorType((batch_size, num_pivots), "float32"))
|
|
with bb.function("renormalize_by_top_p", [probs, top_p, init_pivots]):
|
|
with bb.dataflow():
|
|
cutoff_output = bb.emit(
|
|
relax.call_tir(
|
|
bb.add_func(top_p_pivot(num_pivots, target), "top_p_pivot_cutoff"),
|
|
args=[probs, top_p, init_pivots],
|
|
out_ty=[top_p.ty, top_p.ty],
|
|
)
|
|
)
|
|
final_pivot = cutoff_output[0]
|
|
renorm_sum = cutoff_output[1]
|
|
renormalized_probs = bb.emit_output(
|
|
relax.call_tir(
|
|
bb.add_func(top_p_renorm(target), "top_p_renorm_after_cutoff"),
|
|
args=[probs, final_pivot, renorm_sum],
|
|
out_ty=probs.ty,
|
|
)
|
|
)
|
|
gv = bb.emit_func_output(renormalized_probs)
|
|
return gv
|
|
|
|
|
|
def _attach_take_probs_func(bb: relax.BlockBuilder):
|
|
batch_size = tirx.Var("batch_size", "int64")
|
|
num_samples = tirx.Var("num_samples", "int64")
|
|
num_positions = tirx.Var("num_positions", "int64")
|
|
vocab_size = tirx.Var("vocab_size", "int64")
|
|
unsorted_probs = relax.Var(
|
|
"unsorted_probs", relax.TensorType((batch_size, vocab_size), "float32")
|
|
)
|
|
sorted_indices = relax.Var(
|
|
"sorted_indices", relax.TensorType((batch_size, vocab_size), "int32")
|
|
)
|
|
sample_indices = relax.Var("sample_indices", relax.TensorType((num_samples,), "int32"))
|
|
sampling_results = relax.Var("sampling_result", relax.TensorType((num_samples,), "int32"))
|
|
top_prob_offsets = relax.Var("lobprob_offsets", relax.TensorType((num_positions,), "int32"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def sampler_take_probs_tir(
|
|
var_unsorted_probs: T.handle,
|
|
var_sorted_indices: T.handle,
|
|
var_sample_indices: T.handle,
|
|
var_sampling_results: T.handle,
|
|
var_top_prob_offsets: T.handle,
|
|
var_sampled_values: T.handle,
|
|
var_top_prob_probs: T.handle,
|
|
var_top_prob_indices: T.handle,
|
|
):
|
|
batch_size = T.int32()
|
|
num_samples = T.int32()
|
|
num_positions = T.int32()
|
|
vocab_size = T.int32()
|
|
unsorted_probs = T.match_buffer(var_unsorted_probs, (batch_size, vocab_size), "float32")
|
|
sorted_indices = T.match_buffer(var_sorted_indices, (batch_size, vocab_size), "int32")
|
|
sample_indices = T.match_buffer(var_sample_indices, (num_samples,), "int32")
|
|
sampling_results = T.match_buffer(var_sampling_results, (num_samples,), "int32")
|
|
top_prob_offsets = T.match_buffer(var_top_prob_offsets, (num_positions,), "int32")
|
|
sampled_values = T.match_buffer(var_sampled_values, (num_samples,), "float32")
|
|
top_prob_probs = T.match_buffer(var_top_prob_probs, (num_positions,), "float32")
|
|
top_prob_indices = T.match_buffer(var_top_prob_indices, (num_positions,), "int32")
|
|
for i in T.serial(num_positions):
|
|
with T.sblock("top_prob"):
|
|
vi = T.axis.spatial(num_positions, i)
|
|
# Reads are data-dependent gathers; declare full-buffer read
|
|
# regions explicitly so tirx does not infer data-dependent regions.
|
|
T.reads(
|
|
top_prob_offsets[vi],
|
|
sorted_indices[0:batch_size, 0:vocab_size],
|
|
unsorted_probs[0:batch_size, 0:vocab_size],
|
|
)
|
|
T.writes(top_prob_indices[vi], top_prob_probs[vi])
|
|
row = T.floordiv(top_prob_offsets[vi], vocab_size)
|
|
col = T.floormod(top_prob_offsets[vi], vocab_size)
|
|
top_prob_indices[vi] = sorted_indices[row, col]
|
|
top_prob_probs[vi] = unsorted_probs[row, sorted_indices[row, col]]
|
|
for i in T.serial(num_samples):
|
|
with T.sblock("sample"):
|
|
vj = T.axis.spatial(num_samples, i)
|
|
T.reads(
|
|
sample_indices[vj],
|
|
sampling_results[vj],
|
|
unsorted_probs[0:batch_size, 0:vocab_size],
|
|
)
|
|
T.writes(sampled_values[vj])
|
|
sampled_values[vj] = unsorted_probs[sample_indices[vj], sampling_results[vj]]
|
|
|
|
args = [
|
|
unsorted_probs,
|
|
sorted_indices,
|
|
sample_indices,
|
|
sampling_results,
|
|
top_prob_offsets,
|
|
]
|
|
with bb.function("sampler_take_probs", args):
|
|
with bb.dataflow():
|
|
taken_probs_indices = bb.emit_output(
|
|
relax.call_tir(
|
|
bb.add_func(sampler_take_probs_tir, "sampler_take_probs_tir"),
|
|
args,
|
|
out_ty=[
|
|
relax.TensorType((num_samples,), "float32"),
|
|
relax.TensorType((num_positions,), "float32"),
|
|
relax.TensorType((num_positions,), "int32"),
|
|
],
|
|
)
|
|
)
|
|
gv = bb.emit_func_output(taken_probs_indices)
|
|
return gv
|
|
|
|
|
|
def _attach_batch_verifier(bb: relax.BlockBuilder):
|
|
num_nodes = tirx.Var("num_nodes", "int64")
|
|
nbatch = tirx.Var("nbatch", "int64")
|
|
vocab_size = tirx.Var("vocab_size", "int64")
|
|
draft_probs = relax.Var("draft_probs", relax.TensorType((num_nodes, vocab_size), "float32"))
|
|
draft_tokens = relax.Var("draft_tokens", relax.TensorType((num_nodes,), "int32"))
|
|
model_probs = relax.Var("model_probs", relax.TensorType((num_nodes, vocab_size), "float32"))
|
|
token_tree_first_child = relax.Var(
|
|
"token_tree_first_child", relax.TensorType((num_nodes,), "int32")
|
|
)
|
|
token_tree_next_sibling = relax.Var(
|
|
"token_tree_next_sibling", relax.TensorType((num_nodes,), "int32")
|
|
)
|
|
uniform_samples = relax.Var("uniform_samples", relax.TensorType((num_nodes,), "float32"))
|
|
token_tree_parent_ptr = relax.Var("token_tree_parent_ptr", relax.TensorType((nbatch,), "int32"))
|
|
args = [
|
|
draft_probs,
|
|
draft_tokens,
|
|
model_probs,
|
|
token_tree_first_child,
|
|
token_tree_next_sibling,
|
|
uniform_samples,
|
|
token_tree_parent_ptr,
|
|
]
|
|
with bb.function("sampler_verify_draft_tokens", args):
|
|
with bb.dataflow():
|
|
res = bb.emit_output(
|
|
relax.call_tir_inplace(
|
|
bb.add_func(
|
|
batch_spec_verify(vocab_size),
|
|
"batch_verify_on_gpu_single_kernel",
|
|
),
|
|
args,
|
|
inplace_indices=[
|
|
args.index(model_probs),
|
|
args.index(token_tree_parent_ptr),
|
|
],
|
|
out_ty=[
|
|
model_probs.ty,
|
|
token_tree_parent_ptr.ty,
|
|
],
|
|
)
|
|
)
|
|
gv = bb.emit_func_output(res)
|
|
return gv
|