[tvm] branch unity updated: [Unity][CUTLASS] Fixed memory leak in attention kernel offload (#14723)

[yaxingcai]

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

yaxingcai 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 24480ea573 [Unity][CUTLASS] Fixed memory leak in attention kernel 
offload (#14723)
24480ea573 is described below

commit 24480ea573b8f4be81ac65dd5c1cafef9a18a533
Author: masahi <masahi...@gmail.com>
AuthorDate: Wed Apr 26 16:24:24 2023 +0900

    [Unity][CUTLASS] Fixed memory leak in attention kernel offload (#14723)
    
    fixed memory leak in CUTLASS attention offload
---
 python/tvm/contrib/cutlass/attention_operation.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/python/tvm/contrib/cutlass/attention_operation.py 
b/python/tvm/contrib/cutlass/attention_operation.py
index 7c5b7048a2..1921edb397 100644
--- a/python/tvm/contrib/cutlass/attention_operation.py
+++ b/python/tvm/contrib/cutlass/attention_operation.py
@@ -149,6 +149,10 @@ def instantiate_attention_template(attrs):
 
   CHECK(Attention::check_supported(p));
   kernel_fn<<<p.getBlocksGrid(), p.getThreadsGrid(), smem_bytes>>>(p);
+
+  if (Attention::kNeedsOutputAccumulatorBuffer) {
+    cudaFree(p.output_accum_ptr);
+  }
 """
 
     template = substitute_template(