Fix issues (ROCm#1606)

Author: ruanjm

aiter/ops/attention.py

@@ -636,7 +636,7 @@ def get_mla_metadata_info_v1(
     tile_cnt = batch_size * max_qo_tiles_per_batch
 
     if fast_mode:
-        max_work = tile_cnt + cu_num - 1
+        max_work = (batch_size + cu_num - 1) * max_qo_tiles_per_batch
         max_split_tiles = (
             min(batch_size + cu_num - 1, (cu_num - 1) * 2) * max_qo_tiles_per_batch
         )

op_tests/test_mla_persistent.py

@@ -244,8 +244,8 @@ def test_mla(
         batch_size,
         max_seqlen_qo,
         nhead,
-        q.dtype,
-        kv_buffer.dtype,
+        dtype,
+        kvtype,
         is_sparse=False,
         fast_mode=True if not non_persistent_mode else False,
         num_kv_splits=max_split_per_batch,
@@ -292,9 +292,7 @@ def test_mla(
         max_split_per_batch=max_split_per_batch,
         intra_batch_mode=non_persistent_mode,
         dtype_q=dtype,
-        dtype_kv=(
-            kvtype if dtype == kvtype else dtype
-        ),  # if q bf16 k fp8 should be same as bf16bf16 for dp mode
+        dtype_kv=kvtype,
     )
 
     def test_absorb_decode_bf16():

op_tests/test_mla_sparse.py

@@ -408,8 +408,8 @@ def test_mla(
         batch_size,
         max_seqlen_qo,
         nhead,
-        q.dtype,
-        kv_buffer.dtype,
+        dtype,
+        kvtype,
         is_sparse=True,
         fast_mode=True,
     )
@@ -452,8 +452,8 @@ def test_mla(
         uni_seqlen_qo=decode_qlen,
         fast_mode=True,
         topk=2048,
-        dtype_q=q.dtype,
-        dtype_kv=kv_buffer.dtype,
+        dtype_q=dtype,
+        dtype_kv=kvtype,
     )
 
     # generate kv topk per token & convert indices into per token