fix unit test

Signed-off-by: Lucas Wilkinson <[email protected]>

Author: LucasWilkinson

tests/v1/attention/test_attention_splitting.py

@@ -13,7 +13,7 @@
     split_attn_metadata,
     split_decodes_and_prefills,
 )
-from vllm.v1.worker.ubatch_utils import create_ubatch_slices
+from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 
 
 @pytest.fixture
@@ -294,7 +294,9 @@ def test_prefill_split_across_ubatches(
     qsl_np = common.query_start_loc_cpu.numpy()
     num_tokens = common.num_actual_tokens
 
-    ubatch_slices = create_ubatch_slices(num_scheduled_tokens, split_point)
+    ubatch_slices = maybe_create_ubatch_slices(
+        True, num_scheduled_tokens, num_tokens, batch_spec.batch_size
+    )
     assert len(ubatch_slices) == 2
 
     first_meta = _make_metadata_with_slice(ubatch_slices[0], common)