fix: enhance attention mask type handling in TllmGenFmhaRunnerParams

Updated the setAttentionMaskType method to include a switch-case structure for better handling of attention mask types, ensuring proper mapping and error handling for invalid types.

Signed-off-by: Qixiang Lin <[email protected]>

Author: qixiang-99

Diff for: cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h

@@ -202,6 +202,7 @@ struct MHARunnerFixedParams
         case 2: // tensorrt_llm::kernels::AttentionMaskType::SLIDING_WINDOW_CAUSAL
             attentionMaskType = ContextAttentionMaskType::SLIDING_WINDOW_CAUSAL;
             break;
+        // NOTE: For BIDIRECTIONAL, BIDIRECTIONALGLM, BLOCKSPARSE context phase, CAUSAL mask is used
         case 3: // tensorrt_llm::kernels::AttentionMaskType::BIDIRECTIONAL
             attentionMaskType = ContextAttentionMaskType::CAUSAL;
             break;

Diff for: cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h

@@ -238,8 +238,26 @@ struct TllmGenFmhaRunnerParams
     // set the attention mask type
     TllmGenFmhaRunnerParams& setAttentionMaskType(std::int8_t maskType)
     {
-        TLLM_CHECK_WITH_INFO(maskType >= 0 && maskType <= 3, "Invalid mask type for TrtllmGenAttentionMaskType");
-        mMaskType = static_cast<TrtllmGenAttentionMaskType>(maskType);
+        // maskType is the enum of tensorrt_llm::kernels::ContextAttentionMaskType
+        // convert ContextAttentionMaskType to TrtllmGenAttentionMaskType
+        switch (maskType)
+        {
+        case 0: // tensorrt_llm::kernels::ContextAttentionMaskType::PADDING
+            mMaskType = TrtllmGenAttentionMaskType::Dense;
+            break;
+        case 1: // tensorrt_llm::kernels::ContextAttentionMaskType::CAUSAL
+            mMaskType = TrtllmGenAttentionMaskType::Causal;
+            break;
+        case 2: // tensorrt_llm::kernels::ContextAttentionMaskType::SLIDING_WINDOW_CAUSAL
+            mMaskType = TrtllmGenAttentionMaskType::SlidingWindowCausal;
+            break;
+        case 3: // tensorrt_llm::kernels::ContextAttentionMaskType::CUSTOM_MASK
+            mMaskType = TrtllmGenAttentionMaskType::Custom;
+            break;
+        default:
+            TLLM_THROW("ContextAttentionMaskType %d cannot be mapped to TrtllmGenAttentionMaskType",
+                static_cast<int>(maskType));
+        }
         return *this;
     }
 };