[Unity] Cutlass attention with dynamic sequence length

[yelite]

This PR adds support for the dynamic sequence length in cutlass attention.

cc @cyx-6 @masahi

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @quic-sanirudh _{See #10317 for details}

_{Generated by tvm-bot}

[masahi]

How hard would it to be to support workspace alloc for dynamic attention?

[yelite]

The approach in my mind is to make the max_workspace_size_ of WorkspaceProvider a PrimExpr isntead of IntImm, which consists of a series of max expr to get the actual max workspace size from symbolic shape variables.

I tried to doing this in the most straightforward way and got This IR is not well formed: Symbolic Var a presents in different functions in the same Module.. Apparently it needs more sophisticated analysis to replace the shape variables with valid ones as the workspace size is propagated from composite function to its caller. Also I am not sure if the arg of alloc_tensor can be a non-constant expr too.

[sunggg]

Just for curiosity, what is the rationale behind workspace_size_bytes = 8?

The approach in my mind is to make the max_workspace_size_ of WorkspaceProvider a PrimExpr isntead of IntImm, which consists of a series of max expr to get the actual max workspace size from symbolic shape variables.

Would you elaborate a little more? Did you try to do more sophisticated approach than using the annotation like R.func_attr({"tir_var_upper_bound": {"m": 2048, "n": 2048}})?

I am not sure if the arg of alloc_tensor can be a non-constant expr too.

aloc_tensor can take symbolic shape. See https://github.com/apache/tvm/blob/unity/tests/python/relax/test_transform_static_plan_block_memory.py#L958

[masahi]

workspace_size_bytes = 8

It is just for a temp workaround, it will cause cudaMalloc to be always called.