Add support for Relax RMS norm in cutlass offloading

python/tvm/relax/backend/contrib/cutlass.py

@@ -44,6 +44,7 @@
     make_matmul_pattern,
     make_residual_block_pattern,
     make_rms_norm_pattern,
+    make_rms_norm_tir_pattern,
     make_stacked_attention_pattern,
 )
 from ..utils import has_leaking_intermediate_variables
@@ -469,7 +470,7 @@ def layer_norm_pattern():
     ]
 
 
-def _check_rms_norm(ctx: PatternCheckContext) -> bool:
+def _check_rms_norm_tir(ctx: PatternCheckContext) -> bool:
     rms_norm = ctx.annotated_expr["rms_norm"]
     if "rms_norm" not in rms_norm.args[0].name_hint:
         return False
@@ -483,7 +484,17 @@ def rms_norm_pattern():
         (
             "cutlass.rms_norm",
             *make_rms_norm_pattern(),
-            _check_rms_norm,
+        ),
+    ]
+
+
+def rms_norm_tir_pattern():
+    """Create a RMS norm pattern for CUTLASS."""
+    return [
+        (
+            "cutlass.rms_norm_tir",
+            *make_rms_norm_tir_pattern(),
+            _check_rms_norm_tir,
         ),
     ]
 
@@ -512,6 +523,7 @@ def attention_rewrite_patterns():
         *attention_patterns(),
         *layer_norm_pattern(),
         *rms_norm_pattern(),
+        *rms_norm_tir_pattern(),
     ]
 )
 

python/tvm/relax/backend/patterns.py

@@ -326,7 +326,7 @@ def make_layer_norm_pattern():
     return is_op("relax.nn.layer_norm")(inp, gamma, beta), {}
 
 
-def make_rms_norm_pattern():
+def make_rms_norm_tir_pattern():
     """Create a layer norm pattern."""
     inp = wildcard()
     weight = wildcard()
@@ -336,6 +336,14 @@ def make_rms_norm_pattern():
     return out, annotations
 
 
+def make_rms_norm_pattern():
+    """Create a layer norm pattern."""
+    inp = wildcard()
+    weight = wildcard()
+    out = is_op("relax.nn.rms_norm")(inp, weight)
+    return out, {}
+
+
 def make_attention_rewrite_pattern(
     qkv_layout: str, out_layout: str, with_bias: bool, with_cast: bool, with_kv_repeat: bool = False
 ):