Move activation handling to GatedMLP for LoRA compatibility

- Modified _apply_activation method to accept a for_lora flag, allowing for specific handling of activation during LoRA operations.
- Updated the call to _apply_activation in GatedMLP to pass the for_lora argument, ensuring correct behavior in LoRA scenarios.
- Removed unnecessary dtype casting checks in LoraLayer, simplifying the code.

Signed-off-by: Venky Ganesh <[email protected]>

Author: venkywonka

tensorrt_llm/_torch/modules/gated_mlp.py

@@ -5,6 +5,7 @@
 import torch.nn.functional as F
 from torch import nn
 
+from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
 
 from ..distributed import AllReduceParams
@@ -95,12 +96,23 @@ def __init__(self,
             [LoraModuleType.MLP_GATE_UP],
             [2 * self.intermediate_size // mapping.tp_size])
 
-    def _apply_activation(self, x):
+    def _apply_activation(self, x, *, for_lora: bool = False):
         if self.activation == F.silu:
             if self.down_proj.has_fp8_qdq:
-                return swiglu(x,
-                              quant_scale=self.down_proj.input_scale,
-                              quant_type=torch.float8_e4m3fn)
+                if for_lora:
+
+                    target = torch.bfloat16 if torch.cuda.is_bf16_supported(
+                    ) else torch.float16
+                    logger.debug(
+                        f"GatedMLP._apply_activation: LoRA path active; forcing non-FP8 activation dtype {target} (keeping activations in bf16/fp16), layer_idx={self.layer_idx}"
+                    )
+                    return swiglu(x,
+                                  quant_scale=self.down_proj.input_scale,
+                                  quant_type=target)
+                else:
+                    return swiglu(x,
+                                  quant_scale=self.down_proj.input_scale,
+                                  quant_type=torch.float8_e4m3fn)
             else:
                 return swiglu(x)
         elif callable(self.activation):
@@ -152,7 +164,7 @@ def forward_lora(
         if h1_lora is not None:
             h1 = h1 + h1_lora
 
-        h2 = self._apply_activation(h1)
+        h2 = self._apply_activation(h1, for_lora=True)
         output = self.down_proj(h2,
                                 all_reduce_params=final_all_reduce_params,
                                 lora_params=lora_params,

tensorrt_llm/_torch/peft/lora/layer.py

@@ -3,8 +3,6 @@
 
 import torch
 
-from tensorrt_llm._utils import logger
-
 
 class LoraModuleType(IntEnum):
     """Enum class representing different types of modules that can have LoRA adapters.
@@ -121,15 +119,6 @@ def forward(
             if len(active_lora_module_ids) == 0:
                 return None
             else:
-                # Guard: LoRA custom op only supports FP16/BF16 activations.
-                # If upstream produced FP8 (e.g., FP8 SwiGLU), cast here to avoid runtime failure.
-                if x.dtype not in (torch.float16, torch.bfloat16):
-                    target_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported(
-                    ) else torch.float16
-                    logger.debug(
-                        f"lora_grouped_gemm supports only FP16/BF16. Casting input from {x.dtype} to {target_dtype}."
-                    )
-                    x = x.to(target_dtype).contiguous()
                 lora_outputs = torch.ops.trtllm.lora_grouped_gemm(
                     x,
                     lora_params['host_request_types'][:num_seqs],