wehub-resource-sync
115 lines
4.4 KiB
Python
115 lines
4.4 KiB
Python
import tvm
|
|
from tvm import tirx
|
|
from tvm.relax.frontend.nn import core, modules, spec
|
|
from tvm.script import ir as I
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
from mlc_llm.nn.kv_cache import PagedKVCache, RopeMode
|
|
|
|
# mypy: disable-error-code="attr-defined"
|
|
|
|
|
|
def test_nn_module_paged_kv_cache():
|
|
# fmt: off
|
|
@I.ir_module
|
|
class Module:
|
|
@R.function
|
|
def create_paged_kv_cache(
|
|
max_batch_size: R.Shape(["max_batch_size_1"]),
|
|
max_total_seq_len: R.Shape(["max_total_seq_len_1"]),
|
|
prefill_chunk_size: R.Shape(["prefill_chunk_size_1"]),
|
|
page_size: R.Shape(["page_size_1"]),
|
|
support_sliding_window: R.Shape(["support_sliding_window_1"]),
|
|
) -> R.Object:
|
|
max_batch_size_1 = T.int64()
|
|
max_total_seq_len_1 = T.int64()
|
|
prefill_chunk_size_1 = T.int64()
|
|
page_size_1 = T.int64()
|
|
support_sliding_window_1 = T.int64()
|
|
R.func_attr({"num_input": 5})
|
|
with R.dataflow():
|
|
paged_kv_cache: R.Object = R.call_pure_packed("mlc.create_paged_kv_cache_generic", R.shape([max_batch_size_1, max_total_seq_len_1, prefill_chunk_size_1, page_size_1, support_sliding_window_1]), R.prim_value(32), R.prim_value(32), R.prim_value(32), R.prim_value(128), R.prim_value(1), R.prim_value(1), R.prim_value(10000), R.prim_value(128), R.dtype("float16"), sinfo_args=(R.Object,)) # noqa: E501
|
|
gv1: R.Object = paged_kv_cache
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
@R.function
|
|
def forward(
|
|
cache: R.Object, qkv: R.Tensor((1, 100, 96, 128), dtype="float16")
|
|
) -> R.Tensor((1, 100, 32, 128), dtype="float16"):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
reshape: R.Tensor((100, 96, 128), dtype="float16") = R.reshape(
|
|
qkv, R.shape([100, 96, 128])
|
|
)
|
|
lv = R.call_dps_packed(
|
|
"vm.builtin.attention_kv_cache_attention_with_fused_qkv",
|
|
(cache, R.prim_value(0), R.prim_value(T.float32(1)), reshape),
|
|
out_sinfo=R.Tensor((100, 32, 128), dtype="float16"),
|
|
)
|
|
reshape1: R.Tensor((1, 100, 32, 128), dtype="float16") = R.reshape(
|
|
lv, R.shape([1, 100, 32, 128])
|
|
)
|
|
gv: R.Tensor((1, 100, 32, 128), dtype="float16") = reshape1
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
class PagedKVCacheTest(modules.Module):
|
|
def forward(
|
|
self,
|
|
cache: PagedKVCache,
|
|
qkv: core.Tensor,
|
|
) -> core.Tensor:
|
|
return cache.attention_with_fused_qkv(0, qkv, num_qo_heads=32, sm_scale=128**-0.5)
|
|
|
|
def create_paged_kv_cache(
|
|
self,
|
|
max_batch_size: tirx.Var,
|
|
max_total_seq_len: tirx.Var,
|
|
prefill_chunk_size: tirx.Var,
|
|
page_size: tirx.Var,
|
|
support_sliding_window: tirx.Var,
|
|
) -> PagedKVCache:
|
|
return PagedKVCache.create_generic(
|
|
attn_kind="mha",
|
|
max_batch_size=max_batch_size,
|
|
max_total_seq_len=max_total_seq_len,
|
|
prefill_chunk_size=prefill_chunk_size,
|
|
page_size=page_size,
|
|
support_sliding_window=support_sliding_window,
|
|
num_hidden_layers=32,
|
|
num_attention_heads=32,
|
|
num_key_value_heads=32,
|
|
qk_head_dim=128,
|
|
v_head_dim=128,
|
|
rope_mode=RopeMode.NORMAL,
|
|
rope_scale=1,
|
|
rope_theta=10000,
|
|
rotary_dim=128,
|
|
dtype="float16",
|
|
)
|
|
|
|
export_results = PagedKVCacheTest().export_tvm(
|
|
spec={
|
|
"forward": {
|
|
"cache": spec.Object(object_type=PagedKVCache),
|
|
"qkv": spec.Tensor((1, 100, 96, 128), "float16"),
|
|
},
|
|
"create_paged_kv_cache": {
|
|
"max_batch_size": int,
|
|
"max_total_seq_len": int,
|
|
"prefill_chunk_size": int,
|
|
"page_size": int,
|
|
"support_sliding_window": int,
|
|
},
|
|
},
|
|
)
|
|
tvm_mod = export_results[0]
|
|
tvm.ir.assert_structural_equal(tvm_mod, Module, True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_nn_module_paged_kv_cache()
|