[Relax] Handle dynamic arguments in legalization of nn.attention (#16592

)

Prior to this commit, when using causal_mask="BottomRight" in `R.nn.attention`,
the legalization would assume that the query and key/value sequence lengths
were static integers. This commit updates the legalization to allow dynamic shapes.

python/tvm/relax/transform/legalize_ops/nn.py

@@ -486,7 +486,7 @@ def _te_attention(
         if causal_mask == "TopLeft":
             offset = tir.IntImm("int32", 0)
         elif causal_mask == "BottomRight":
-            offset = tir.IntImm("int32", abs(seq_len - seq_len_kv))
+            offset = tir.abs(seq_len - seq_len_kv).astype("int32")
         else:
             raise NotImplementedError()
         p_masked = topi.trilu(p, k=offset, upper=False)

tests/python/relax/test_transform_legalize_ops_nn.py

@@ -3270,6 +3270,30 @@ def main(q: R.Tensor((4, 16, 32, 8), dtype="float32"), k: R.Tensor((4, 8, 32, 8)
     tvm.ir.assert_structural_equal(mod, Expected)
 
 
+def test_dynamic_attention():
+    """The sequence lengths may be dynamic
+
+    In previous implementations, the `seq_len` and `seq_len_kv` were
+    assumed to be static integers, and produced an exception during
+    legalization.
+    """
+
+    @tvm.script.ir_module
+    class Attention:
+        @R.function
+        def main(
+            q: R.Tensor((4, "seq_len", 32, 8), "float32"),
+            k: R.Tensor((4, "seq_len_kv", 32, 8), "float32"),
+            v: R.Tensor((4, "seq_len_kv", 32, 16), "float32"),
+            bias: R.Tensor((4, 32, "seq_len", "seq_len_kv"), "float32"),
+        ):
+            scale = T.FloatImm("float32", 0.1)
+            gv = R.nn.attention(q, k, v, bias, scale=scale, causal_mask="BottomRight")
+            return gv
+
+    LegalizeOps()(Attention)
+
+
 def test_nll_loss():
     # fmt: off
     @tvm.script.ir_module