fix black code stype

Author: Zzz9990

aiter/mla.py

@@ -196,9 +196,15 @@ def mla_decode_fwd(
             MAYBE_FINAL_OUT = False
 
         logits = (
-            o.view((total_s, num_kv_splits, nhead, v_head_dim)) if (num_kv_splits == 1
-                and (q.dtype == dtypes.fp8 or (q.dtype == dtypes.bf16 and max_seqlen_q == 4))) else
-            torch.empty(
+            o.view((total_s, num_kv_splits, nhead, v_head_dim))
+            if (
+                num_kv_splits == 1
+                and (
+                    q.dtype == dtypes.fp8
+                    or (q.dtype == dtypes.bf16 and max_seqlen_q == 4)
+                )
+            )
+            else torch.empty(
                 (total_s, num_kv_splits, nhead, v_head_dim),
                 dtype=dtypes.fp32,
                 device=device,
@@ -230,7 +236,9 @@ def mla_decode_fwd(
             kv_scale,
         )
 
-        if (num_kv_splits == 1 and (q.dtype == dtypes.fp8 or (q.dtype == dtypes.bf16 and max_seqlen_q == 4))):
+        if num_kv_splits == 1 and (
+            q.dtype == dtypes.fp8 or (q.dtype == dtypes.bf16 and max_seqlen_q == 4)
+        ):
             return logits.view(total_s, nhead, v_head_dim), attn_lse
 
         Lv = v_head_dim

op_tests/test_mla_persistent.py

@@ -277,7 +277,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,
+        max_split_per_batch=max_split_per_batch,
     )
 
     def test_absorb_decode_bf16():