Use a more robust implementation to keep ActivationType aligned across layers

Signed-off-by: Neta Zmora <96238833+nzmora-nvidia@users.noreply.github.com>

Author: nzmora-nvidia

cpp/tensorrt_llm/kernels/cutlass_kernels/include/common.h

@@ -29,8 +29,8 @@ enum class ActivationType
     Swiglu,
     Geglu,
     SwigluBias,
-    Identity,
     Relu2,
+    Identity,
     InvalidType
 };
 

tensorrt_llm/_torch/utils.py

@@ -40,8 +40,8 @@ class ActivationType(IntEnum):
     Swiglu = 3
     Geglu = 4
     SwigluBias = 5
-    Identity = 6
-    Relu2 = 7
+    Relu2 = 6
+    Identity = 7
     InvalidType = 8
 
 

tensorrt_llm/layers/moe.py

@@ -20,6 +20,7 @@
 import tensorrt as trt
 import torch
 
+from tensorrt_llm._torch.utils import ActivationType
 from tensorrt_llm._utils import (get_init_params, str_dtype_to_torch,
                                  str_dtype_to_trt)
 from tensorrt_llm.layers.lora import LoraParams
@@ -49,14 +50,15 @@
 
 activation_str_to_int_map = {
     # [WARNING] Keep the below in sync with cpp/tensorrt_llm/kernels/cutlass_kernels/include/common.h
-    "gelu": 0,
-    "gelu_new": 0,
-    "relu": 1,
-    "silu": 2,
-    "swiglu": 3,
-    "geglu": 4,
-    "swiglu_bias": 5,
-    "identity": 6,
+    "gelu": int(ActivationType.Gelu),
+    "gelu_new": int(ActivationType.Gelu),
+    "relu": int(ActivationType.Relu),
+    "silu": int(ActivationType.Silu),
+    "swiglu": int(ActivationType.Swiglu),
+    "geglu": int(ActivationType.Geglu),
+    "swiglu_bias": int(ActivationType.SwigluBias),
+    "identity": int(ActivationType.Identity),
+    "relu2": int(ActivationType.Relu2),
 }