disable skip for short seqs

Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com>

Author: pengbowang-nv

cpp/kernels/xqa/mha_sm90.cu

@@ -786,7 +786,8 @@ CUBIN_EXPORT __global__
     static_assert(multiBlockMinNbTiles >= multiBlockMinNbTilesPerCta * 2);
     assert(isMultiBlockMode == (nbSubSeq > 1));
 #if SKIP_SOFTMAX_ATTN
-    float const skipSoftmaxThreshold = skipSoftmaxThresholdScaleFactor / cacheSeqLen;
+    bool const disableSkipForShortSeq = (cacheSeqLen < skipSoftmaxThresholdScaleFactor);
+    float const skipSoftmaxThreshold = disableSkipForShortSeq ? 0.0f : skipSoftmaxThresholdScaleFactor / cacheSeqLen;
 #endif
     if (idxSubSeq >= nbSubSeq)
     {
@@ -1001,7 +1002,7 @@ CUBIN_EXPORT __global__
             auto& skipSoftmaxXBar = smem.skipSoftmaxXBar[idxXBuf];
             skipSoftmaxXBar.consumed.arrive_and_wait();
 
-            bool const maybeSkip = idxIter != 0;
+            bool const maybeSkip = !disableSkipForShortSeq && idxIter != 0;
             RegColWiseVec const colMax = computeWarpGrpColMax_sync(smem.gemm0WarpGrpBar, smem.gemm0CurrentSeqMax, acc,
                 skipSoftmaxThreshold, &smem.skipSoftmaxVotesGemm0ToV[idxXBuf], maybeSkip);
             bool const shouldSkipSoftmaxAttn = static_cast<bool>(smem.skipSoftmaxVotesGemm0ToV[idxXBuf]);
@@ -2142,7 +2143,8 @@ __device__ inline RegColWiseVec computeWarpGrpColMax_sync(
     {
         *smemSkipVote = maybeSkip ? 1U : 0U; // will sync before vote
     }
-    float const lnThreshold = log(skipSoftmaxThreshold);
+    float const lnThreshold
+        = log(skipSoftmaxThreshold); // this can be -inf, but should be safe as we only use it for comparison
 #endif
 
     auto colMax = RegColWiseVec::filled(Vec<float, 2>::filled(safeInitRowMax));

cpp/kernels/xqa/test/refAttention.cpp

@@ -69,7 +69,8 @@ Eigen::Matrix<float, headGrpSize, validElemsPerHead, Eigen::RowMajor> refFlashAt
     {
         skipRowMaxs[i].fill(-INFINITY);
     }
-    float skipSoftmaxThreshold = skipSoftmaxThresholdScaleFactor / seqLen;
+    bool const disableSkipForShortSeq = (seqLen < skipSoftmaxThresholdScaleFactor);
+    float const skipSoftmaxThreshold = disableSkipForShortSeq ? 0.0f : skipSoftmaxThresholdScaleFactor / seqLen;
 
     for (uint32_t idxTile = idxTileBeg; idxTile < nbTiles; idxTile++)
     {
@@ -103,7 +104,7 @@ Eigen::Matrix<float, headGrpSize, validElemsPerHead, Eigen::RowMajor> refFlashAt
         auto const prevSkipRowMax = skipRowMaxs[idxTile % nbSubSeq];
         skipRowMaxs[idxTile % nbSubSeq] = localRowMax.cwiseMax(skipRowMaxs[idxTile % nbSubSeq]).eval();
 
-        if (skipSoftmaxThreshold > 0)
+        if (!disableSkipForShortSeq && skipSoftmaxThreshold > 0)
         {
             *totalBlockCount += 1;
             auto const skipSoftmaxMask = ((localRowMax - prevSkipRowMax).array() < std::log(skipSoftmaxThreshold));

cpp/kernels/xqa/test/test.cpp

@@ -1312,10 +1312,11 @@ TEST(RefCheck, llama_V2_70b)
 #endif
 #if SKIP_SOFTMAX_ATTN
     runTest<1>(32, 2048, false, true, false, false, false, ~0U, 1U << 30, 0.f);
-    runTest<4>(32, 1538, false, true, false, false, false, ~0U, 1U << 30, 55.f);
+    runTest<4>(32, 1538, false, true, false, false, false, ~0U, 1U << 30, 1280.f);
     runTest<2>(32, 4096, false, true, false, false, false, ~0U, 1U << 30, 125.f);
     runTest<4>(32, 300, false, true, false, false, false, ~0U, 1U << 30, 80.f);
-    runTest<4>(32, 500, false, true, false, false, false, ~0U, 1U << 30, 455.f);
+    runTest<4>(32, 500, false, true, false, false, false, ~0U, 1U << 30, 501.0f);
+    runTest<4>(32, 500, false, true, false, false, false, ~0U, 1U << 30, 500.f);
 #endif
     runTest<8>(120, 367, false, true);
     runTest<8>(1792, 2048, false, true);