MemEff: BUGFIX when a batch has no keys

ghstack-source-id: 84e11de209464a6a51347f9ba6b08258782bf139
Pull Request resolved: #492

tests/test_mem_eff_attention.py

@@ -346,10 +346,11 @@ def test_cu_seqlen_forward(
     scale = 3
     # Reduce batch size to speedup tests
     batch_size = min(batch_size, 20)
+    q_lens = [r.randint(1, max_q_len) for b in range(batch_size)] + [0, 0, 16]
+    kv_lens = [r.randint(1, max_kv_len) for b in range(batch_size)] + [0, 16, 0]
 
-    for batch_id in range(batch_size):
-        q_len = r.randint(1, max_q_len)
-        kv_len = r.randint(1, max_kv_len)
+    for batch_id in range(len(q_lens)):
+        q_len, kv_len = q_lens[batch_id], kv_lens[batch_id]
 
         all_q.append(
             torch.randn((1, q_len, num_heads, k), device=device, dtype=dtype) * scale

xformers/components/attention/csrc/cuda/mem_eff_attention/kernel_forward.h

@@ -801,6 +801,24 @@ struct AttentionKernel {
             cutlass::platform::numeric_limits<accum_t>::infinity();
       }
     }
+
+    // Special case when we have no keys BUT num_queries>0
+    // we still need to fill the output with zeros
+    if (__builtin_expect(p.num_keys <= 0, 0)) {
+      using Epilogue = typename MM1::DefaultEpilogue;
+      using OutputTileIterator = typename MM1::OutputTileIterator;
+      typename OutputTileIterator::Fragment zero;
+      zero.clear();
+      CUTLASS_PRAGMA_UNROLL
+      for (int col = 0; col < p.head_dim_value; col += MM1::Mma::Shape::kN) {
+        auto output_it = createOutputIter(col);
+        CUTLASS_PRAGMA_UNROLL
+        for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) {
+          output_it.store(zero);
+          ++output_it;
+        }
+      }
+    }
   }
 
   static CUTLASS_DEVICE int8_t lane_id() {