Merge pull request #122 from foundation-model-stack/chunked_prefill

Chunked prefill support for paged attention

Author: JRosenkranz

aiu_fms_testing_utils/scripts/drive_paged_programs.py

@@ -128,6 +128,12 @@
     choices=["paged", "paged_fp8"],
     help="The attention type to use",
 )
+parser.add_argument(
+    "--prefill_chunk_size",
+    type=int,
+    default=0,
+    help="if > 0, activate chunked prefill, with chunk_size=this_argument. Only works with paged attention variants.",
+)
 parser.add_argument(
     "--stagger_load",
     type=int,
@@ -397,6 +403,7 @@ def __load_validation_info(
     max_new_tokens=max_new_tokens,
     compile_dynamic_sendnn=True,
     stagger_update_lazyhandle=args.stagger_update_lazyhandle,
+    prefill_chunk_size=args.prefill_chunk_size,
     **extra_kwargs,
 )
 
@@ -682,6 +689,7 @@ def __metric_calculator(r: torch.Tensor, t: torch.Tensor):
                 GoldenTokenHook(cpu_validation_info.get_info("tokens")),
                 last_n_tokens=64,
                 timing=TIMING,
+                prefill_chunk_size=args.prefill_chunk_size,
                 **extra_kwargs,
             )
 
@@ -726,6 +734,7 @@ def __metric_calculator(r: torch.Tensor, t: torch.Tensor):
                 None,
                 last_n_tokens=64,
                 timing=TIMING,
+                prefill_chunk_size=args.prefill_chunk_size,
                 **extra_kwargs,
             )
 
@@ -768,6 +777,7 @@ def __metric_calculator(r: torch.Tensor, t: torch.Tensor):
             None,
             last_n_tokens=64,
             timing=TIMING,
+            prefill_chunk_size=args.prefill_chunk_size,
             **extra_kwargs,
         )
 

aiu_fms_testing_utils/scripts/inference.py

@@ -239,6 +239,12 @@
     default="sdpa",
     help="which backend attention to use in mha",
 )
