fix tests

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

Author: LucasWilkinson

tests/v1/attention/test_attention_splitting.py

@@ -294,10 +294,14 @@ def test_prefill_split_across_ubatches(
     qsl_np = common.query_start_loc_cpu.numpy()
     num_tokens = common.num_actual_tokens
 
-    ubatch_slices = maybe_create_ubatch_slices(
-        True, num_scheduled_tokens, num_tokens, batch_spec.batch_size
+    ubatch_slices, _ = maybe_create_ubatch_slices(
+        True,
+        num_scheduled_tokens,
+        num_tokens,
+        batch_spec.batch_size,
+        split_point=split_point,
     )
-    assert len(ubatch_slices) == 2
+    assert ubatch_slices is not None and len(ubatch_slices) == 2
 
     first_meta = _make_metadata_with_slice(ubatch_slices[0], common)
     second_meta = _make_metadata_with_slice(ubatch_slices[1], common)

vllm/v1/worker/ubatch_utils.py

@@ -65,11 +65,13 @@ def maybe_create_ubatch_slices(
     num_scheduled_tokens: np.ndarray,
     num_tokens_padded: int,
     num_reqs_padded: int,
+    split_point: int | None = None,
 ) -> tuple[UBatchSlices | None, UBatchSlices | None]:
     if not should_ubatch:
         return None, None
 
-    split_point = int(num_tokens_padded) // 2
+    if split_point is None:
+        split_point = int(num_tokens_padded) // 2
 
     # TODO(lucas): Refactor the gpu_model_runner.py so we can pass
     # in cu_num_tokens directly (i.e. query_start_loc)