reduce split kv amount

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

vllm/attention/backends/triton_mla.py

@@ -622,8 +622,6 @@ def build(self, seq_lens: List[int], query_lens: List[int],
             self.multimodal_placeholder_maps.items()
         }
 
-        num_kv_splits = 8
-
         return TritonMLAMetadata(
             num_prefills=self.num_prefills,
             slot_mapping=slot_mapping_tensor,
@@ -643,7 +641,7 @@ def build(self, seq_lens: List[int], query_lens: List[int],
             context_lens_tensor=context_lens_tensor,
             block_tables=block_tables,
             use_cuda_graph=use_captured_graph,
-            num_kv_splits=num_kv_splits,
+            num_kv_splits=4,  # TODO(lucas) add heuristic
             head_dim=self.runner.model_config.get_head_size(),
         )