Modify comments

NVIDIA · Aug 23, 2021 · b0b017f · b0b017f
1 parent c6b7763
commit b0b017f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/megatron/model/fused_softmax.py b/megatron/model/fused_softmax.py
@@ -138,8 +138,8 @@ def is_kernel_available(self, mask, b, np, sq, sk):
             self.scaled_masked_softmax_fusion  # user want to fuse
             and self.input_in_float16  # input must be fp16
             and mask is not None  # mask tensor must not be None
-            and 16 < sk <= 2048  # sq must be 16 ~ 2048
-            and sq % 4 == 0  # sk must be divisor of 4
+            and 16 < sk <= 2048  # sk must be 16 ~ 2048
+            and sq % 4 == 0  # sq must be divisor of 4
             and attn_batches % 4 == 0  # np * b must be divisor of 4
         ):
             if 0 <= sk <= 2048: