Update transformer.py

megatron/legacy/model/transformer.py

@@ -582,13 +582,13 @@ class ParallelAttention(MegatronModule):
         else:
             kv_projection_size = args.kv_channels * args.num_attention_heads
 
-        self.use_flash_attn = (args.use_flash_attn_ck or args.use_flash_attn_triton) \
+        self.use_flash_attn = (args.use_flash_attn_cutlass or args.use_flash_attn_triton) \
             and attention_type == AttnType.self_attn \
             and self.attn_mask_type == AttnMaskType.causal
         self.use_flash_attn_triton = args.use_flash_attn_triton
 
         if self.use_flash_attn:
-            if args.use_flash_attn_ck:
+            if args.use_flash_attn_cutlass:
                 if flash_attn_unpadded_func is None:
                     raise ImportError('FlashAttention is not installed, please install with '
                                     'pip install flash-attn')