refactor split func

Signed-off-by: ganyi <ygan@amd.com>

Author: ganyi1996ppo

vllm/v1/attention/backends/utils.py

@@ -761,35 +761,24 @@ def split_decodes_prefills_and_extends(
 
     query_lens = query_start_loc[1:] - query_start_loc[:-1]
     is_prefill = query_lens > decode_threshold
-    if not torch.any(is_prefill):
-        return num_reqs, 0, 0, num_tokens, 0, 0
-
+    is_pure_prefill = (seq_lens == query_lens) & is_prefill
     first_prefill = is_prefill.int().argmax(dim=-1).item()
-    assert torch.all(query_lens[first_prefill:] > decode_threshold)
-    assert torch.all(query_lens[:first_prefill] <= decode_threshold)
-
+    first_pure_prefill = is_pure_prefill.int().argmax(dim=-1).item()
     num_decodes = first_prefill
     num_decode_tokens = query_start_loc[first_prefill].item()
-
-    query_lens_prefill = query_lens[first_prefill:]
-    seq_lens_prefill = seq_lens[first_prefill:]
-    is_extend = seq_lens_prefill != query_lens_prefill
-
-    if torch.all(is_extend):
-        num_extends = num_reqs - num_decodes
-        num_extend_tokens = num_tokens - num_decode_tokens
-        return (num_decodes, num_extends, 0, num_decode_tokens, num_extend_tokens, 0)
+    if not torch.any(is_prefill):
+        return (num_decodes, 0, 0, num_decode_tokens, 0, 0)
 
     num_prefills = num_reqs - num_decodes
-    first_extend = is_extend.int().argmax(dim=-1).item()
+    num_prefill_tokens = num_tokens - num_decode_tokens
+    if not torch.any(is_pure_prefill):
+        return (num_decodes, num_prefills, 0, num_decode_tokens, num_prefill_tokens, 0)
 
-    num_extends = first_extend
-    num_pure_prefills = num_prefills - first_extend
+    num_extends = first_pure_prefill - num_decodes
+    num_pure_prefills = num_reqs - first_pure_prefill
 
-    num_extend_tokens = (
-        query_start_loc[num_extends + num_decodes].item() - num_decode_tokens
-    )
-    num_pure_prefill_tokens = num_tokens - num_decode_tokens - num_extend_tokens
+    num_pure_prefill_tokens = num_tokens - query_start_loc[first_pure_prefill]
+    num_extend_tokens = num_prefill_tokens - num_pure_prefill_tokens
     return (
         num_decodes,
         num_extends,
@@ -875,28 +864,6 @@ def reorder_batch_to_split_decodes_and_prefills(
     # NOTE for now we loosely use "decode" to mean requests where attention is
     #  likely memory-bound and "prefill" to mean requests where attention is
     #  likely compute-bound,
-    # rid = dist.get_rank()
-    rid = 0
-
-    def print_order():
-        if rid == 0:
-            num_scheduled_tokens = [
-                scheduler_output.num_scheduled_tokens[id] for id in input_batch.req_ids
-            ]
-            num_scheduled_tokens_np = np.array(num_scheduled_tokens)
-            num_computed_tokens_np = input_batch.num_computed_tokens_cpu[:num_reqs]
-            print("num scheduled tokens: ", num_scheduled_tokens_np, flush=True)
-            print("num computed tokens: ", num_computed_tokens_np, flush=True)
-            is_decode = num_scheduled_tokens_np <= decode_threshold
-            is_extend = (~is_decode) & (num_computed_tokens_np > 0)
-            is_prefill = (~is_decode) & (num_computed_tokens_np == 0)
-            idx = np.arange(0, is_decode.shape[0])
-            decodes = idx[is_decode]
-            extends = idx[is_extend]
-            prefills = idx[is_prefill]
-            print("decode: ", decodes, flush=True)
-            print("extends: ", extends, flush=True)
-            print("prefills: ", prefills, flush=True)
 
     num_reqs = len(input_batch.req_ids)
     num_scheduled_tokens = [