   params.inv_sqrt_dh = inv_sqrt_dh;
 
   switch (dim_head) {
+    case 10:
+      fmha_launch_kernel<T, 10, 32>(params, dev_ctx.stream());
+      break;
+    case 26:
+      fmha_launch_kernel<T, 26, 32>(params, dev_ctx.stream());
+      break;
     case 32:
       fmha_launch_kernel<T, 32, 32>(params, dev_ctx.stream());
       break;

    case 128:
      fmha_launch_kernel<T, 128, 128>(params, dev_ctx.stream());
      break;
    default:
      PADDLE_THROW(platform::errors::Unimplemented(
          "dim_head = %d is unsupport, only support "
          "dim_head = 32, 64, 96 or 128 for now.",
          dim_head));
  }
}


template <typename T>
inline size_t smem_size_in_bytes(
    const Masked_multihead_attention_params<T> &params, int dim_head,
    int threads_per_value, int threads_per_block) {
  size_t qk_sz = div_up(params.timestep + 1, 4) * 16;
  size_t logits_sz = 0;

