[tvm] branch unity updated: [Unity] Allow KVCache Access without Shape (#14726)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/unity by this push:
     new 9bff70ab06 [Unity] Allow KVCache Access without Shape (#14726)
9bff70ab06 is described below

commit 9bff70ab063c9ea75afa1175d0c96887047bc8c9
Author: Junru Shao <junrus...@apache.org>
AuthorDate: Wed Apr 26 13:03:52 2023 -0700

    [Unity] Allow KVCache Access without Shape (#14726)
    
    This PR makes the second argument of `vm.builtin.attention_kv_cache_view`
    optional, making it easier when debugging without the model config (i.e.
    when `num_attention_heads` and `head_dim` are unknown).
---
 src/runtime/relax_vm/lm_support.cc | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/src/runtime/relax_vm/lm_support.cc 
b/src/runtime/relax_vm/lm_support.cc
index 9470b86e1a..8b867cc602 100644
--- a/src/runtime/relax_vm/lm_support.cc
+++ b/src/runtime/relax_vm/lm_support.cc
@@ -165,7 +165,24 @@ NDArray AttentionKVCacheView(AttentionKVCache cache, 
ShapeTuple shape) {
   return cache->View(shape);
 }
 
-TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_view").set_body_typed(AttentionKVCacheView);
+TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_view")
+    .set_body([](TVMArgs args, TVMRetValue* rv) {
+      CHECK(args.size() == 1 || args.size() == 2)
+          << "ValueError: `vm.builtin.attention_kv_cache_view` expects 1 or 2 
arguments, but got "
+          << args.size() << ".";
+      AttentionKVCache cache = args[0];
+      if (args.size() == 2) {
+        ShapeTuple shape = args[1];
+        *rv = cache->View(shape);
+      } else {
+        std::vector<ShapeTuple::index_type> shape;
+        shape.push_back(cache->fill_count);
+        for (int i = 1; i < cache->data->ndim; ++i) {
+          shape.push_back(cache->data->shape[i]);
+        }
+        *rv = cache->View(ShapeTuple(shape));
+      }
+    });
 
 void AttentionKVCacheArrayClear(Array<AttentionKVCache> caches) {
   for (AttentionKVCache cache : caches) {