add multi-item scoring (#1015)

Co-authored with Qingquan Song (@qingquansong) and Ziang Li (@zianglih )

**Multi-item scoring** 
1. concatenate multiple candidates of a same member with all ranking
candidates with delimiter separation.
<member prefix (profile & history)> + <delimiter> + <item1> +
<delimiter> + item 2 + ... + item N <delimiter>
2. Extract the logits of the hidden states of the tokens before each
delimiter token and extract the log prob of given label tokens. For each
single prompt, output returned will be a 2D list with shape N * K where
N is the number of candidate it contains and K is the number of choices
we provided to the server engine (e.g., 2 for ["Yes", "No"])) (mainly
done in the logit processor)

![image](https://github.com/user-attachments/assets/837ea448-99e0-4e75-b6ec-5d7bf89b53dc)

The PR optimized the multi-item scoring attention by passing four new
args and use it to check the masking condition. The provided args are:
```
prefix_len_ptr :Optional[torch.Tensor]
    prefix length. A uint32 1D tensor indicating the prefix length of each prompt. The tensor size is equal to the batch size.
token_pos_in_items_ptr : Optional[float]
    A uint16 1D tensor (it will be converted to uint16 in flashinfer) indicating the token position of each item and started from 0 (delimiter)
    for each item. E.g., if we have 3 items of length 3, 2, 4 respectively for this member. This vector will be looking like
    `[0, 1, 2, 3, 0, 1, 2, 0, 1, 2, 3, 4, 0]` with 4 delimiters indexed as 0. For batch size > 1,
    we will concat them as 1D with zero paddings to make sure each has the same length, the padding length is defined by
    `token_pos_in_items_len` - length of the raw `token_pos_in_items_ptr` for each prompt.
token_pos_in_items_len : Optional[int]
    zero padding length for `token_pos_in_items_ptr` to better handle the bsz > 1 case. Still using the above 3,2,4 example.
    If we set `token_pos_in_items_len` to be 20, it will be  `[0, 1, 2, 3, 0, 1, 2, 0, 1, 2, 3, 4, 0, 0, 0, 0, 0, 0, 0, 0]`
    with 7 padded zeros. (note there're 8 zeros in the end where the first one is the delimiter token 0 in the end of the prompt)
max_item_len_ptr : Optional[float]
    a uint16 vector contains the max token length of all items for each prompt
```

**Optimizations**
1. Implement efficient multi-item scoring mask for FA2 and FA3.
2. Enhance FA3  to support batch-idx for the multi-item scoring mask.
3. Implement skip tiles for FA2 and FA3 multi-item scoring
4. Optimize mask by preloading to L1 cache for thread register.

---------

Co-authored-by: qingquansong <[email protected]>
Co-authored-by: zianglih <[email protected]>
Co-authored-by: Zihao Ye <[email protected]>

Author: 4 people

aot_build_utils/generate_aot_default_additional_params_header.py

@@ -125,23 +125,35 @@ def get_aot_default_additional_params_header_str() -> str:
 
     ret += generate_macro_entry(
         "BATCH_PREFILL",
-        ["maybe_custom_mask", "maybe_mask_indptr", "maybe_alibi_slopes"],
-        ["uint8_t", "int32_t", "float"],
+        [
+            "maybe_custom_mask",
+            "maybe_mask_indptr",
+            "maybe_alibi_slopes",
+            "maybe_prefix_len_ptr",
+            "maybe_token_pos_in_items_ptr",
+            "maybe_max_item_len_ptr",
+        ],
+        ["uint8_t", "int32_t", "float", "uint32_t", "uint16_t", "uint16_t"],
         [
             "logits_soft_cap",
             "sm_scale",
             "rope_rcp_scale",
             "rope_rcp_theta",
+            "token_pos_in_items_len",
         ],
-        ["double", "double", "double", "double"],
+        ["double", "double", "double", "double", "int64_t"],
     )
 
     ret += generate_macro_entry(
         "BATCH_PREFILL_SM90",
-        [],
-        [],
-        ["logits_soft_cap", "sm_scale"],
-        ["double", "double"],
+        [
+            "maybe_prefix_len_ptr",
+            "maybe_token_pos_in_items_ptr",
+            "maybe_max_item_len_ptr",
+        ],
+        ["uint32_t", "uint16_t", "uint16_t"],
+        ["logits_soft_cap", "sm_scale", "token_pos_in_items_len"],
+        ["double", "double", "int64_t"],
         is_sm90_template=True,
     )
 

aot_build_utils/literal_map.py

@@ -18,6 +18,7 @@
     0: "MaskMode::kNone",
     1: "MaskMode::kCausal",
     2: "MaskMode::kCustom",
+    3: "MaskMode::kMultiItemScoring",
 }
 
 pos_encoding_mode_literal = {

csrc/batch_prefill_sm90.cu

@@ -141,6 +141,8 @@ void BatchPrefillWithRaggedKVCacheSM90Run(at::Tensor float_workspace_buffer,
             GetPtrFromBaseOffset<IdType>(int_buffer_ptr, plan_info.head_indices_offset);
         params.work_indptr =
             GetPtrFromBaseOffset<IdType>(int_buffer_ptr, plan_info.work_indptr_offset);
+        params.batch_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer_ptr, plan_info.batch_indices_offset);
 
         ADDITIONAL_PARAMS_SETTER
 
@@ -238,6 +240,8 @@ void BatchPrefillWithPagedKVCacheSM90Run(
             GetPtrFromBaseOffset<IdType>(int_buffer_ptr, plan_info.head_indices_offset);
         params.work_indptr =
             GetPtrFromBaseOffset<IdType>(int_buffer_ptr, plan_info.work_indptr_offset);
+        params.batch_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer_ptr, plan_info.batch_indices_offset);
         params.kv_indices = static_cast<IdType*>(paged_kv_indices.data_ptr());
 
         ADDITIONAL_PARAMS_SETTER

csrc/batch_prefill_sm90_customize_config.jinja

@@ -43,6 +43,7 @@ struct RaggedParams {
   IdType* kv_lens;
   IdType* head_indices;
   IdType* work_indptr;
+  IdType* batch_indices;
 
   struct AdditionalParams {
     {{ additional_params_decl }}
@@ -88,6 +89,7 @@ struct PagedParams {
   IdType* kv_lens;
   IdType* head_indices;
   IdType* work_indptr;
+  IdType* batch_indices;
 
   struct AdditionalParams {
     {{ additional_params_decl }}

flashinfer/decode.py

@@ -1170,13 +1170,17 @@ def run(
                     None,  # packed_custom_mask
                     None,  # mask_indptr_buf
                     _get_cache_alibi_slopes_buf(q.shape[1], q.device),
+                    None,  # maybe_prefix_len_ptr
+                    None,  # maybe_token_pos_in_items_ptr
+                    None,  # maybe_max_item_len_ptr
                     logits_soft_cap,
                     sm_scale,
                     None,  # scale_q, not supported yet
                     None,  # scale_k
                     None,  # scale_v
                     rope_scale,
                     rope_theta,
+                    0,  # token_pos_in_items_len
                 ]
 
             self._cached_module.paged_run(*run_args)

flashinfer/jit/attention/pytorch.py

@@ -645,8 +645,8 @@ def gen_customize_pod_module(
 
     source_paths = []
 
-    for mask_mode_p in [0, 1, 2]:
-        for mask_mode_d in [0, 1, 2]:
+    for mask_mode_p in [0, 1, 2, 3]:
+        for mask_mode_d in [0, 1, 2, 3]:
             kwargs["mask_mode_p"] = mask_mode_literal[mask_mode_p]
             kwargs["mask_mode_d"] = mask_mode_literal[mask_mode_d]
 
@@ -759,27 +759,42 @@ def gen_batch_prefill_module(
             "maybe_custom_mask",
             "maybe_mask_indptr",
             "maybe_alibi_slopes",
+            "maybe_prefix_len_ptr",
+            "maybe_token_pos_in_items_ptr",
+            "maybe_max_item_len_ptr",
         ]
         additional_tensor_dtypes = [
             "uint8_t",
             "int32_t",
             "float",
+            "uint32_t",
+            "uint16_t",
+            "uint16_t",
         ]  # NOTE(Zihao): int32_t should follow dtype_idx
         additional_scalar_names = [
             "logits_soft_cap",
             "sm_scale",
             "rope_rcp_scale",
             "rope_rcp_theta",
+            "token_pos_in_items_len",
         ]
-        additional_scalar_dtypes = ["double", "double", "double", "double"]
+        additional_scalar_dtypes = ["double", "double", "double", "double", "int64_t"]
         variant_name = f"DefaultAttention<use_custom_mask, {str(use_sliding_window).lower()}, {str(use_logits_soft_cap).lower()}, {str(pos_encoding_mode == 2).lower()}>"
-        variant_decl = f"#include<flashinfer/attention/variants.cuh>"
+        variant_decl = "#include<flashinfer/attention/variants.cuh>"
     else:
         if not fp8_enabled:
-            additional_tensor_names = []
-            additional_tensor_dtypes = []
-            additional_scalar_names = ["logits_soft_cap", "sm_scale"]
-            additional_scalar_dtypes = ["double", "double"]
+            additional_tensor_names = [
+                "maybe_prefix_len_ptr",
+                "maybe_token_pos_in_items_ptr",
+                "maybe_max_item_len_ptr",
+            ]
+            additional_tensor_dtypes = ["uint32_t", "uint16_t", "uint16_t"]
+            additional_scalar_names = [
+                "logits_soft_cap",
+                "sm_scale",
+                "token_pos_in_items_len",
+            ]
+            additional_scalar_dtypes = ["double", "double", "int64_t"]
             variant_name = f"DefaultAttention<{str(use_logits_soft_cap).lower()}>"
             variant_decl = f"#include<flashinfer/attention/hopper/variants.cuh>"
         else:
@@ -961,7 +976,7 @@ def gen_customize_single_prefill_module(
         os.makedirs(gen_directory, exist_ok=True)
 
         source_paths = []
-        for mask_mode in [0, 1, 2]:
+        for mask_mode in [0, 1, 2, 3]:
             filename = f"single_prefill_kernel_mask_{mask_mode}.cu"
             dest_path = gen_directory / filename
             source_paths.append(dest_path)
@@ -1025,7 +1040,7 @@ def gen_customize_single_prefill_module(
         os.makedirs(gen_directory, exist_ok=True)
 
         source_paths = []
-        for mask_mode in [0, 1, 2]:
+        for mask_mode in [0, 1, 2, 3]:
             filename = f"single_prefill_sm90_kernel_mask_{mask_mode}.cu"
             dest_path = gen_directory / filename
             source_paths.append(dest_path)
@@ -1209,7 +1224,7 @@ def gen_customize_batch_prefill_module(
         os.makedirs(gen_directory, exist_ok=True)
 
         source_paths = []
-        for mask_mode in [0, 1, 2]:
+        for mask_mode in [0, 1, 2, 3]:
             dest_path = (
                 gen_directory / f"batch_prefill_paged_kernel_mask_{mask_mode}.cu"
             )
@@ -1286,7 +1301,7 @@ def gen_customize_batch_prefill_module(
         generated_inc_str = config_templ.render(**kwargs)
 
         source_paths = []
-        for mask_mode in [0, 1, 2]:
+        for mask_mode in [0, 1, 2, 3]:
             filename = f"batch_prefill_paged_sm90_kernel_mask_{mask_mode}.cu"
             dest_path = gen_directory / filename
             source_paths.append(dest_path)

flashinfer/jit/utils.py

@@ -98,4 +98,5 @@ def wrapper(func, args):
     0: "MaskMode::kNone",
     1: "MaskMode::kCausal",
     2: "MaskMode::kCustom",
+    3: "MaskMode::kMultiItemScoring",
 }