You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "Hzfengsy (via GitHub)" <gi...@apache.org> on 2023/12/28 12:38:49 UTC
[PR] [Unity][Frontend] Introduce StateCache for State Space modules [tvm]
Hzfengsy opened a new pull request, #16290:
URL: https://github.com/apache/tvm/pull/16290
We introduce the StateCache to support State Space modules. Unlike KVCache, StateCache does not need the operation `append` but requires the init value.
cc @junrushao @tqchen
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Frontend] Introduce StructuredState for State Space modules [tvm]
Posted by "Hzfengsy (via GitHub)" <gi...@apache.org>.
Hzfengsy commented on PR #16290:
URL: https://github.com/apache/tvm/pull/16290#issuecomment-1951889406
Migrate to https://github.com/apache/tvm/pull/16568
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Frontend] Introduce StateCache for State Space modules [tvm]
Posted by "Hzfengsy (via GitHub)" <gi...@apache.org>.
Hzfengsy commented on code in PR #16290:
URL: https://github.com/apache/tvm/pull/16290#discussion_r1437614883
##########
tests/python/relax/test_frontend_nn_modules.py:
##########
@@ -461,6 +461,61 @@ def forward(self, x: core.Tensor) -> core.Tensor:
assert_structural_equal(tvm_mod, Module, True)
+def test_state_cache():
+ @I.ir_module
+ class Module:
+ @R.function
+ def _initialize_effect() -> R.Tuple(R.Object, R.Object):
+ with R.dataflow():
+ _io: R.Object = R.null_value()
+ lv: R.Tensor((1, 8), dtype="float32") = R.zeros(R.shape([1, 8]), dtype="float32")
+ cache: R.Object = R.call_packed(
+ "vm.builtin.attention_kv_cache_create",
+ lv,
+ R.shape([1, 8]),
+ R.prim_value(1),
+ sinfo_args=(R.Object,),
+ )
+ lv1: R.Tuple(R.Object, R.Object) = _io, cache
+ gv: R.Tuple(R.Object, R.Object) = lv1
+ R.output(gv)
+ return gv
+
+ @R.function
+ def forward(
+ x: R.Tensor((1, 8), dtype="float32"), _io: R.Object, cache: R.Object
+ ) -> R.Tuple(R.Tensor((1, 8), dtype="float32"), R.Tuple(R.Object, R.Object)):
+ R.func_attr({"num_input": 3})
+ with R.dataflow():
+ lv2: R.Object = R.call_packed(
+ "vm.builtin.attention_kv_cache_update", cache, x, sinfo_args=(R.Object,)
+ )
+ lv3: R.Tensor((1, 8), dtype="float32") = R.call_packed(
+ "vm.builtin.attention_kv_cache_view",
+ lv2,
+ R.shape([1, 8]),
+ sinfo_args=(R.Tensor((1, 8), dtype="float32"),),
+ )
+ gv1: R.Tuple(
+ R.Tensor((1, 8), dtype="float32"), R.Tuple(R.Object, R.Object)
+ ) = lv3, (_io, lv2)
+ R.output(gv1)
+ return gv1
+
+ class StateCacheTest(modules.Module):
+ def __init__(self) -> None:
+ self.cache = modules.StateCache(relax.op.zeros((1, 8), "float32"))
Review Comment:
I use `relax.op.zeros` here instead of `nn.op.zeros` here, since there is no blockbuilder when we initialize the Module. it indeed a little bit tricky to involve raw relax expr here, love to see if there is a better solution
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Frontend] Introduce StructuredState for State Space modules [tvm]
Posted by "Hzfengsy (via GitHub)" <gi...@apache.org>.
Hzfengsy closed pull request #16290: [Unity][Frontend] Introduce StructuredState for State Space modules
URL: https://github.com/apache/tvm/pull/16290
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org