use VLLM_USE_TRITON_CAT during the prefill phase

vllm/v1/attention/backends/mla/common.py

@@ -982,7 +982,10 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
         k_nope, v = kv_nope\
             .split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
 
-        k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
+        if envs.VLLM_USE_TRITON_CAT:
+            k = concat_helper((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
+        else:
+            k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
 
         output = self._flash_attn_varlen_diff_headdims(
             q=q,