[Bugfix] Fix hanging issue with cutlass_mla by limiting the kv_split for larger batch size

Signed-off-by: Alexander Matveev <amatveev@redhat.com>

Author: alexm-redhat

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

@@ -133,6 +133,13 @@ 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 is large and there is more
+    // than 4 kv_splits. Discuss with NVIDIA how this can be fixed.
+    if (B >= 8) {
+      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);