[Bugfix][B200] Fix cutlass_mla hang (vllm-project#24966)

Signed-off-by: Alexander Matveev <amatveev@redhat.com>
Co-authored-by: Michael Goin <mgoin64@gmail.com>

Author: 2 people

csrc/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp

@@ -133,6 +133,14 @@ class MLA {
     // printf("    sm_count = %d\n", sm_count);
     int max_splits = ceil_div(K, 128);
     max_splits = min(16, max_splits);
+
+    // TODO: This avoids a hang when the batch size larger than 1 and 
+    // there is more than 4 kv_splits. 
+    // Discuss with NVIDIA how this can be fixed.
+    if (B > 1) {
+      max_splits = min(2, max_splits);
+    }
+    
     // printf("    max_splits = %d\n", max_splits);
     int sms_per_batch = max(1, sm_count / B);
     // printf("    sms_per_batch = %d\n", sms_per_batch);