+parser.add_argument(
+    "--prefill_chunk_size",
+    type=int,
+    default=0,
+    help="if > 0, activate chunked prefill, with chunk_size=this_argument. Only works with paged attention variants",
+)
 parser.add_argument(
     "--stagger_load",
     type=int,
@@ -782,6 +788,8 @@ def infer(use_cache, do_sample, warmup):
     if attn_name == "sdpa_causal":
         attention_specific_kwargs["contiguous_cache"] = True
         attention_specific_kwargs["max_seq_len"] = ids.shape[1] + args.max_new_tokens
+    elif "paged" in attn_name:
+        attention_specific_kwargs["prefill_chunk_size"] = args.prefill_chunk_size
 
     result = generate(
         model,
@@ -844,6 +852,7 @@ def infer(use_cache, do_sample, warmup):
                 args.compile_dynamic_sendnn,
                 use_cache=cache,
                 stagger_update_lazyhandle=args.stagger_update_lazyhandle,
+                prefill_chunk_size=args.prefill_chunk_size,
                 **extra_generation_kwargs,
             )
         if (

aiu_fms_testing_utils/testing/validation.py

@@ -263,11 +263,14 @@ def extract_validation_information(
     eos_token_id=None,
     last_n_tokens=0,
     timing="",
+    prefill_chunk_size=0,
     **extra_kwargs,
 ):
     attention_specific_kwargs = {}
     if "paged" in extra_kwargs.get("attn_name", "sdpa"):
         from aiu_fms_testing_utils.utils.paged import generate
+
+        attention_specific_kwargs["prefill_chunk_size"] = prefill_chunk_size
     else:
         # TODO: Add a unified generation dependent on attn_type
         from fms.utils.generation import generate

aiu_fms_testing_utils/utils/__init__.py

@@ -54,6 +54,7 @@ def warmup_model(
     compile_dynamic_sendnn: bool = False,
     use_cache: bool = True,
     stagger_update_lazyhandle: int = 0,
+    prefill_chunk_size: int = 0,
     **extra_kwargs,
 ):
     import torch_sendnn
@@ -62,6 +63,8 @@ def warmup_model(
     attn_name = extra_kwargs.get("attn_name", "sdpa")
     if "paged" in attn_name:
         from aiu_fms_testing_utils.utils.paged import generate, adjust_inputs_to_batch
+
+        attention_specific_kwargs["prefill_chunk_size"] = prefill_chunk_size
     else:
         # TODO: Add a unified generation dependent on attn_type
         from fms.utils.generation import generate

aiu_fms_testing_utils/utils/paged.py

@@ -1,3 +1,4 @@
+import math
 import os
 import random
 import time
@@ -35,6 +36,7 @@ def generate(
     do_sample: bool = True,
     num_beams: int = 1,
     use_cache: bool = False,
+    prefill_chunk_size: int = 0,
     eos_token_id: Optional[int] = None,
     timing: str = "",
     post_iteration_hook: Optional[
@@ -117,6 +119,10 @@ def generate(
     max_possible_context_length = input_ids.size(1) + max_new_tokens
 
     BLOCK_SIZE = 64
+    if prefill_chunk_size > 0:
+        assert prefill_chunk_size % BLOCK_SIZE == 0, (
+            "Chunk size must be a multiple of the page size"
+        )
 
     # these variables are guaranteed to be set in another location (inference.py, test_decoders.py, etc.)
     # if we set these variables here, we run the risk of warming up and generating with different sizes
@@ -229,7 +235,6 @@ def generate(
         for pos_i in range(input_ids.shape[1] - seq_tkv, input_ids.shape[1]):
             # we may have already popped a block, so index to the proper block
             block_number = block_table_i[pos_i // BLOCK_SIZE]
-
             block_offset = pos_i % BLOCK_SIZE
             slot = block_number * BLOCK_SIZE + block_offset
             slot_mapping_i.append(slot)
@@ -264,57 +269,149 @@ def generate(
                 # remove extra pads from the input_ids, slot_mapping, position_ids, mask to account for empty pages
                 # each input should be padded to its smallest multiple of BLOCK_SIZE (64)
                 # we need to clone these tensors to ensure the pointer offset is 0
-                input_ids_i = input_ids[seq_i][-current_tkv:].unsqueeze(0).clone()
-                slot_mapping_i = (
+                input_ids_seq = input_ids[seq_i][-current_tkv:].unsqueeze(0).clone()
+                slot_mapping_seq = (
                     torch.tensor(slot_mapping[seq_i][-current_tkv:], dtype=torch.int64)
                     .unsqueeze(0)
                     .clone()
                 )
-                position_ids_i = (
+                position_ids_seq = (
                     kwargs["position_ids"][seq_i][-current_tkv:].unsqueeze(0).clone()
                 )
 
                 # This view will result in a discontiguous tensor (creates a new graph during compile)
                 # For this reason, we must explicitly make contiguous
-                mask_i = (
+                mask_seq = (
                     kwargs["mask"][seq_i][:, -current_tkv:, -current_tkv:]
                     .unsqueeze(0)
                     .contiguous()
                 )
 
-                # batch dynamic
-                torch._dynamo.mark_static(input_ids_i, 0)
-                torch._dynamo.mark_static(slot_mapping_i, 0)
-                torch._dynamo.mark_static(position_ids_i, 0)
-                torch._dynamo.mark_static(mask_i, 0)
-
-                # seq dynamic
-                torch._dynamo.mark_dynamic(input_ids_i, 1)
-                torch._dynamo.mark_dynamic(slot_mapping_i, 1)
-                torch._dynamo.mark_dynamic(position_ids_i, 1)
-                torch._dynamo.mark_dynamic(mask_i, 2)
-                torch._dynamo.mark_dynamic(mask_i, 3)
-
                 # FP8 per-sentence scale handling
                 if "fp8" in kwargs["attn_name"]:
                     for layer_idx, (t1, t2) in enumerate(current_kv_cache):
                         t1._scale = current_kv_scales[layer_idx][0][seq_i].reshape(-1)
                         t2._scale = current_kv_scales[layer_idx][1][seq_i].reshape(-1)
 
                 last_n_tokens = kwargs.get("last_n_tokens", 0)
-                output, current_kv_cache = model(
-                    input_ids_i,
-                    slot_mapping=slot_mapping_i,
-                    position_ids=position_ids_i,
-                    mask=mask_i,
-                    past_key_value_states=current_kv_cache,
-                    use_cache=kwargs["use_cache"],
-                    last_n_tokens=last_n_tokens,
-                    attn_name=kwargs["attn_name"],
-                )
 
-                # only last token must be handled here to properly stack the tensors
-                output = output[:, -1, :]
+                if prefill_chunk_size > 0:
+                    left_padded_prompt_mask_seq_chunk = None
+                    # Chunked prefill
+                    for chunk_j in range(math.ceil(current_tkv / prefill_chunk_size)):
+                        chunk_start = -current_tkv + chunk_j * prefill_chunk_size
+                        chunk_end = -current_tkv + min(
+                            (chunk_j + 1) * prefill_chunk_size, current_tkv
+                        )
+
+                        ids_length = input_ids[seq_i].shape[0]
+                        input_ids_seq_chunk = (
+                            input_ids[seq_i][
+                                chunk_start + ids_length : chunk_end + ids_length
+                            ]
+                            .unsqueeze(0)
+                            .clone()
+                        )
+                        slots_length = len(slot_mapping[seq_i])
+                        slot_mapping_seq_chunk = (
+                            torch.tensor(
+                                slot_mapping[seq_i][
+                                    chunk_start + slots_length : chunk_end
+                                    + slots_length
+                                ],
+                                dtype=torch.int64,
+                            )
+                            .unsqueeze(0)
+                            .clone()
+                        )
+                        pids_length = kwargs["position_ids"][seq_i].shape[0]
+                        position_ids_seq_chunk = (
+                            kwargs["position_ids"][seq_i][
+                                chunk_start + pids_length : chunk_end + pids_length
+                            ]
+                            .unsqueeze(0)
+                            .clone()
+                        )
+
+                        # This view will result in a discontiguous tensor (creates a new graph during compile)
+                        # For this reason, we must explicitly make contiguous
+                        if left_padded_prompt_mask_seq_chunk is None:
+                            left_padded_prompt_mask_seq_chunk = (
+                                position_ids_seq_chunk == 0
+                            ).sum(dim=1) - 1
+                        current_tkv_mask_seq_chunk = torch.min(
+                            torch.tensor(
+                                (chunk_j + 1) * prefill_chunk_size, dtype=torch.int64
+                            ),
+                            current_tkv,
+                        ).unsqueeze(0)
+
+                        table_length = len(block_table[seq_i])
+                        block_start = -current_tkv // BLOCK_SIZE + table_length
+                        block_end = chunk_end // BLOCK_SIZE + table_length
+                        block_table_seq_chunk = torch.tensor(
+                            block_table[seq_i][block_start:block_end], dtype=torch.int64
+                        ).unsqueeze(0)
+
+                        chunked_kwargs = {
+                            "slot_mapping": slot_mapping_seq_chunk,
+                            "position_ids": position_ids_seq_chunk,
+                            "past_key_value_states": current_kv_cache,
+                            "use_cache": kwargs["use_cache"],
+                            "last_n_tokens": kwargs["last_n_tokens"],
+                            "attn_name": kwargs["attn_name"],
+                            "left_padded_prompt_mask": left_padded_prompt_mask_seq_chunk,
+                            "current_tkv_mask": current_tkv_mask_seq_chunk,
+                            "block_table": block_table_seq_chunk,
+                        }
+
+                        # batch static
+                        torch._dynamo.mark_static(input_ids_seq_chunk, 0)
+                        torch._dynamo.mark_static(slot_mapping_seq_chunk, 0)
+                        torch._dynamo.mark_static(position_ids_seq_chunk, 0)
+                        torch._dynamo.mark_static(block_table_seq_chunk, 0)
+
+                        # seq dynamic
+                        torch._dynamo.mark_dynamic(input_ids_seq_chunk, 1)
+                        torch._dynamo.mark_dynamic(slot_mapping_seq_chunk, 1)
+                        torch._dynamo.mark_dynamic(position_ids_seq_chunk, 1)
+                        torch._dynamo.mark_dynamic(block_table_seq_chunk, 1)
+
+                        logits, current_kv_cache = model(
+                            input_ids_seq_chunk, **chunked_kwargs
+                        )
+
+                        # only last token must be handled here to properly stack the tensors
+                        logits = logits[:, -1, :]
+
+                        output = (logits, current_kv_cache)
+
+                else:
+                    # batch static
+                    torch._dynamo.mark_static(input_ids_seq, 0)
+                    torch._dynamo.mark_static(slot_mapping_seq, 0)
+                    torch._dynamo.mark_static(position_ids_seq, 0)
+                    torch._dynamo.mark_static(mask_seq, 0)
+
+                    # seq dynamic
+                    torch._dynamo.mark_dynamic(input_ids_seq, 1)
+                    torch._dynamo.mark_dynamic(slot_mapping_seq, 1)
+                    torch._dynamo.mark_dynamic(position_ids_seq, 1)
+                    torch._dynamo.mark_dynamic(mask_seq, 2)
+                    torch._dynamo.mark_dynamic(mask_seq, 3)
+                    output, current_kv_cache = model(
+                        input_ids_seq,
+                        slot_mapping=slot_mapping_seq,
+                        position_ids=position_ids_seq,
+                        mask=mask_seq,
+                        past_key_value_states=current_kv_cache,
+                        use_cache=kwargs["use_cache"],
+                        last_n_tokens=last_n_tokens,
+                        attn_name=kwargs["attn_name"],
+                    )
+
+                    # only last token must be handled here to properly stack the tensors
+                    output = output[:, -1, :]
 
                 # TODO: Figure out how to do this cleanly
                 if "fp8" in kwargs["attn_name"]: