add max split num to sparse mla (ROCm#1617)

Author: Zzz9990

op_tests/test_mla_sparse.py

@@ -326,6 +326,7 @@ def test_mla(
     page_size,
     varlen,
     decode_qlen,
+    max_split_per_batch,
 ):
     ret = {}
 
@@ -412,6 +413,7 @@ def test_mla(
         kvtype,
         is_sparse=True,
         fast_mode=True,
+        num_kv_splits=max_split_per_batch,
     )
 
     # aiter implementation
@@ -451,6 +453,7 @@ def test_mla(
         max_seqlen_qo=int(max_seqlen_qo),
         uni_seqlen_qo=decode_qlen,
         fast_mode=True,
+        max_split_per_batch=max_split_per_batch,
         topk=2048,
         dtype_q=dtype,
         dtype_kv=kvtype,
@@ -501,6 +504,7 @@ def test_sparse_mla_bf16():
             kv_last_page_lens,
             1,
             sm_scale,
+            num_kv_splits=max_split_per_batch,
             work_meta_data=work_meta_data,
             work_indptr=work_indptr,
             work_info_set=work_info_set,
@@ -561,6 +565,7 @@ def test_absorb_decode_fp8():
             kv_last_page_lens,
             1,
             sm_scale,
+            num_kv_splits=max_split_per_batch,
             q_scale=q_scale,
             kv_scale=kv_scale,
             work_meta_data=work_meta_data,
@@ -721,6 +726,15 @@ def test_absorb_decode_fp8():
     help="""Number of heads.
     e.g.: -n 16,1""",
 )
+parser.add_argument(
+    "-ms",
+    "--max_split_per_batch",
+    type=int,
+    nargs="*",
+    default=[16],
+    help="""kv seqlens max split num for per batch.
+    e.g.: -ms 32""",
+)
 parser.add_argument(
     "--varlen",
     action="store_true",
@@ -738,8 +752,8 @@ def test_absorb_decode_fp8():
 
 for nhead, decode_qlen in list_nhead:
     df = []
-    for dtype, kvtype, ctx_len, batch_size in itertools.product(
-        list_dtype, l_kv_dtype, args.ctxLen, args.batchSize
+    for dtype, kvtype, ctx_len, batch_size, max_split_per_batch in itertools.product(
+        list_dtype, l_kv_dtype, args.ctxLen, args.batchSize, args.max_split_per_batch
     ):
         ret = test_mla(
             ctx_len,
@@ -754,6 +768,7 @@ def test_absorb_decode_fp8():
             args.block_size,
             varlen=args.varlen,
             decode_qlen=decode_qlen,
+            max_split_per_batch=max_split_per_batch,
         )
         df.append(ret)
     df = pd.DataFrame(df)