fix CI

Signed-off-by: Dylan Chen <[email protected]>

Author: DylanChen-NV

cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplJIT/decoderXQAImplJIT.cpp

@@ -124,6 +124,14 @@ bool DecoderXQAImplJIT::shouldUse(XQAParams const& umbrellaXQAParams, bool forCo
         bool hasPerfGain = mayHavePerfGain(xqaParams);
         if (!hasPerfGain)
         {
+            if (xqaParams.kv_cache_data_type == DATA_TYPE_E4M3
+                && (xqaParams.data_type == DATA_TYPE_BF16 || xqaParams.data_type == DATA_TYPE_FP16))
+            {
+                TLLM_LOG_DEBUG(
+                    "JIT XQA is selected in the generation phase for fp16/bf16 input and e4m3 kv cache because MMHA "
+                    "does not support this combination.");
+                return true;
+            }
             TLLM_LOG_DEBUG("JIT XQA is not used: maybe no performance gain");
             return false;
         }

cpp/tensorrt_llm/kernels/xqaDispatcher.cpp

@@ -244,15 +244,6 @@ bool XqaDispatcher::shouldUse(XQAParams const& params)
         return true;
     }
 
-    if (params.kv_cache_data_type == DATA_TYPE_E4M3
-        && (params.data_type == DATA_TYPE_BF16 || params.data_type == DATA_TYPE_FP16))
-    {
-        TLLM_LOG_DEBUG(
-            "XQA kernels are selected in the generation phase for fp16/bf16 input and e4m3 kv cache because MMHA does "
-            "not support this combination.");
-        return true;
-    }
-
     return mDecoderXqaRunner->shouldUse(params, /*forConfigurePlugin=*/false);
 }