[Refactor] Unify JIT/Customization/AOT mode (#748)

This PR implements the #706 to unify the codebase for (1) JIT
compilation of default attention (2) JIT compilation of customized
attention (3) AOT compilation of default attention, and supports
customized attention for batch prefill/decode (both fa2/fa3 template).

More specifically:
1. All template files are stored in standalone Jinja files instead of
embedded python strings.
2. All attention modes use the same set of codebase. Default attentions
are instantiated as special forms of customized attention where
additional parameters are hard-coded.
3. The name of optional additional tensor parameters should start with
`maybe_`.
4. For FA3 template, additional parameters are set in an
`AdditionalParams` structure that will be passed to MainloopParams, so
that we can avoid passing the entire kernel parameter class where many
of the members are duplicate of MainloopParams and EpilogueParams.
5. The customized batch prefill/decode examples are added and tested.
6. We change the arguments order of pytorch bindings to unify customized
attention and default attention interface. The APIs exposed to user is
unchanged.

cc @hyhieu @merrymercy for visibility.

## Milestones
* [x] JIT default attention
* [x] JIT customized attention
* [x] AOT default attention
* [x] Check all unittests.
* [ ] C++ tests/benchmarks

## C++/Python Interface of PyTorch Bindings

### Single Decode/Prefill Atteniton Kernels

#### Decode C++ interface:
```cpp
#define AOT_ADDITIONAL_FUNC_PARAMS , std::optional<at::Tensor> maybe_alibi_slopes, \
    float logits_soft_cap, float sm_scale, float rope_scale, float rope_theta
void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp,
                                 at::Tensor o, unsigned int layout, int window_left
                                 ADDITIONAL_FUNC_PARAMS,
                                 int64_t cuda_stream);
```

#### Decode python interface:
```python
def single_decode_with_kv_cache(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, tmp: torch.Tensor,
                                 o: torch.Tensor, layout: int, window_left: int, *args,
                                 cuda_stream: int = 0) -> None:
    pass
```

For default attention, `*args` is expanded to `maybe_alibi_slopes,
logits_soft_cap, sm_scale, rope_scale, rope_theta`.

#### Prefill C++ interface:
```cpp
#define AOT_ADDITIONAL_FUNC_PARAMS , std::optional<at::Tensor> maybe_custom_mask, std::optional<at::Tensor> maybe_alibi_slopes, \
    float logits_soft_cap, float sm_scale, float rope_scale, float rope_theta
void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k,
                                  at::Tensor v, at::Tensor tmp, at::Tensor o, std::optional<at::Tensor> maybe_lse,
                                  unsigned int mask_mode_code, unsigned int layout, int32_t window_left
                                  ADDITIONAL_FUNC_PARAMS,
                                  int64_t cuda_stream);
```

#### Prefill python interface:
```python
def single_prefill_with_kv_cache(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, tmp: torch.Tensor,
                                     o: torch.Tensor, maybe_lse: Optional[torch.Tensor], mask_mode_code: int,
                                     layout: int, window_left: int, *args,
                                     cuda_stream: int = 0) -> None:
    pass
```

For default attention, `*args` is expanded to `maybe_custom_mask,
maybe_alibi_slopes, logits_soft_cap, sm_scale, rope_scale, rope_theta`
for fa2 template, and `logits_soft_cap, sm_scale` for fa3 template.

### Batch Decode/Prefill Attention Kernels

#### Decode c++ interface:
```cpp
#define AOT_ADDITIONAL_FUNC_PARAMS , std::optional<at::Tensor> maybe_alibi_slopes, \
    float logits_soft_cap, float sm_scale, float rope_scale, float rope_theta
std::vector<int64_t> BatchDecodeWithPagedKVCachePlan(
    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, unsigned int batch_size,
    unsigned int num_qo_heads, unsigned int num_kv_heads, unsigned int page_size,
    bool enable_cuda_graph, bool use_logits_soft_cap, unsigned int head_dim, at::Tensor empty_q_data,
    at::Tensor empty_kv_data, int64_t cuda_stream);

void BatchDecodeWithPagedKVCacheRun(
    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
    std::vector<int64_t> plan_info_vec, at::Tensor q, at::Tensor paged_k_cache,
    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices,
    at::Tensor paged_kv_last_page_len, at::Tensor o, std::optional<at::Tensor> maybe_lse,
    unsigned int kv_layout_code, int window_left
    ADDITIONAL_FUNC_PARAMS,
    int64_t cuda_stream);
```

#### Decode python interface
```python
def batch_decode_with_paged_kv_cache_plan(
    float_workspace_buffer: torch.Tensor, int_workspace_buffer: torch.Tensor,
    page_locked_int_workspace_buffer: torch.Tensor, indptr: torch.Tensor, batch_size: int,
    num_qo_heads: int, num_kv_heads: int, page_size: int, enable_cuda_graph: bool,
    use_logits_soft_cap: bool, head_dim: int, empty_q_data: torch.Tensor,
    empty_kv_data: torch.Tensor, cuda_stream: int) -> List[int]:
    pass

def batch_decode_with_paged_kv_cache_run(
    float_workspace_buffer: torch.Tensor, int_workspace_buffer: torch.Tensor,
    plan_info_vec: List[int], q: torch.Tensor, paged_k_cache: torch.Tensor,
    paged_v_cache: torch.Tensor, paged_kv_indptr: torch.Tensor, paged_kv_indices: torch.Tensor,
    paged_kv_last_page_len: torch.Tensor, o: torch.Tensor, maybe_lse: Optional[torch.Tensor],
    kv_layout_code: int, window_left: int, *args,
    cuda_stream: int) -> None:
    pass
```

For default attention, `*args` is expanded to `maybe_alibi_slopes,
logits_soft_cap, sm_scale, rope_scale, rope_theta`.

#### Prefill C++ interface:
```cpp
#define AOT_ADDITIONAL_FUNC_PARAMS , std::optional<at::Tensor> maybe_custom_mask, std::optional<at::Tensor> maybe_mask_indptr, std::optional<at::Tensor> maybe_alibi_slopes, \
    float logits_soft_cap, float sm_scale, float rope_scale, float rope_theta
std::vector<int64_t> BatchPrefillWithKVCachePlan(
    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
    at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor kv_len_arr,
    unsigned total_num_rows, unsigned int batch_size, unsigned int num_qo_heads,
    unsigned int num_kv_heads, unsigned int page_size, bool enable_cuda_graph,
    unsigned int head_dim, bool causal,
    int64_t cuda_stream);

void BatchPrefillWithRaggedKVCacheRun(
    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
    std::vector<int64_t> plan_info_vec, at::Tensor q, at::Tensor k, at::Tensor v,
    at::Tensor qo_indptr, at::Tensor kv_indptr,
    at::Tensor o, std::optional<at::Tensor> maybe_lse,
    unsigned int mask_mode_code, unsigned int layout, int32_t window_left
    ADDITIONAL_FUNC_PARAMS,
    int64_t cuda_stream);

void BatchPrefillWithPagedKVCacheRun(
    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
    std::vector<int64_t> plan_info_vec, at::Tensor q, at::Tensor paged_k_cache,
    at::Tensor paged_v_cache, at::Tensor qo_indptr, at::Tensor paged_kv_indptr,
    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
    at::Tensor o, std::optional<at::Tensor> maybe_lse,
    unsigned int mask_mode_code, unsigned int layout, int32_t window_left
    ADDITIONAL_FUNC_PARAMS,
    int64_t cuda_stream);
```

#### Prefill python interface:
```python
def batch_prefill_with_kv_cache_plan(
    float_workspace_buffer: torch.Tensor, int_workspace_buffer: torch.Tensor,
    page_locked_int_workspace_buffer: torch.Tensor, qo_indptr: torch.Tensor, kv_indptr: torch.Tensor, kv_len_arr: torch.Tensor,
    total_num_rows: int, batch_size: int, num_qo_heads: int, num_kv_heads: int, page_size: int,
    enable_cuda_graph: bool, head_dim: int, causal: bool, cuda_stream: int) -> List[int]:
    pass

def batch_prefill_with_ragged_kv_cache_jit_run(
    float_workspace_buffer: torch.Tensor, int_workspace_buffer: torch.Tensor,
    plan_info_vec: List[int], q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
    qo_indptr: torch.Tensor, kv_indptr: torch.Tensor, o: torch.Tensor,
    maybe_lse: Optional[torch.Tensor], mask_mode_code: int, layout: int, window_left: int,
    *args, cuda_stream: int) -> None:
    pass

def batch_prefill_with_paged_kv_cache_jit_run(
    float_workspace_buffer: torch.Tensor, int_workspace_buffer: torch.Tensor,
    plan_info_vec: List[int], q: torch.Tensor, paged_k_cache: torch.Tensor, paged_v_cache: torch.Tensor,
    qo_indptr: torch.Tensor, paged_kv_indptr: torch.Tensor, paged_kv_indices: torch.Tensor,
    paged_kv_last_page_len: torch.Tensor, o: torch.Tensor, maybe_lse: Optional[torch.Tensor],
    mask_mode_code: int, layout: int, window_left: int, *args,
    cuda_stream: int) -> None:
    pass
```

The `*args` is expanded to `maybe_custom_mask, maybe_mask_indptr,
maybe_alibi_slopes, logits_soft_cap, sm_scale, rope_scale, rope_theta`
for fa2 template, and `logits_soft_cap, sm_scale` for fa3 template.

Author: yzh119

CMakeLists.txt

@@ -66,7 +66,7 @@ flashinfer_option(FLASHINFER_TVM_SOURCE_DIR
 flashinfer_option(FLASHINFER_GEN_HEAD_DIMS "Head dims to enable" 64 128 256)
 flashinfer_option(FLASHINFER_GEN_POS_ENCODING_MODES "Pos encodings to enable" 0
                   1 2)
-flashinfer_option(FLASHINFER_GEN_ALLOW_FP16_QK_REDUCTIONS
+flashinfer_option(FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS
                   "QK reductions to enable" "false" "true")
 flashinfer_option(FLASHINFER_GEN_MASK_MODES "Mask modes to enable" 0 1 2)
 
@@ -126,34 +126,28 @@ endif(FLASHINFER_ENABLE_BF16)
 # generate kernel inst
 set(HEAD_DIMS ${FLASHINFER_GEN_HEAD_DIMS})
 set(POS_ENCODING_MODES ${FLASHINFER_GEN_POS_ENCODING_MODES})
-set(ALLOW_FP16_QK_REDUCTIONS ${FLASHINFER_GEN_ALLOW_FP16_QK_REDUCTIONS})
+set(USE_FP16_QK_REDUCTIONS ${FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS})
 set(MASK_MODES ${FLASHINFER_GEN_MASK_MODES})
 
 # log options
 message(STATUS "FLASHINFER_HEAD_DIMS=${HEAD_DIMS}")
 message(STATUS "FLASHINFER_POS_ENCODING_MODES=${POS_ENCODING_MODES}")
-message(
-  STATUS "FLASHINFER_ALLOW_FP16_QK_REDUCTIONS=${ALLOW_FP16_QK_REDUCTIONS}")
+message(STATUS "FLASHINFER_USE_FP16_QK_REDUCTIONS=${USE_FP16_QK_REDUCTIONS}")
 message(STATUS "FLASHINFER_MASK_MODES=${MASK_MODES}")
 
 file(MAKE_DIRECTORY ${PROJECT_SOURCE_DIR}/src/generated)
 
 set(AOT_GENERATE_COMMAND
-  ${Python3_EXECUTABLE}
-  -m aot_build_utils.generate
-  --path ${PROJECT_SOURCE_DIR}/src/generated
-  --head_dims ${HEAD_DIMS}
-  --pos_encoding_modes ${POS_ENCODING_MODES}
-  --allow_fp16_qk_reductions ${ALLOW_FP16_QK_REDUCTIONS}
-  --mask_modes ${MASK_MODES}
-  --enable_f16 ${FLASHINFER_ENABLE_F16}
-  --enable_bf16 ${FLASHINFER_ENABLE_BF16}
-  --enable_fp8_e4m3 ${FLASHINFER_ENABLE_FP8_E4M3}
-  --enable_fp8_e5m2 ${FLASHINFER_ENABLE_FP8_E5M2})
-
-execute_process(
-  COMMAND ${AOT_GENERATE_COMMAND}
-  WORKING_DIRECTORY ${PROJECT_SOURCE_DIR})
+    ${Python3_EXECUTABLE} -m aot_build_utils.generate --path
+    ${PROJECT_SOURCE_DIR}/src/generated --head_dims ${HEAD_DIMS}
+    --pos_encoding_modes ${POS_ENCODING_MODES} --use_fp16_qk_reductions
+    ${USE_FP16_QK_REDUCTIONS} --mask_modes ${MASK_MODES} --enable_f16
+    ${FLASHINFER_ENABLE_F16} --enable_bf16 ${FLASHINFER_ENABLE_BF16}
+    --enable_fp8_e4m3 ${FLASHINFER_ENABLE_FP8_E4M3} --enable_fp8_e5m2
+    ${FLASHINFER_ENABLE_FP8_E5M2})
+
+execute_process(COMMAND ${AOT_GENERATE_COMMAND}
+                WORKING_DIRECTORY ${PROJECT_SOURCE_DIR})
 
 file(GLOB_RECURSE FLASHINFER_GENERATORS
      ${PROJECT_SOURCE_DIR}/aot_build_utils/*.py)
@@ -175,13 +169,13 @@ add_custom_target(dispatch_inc DEPENDS ${DISPATCH_INC_FILE})
 
 add_library(decode_kernels STATIC ${DECODE_KERNELS_SRCS})
 target_include_directories(decode_kernels PRIVATE ${FLASHINFER_INCLUDE_DIR})
-target_compile_options(decode_kernels PRIVATE
-  -Xcompiler=-fPIC --fatbin-options -compress-all)
+target_compile_options(decode_kernels PRIVATE -Xcompiler=-fPIC --fatbin-options
+                                              -compress-all)
 
 add_library(prefill_kernels STATIC ${PREFILL_KERNELS_SRCS})
 target_include_directories(prefill_kernels PRIVATE ${FLASHINFER_INCLUDE_DIR})
-target_compile_options(prefill_kernels PRIVATE
-  -Xcompiler=-fPIC --fatbin-options -compress-all)
+target_compile_options(prefill_kernels PRIVATE -Xcompiler=-fPIC
+                                               --fatbin-options -compress-all)
 
 if(FLASHINFER_DECODE)
   message(STATUS "Compile single decode kernel benchmarks.")

aot_build_utils/generate.py

@@ -20,6 +20,7 @@
 from typing import List
 
 from . import (
+    generate_aot_default_additional_params_header,
     generate_batch_paged_decode_inst,
     generate_batch_paged_prefill_inst,
     generate_batch_ragged_prefill_inst,
@@ -38,7 +39,7 @@ def write_if_different(path: Path, content: str) -> None:
     path: Path = args.path
     head_dims: List[int] = args.head_dims
     pos_encoding_modes: List[int] = args.pos_encoding_modes
-    allow_fp16_qk_reductions: List[int] = args.allow_fp16_qk_reductions
+    use_fp16_qk_reductions: List[int] = args.use_fp16_qk_reductions
     mask_modes: List[int] = args.mask_modes
     enable_f16: bool = args.enable_f16
     enable_bf16: bool = args.enable_bf16
@@ -54,12 +55,17 @@ def write_if_different(path: Path, content: str) -> None:
             argparse.Namespace(
                 head_dims=head_dims,
                 pos_encoding_modes=pos_encoding_modes,
-                allow_fp16_qk_reductions=allow_fp16_qk_reductions,
+                use_fp16_qk_reductions=use_fp16_qk_reductions,
                 mask_modes=mask_modes,
             )
         ),
     )
 
+    write_if_different(
+        path / "aot_default_additional_params.h",
+        generate_aot_default_additional_params_header.get_aot_default_additional_params_header_str(),
+    )
+
     idtypes = ["i32"]
     prefill_dtypes = []
     decode_dtypes = []
@@ -150,22 +156,22 @@ def write_if_different(path: Path, content: str) -> None:
     for (
         head_dim,
         pos_encoding_mode,
-        allow_fp16_qk_reduction,
+        use_fp16_qk_reduction,
         mask_mode,
     ) in product(
         head_dims,
         pos_encoding_modes,
-        allow_fp16_qk_reductions,
+        use_fp16_qk_reductions,
         mask_modes,
     ):
         for dtype_q, dtype_kv in list(zip(prefill_dtypes, prefill_dtypes)) + list(
             product(prefill_dtypes, fp8_dtypes)
         ):
-            fname = f"single_prefill_head_{head_dim}_posenc_{pos_encoding_mode}_fp16qkred_{allow_fp16_qk_reduction}_mask_{mask_mode}_dtypeq_{dtype_q}_dtypekv_{dtype_kv}_dtypeout_{dtype_q}.cu"
+            fname = f"single_prefill_head_{head_dim}_posenc_{pos_encoding_mode}_fp16qkred_{use_fp16_qk_reduction}_mask_{mask_mode}_dtypeq_{dtype_q}_dtypekv_{dtype_kv}_dtypeout_{dtype_q}.cu"
             content = generate_single_prefill_inst.get_cu_file_str(
                 head_dim,
                 pos_encoding_mode,
-                allow_fp16_qk_reduction,
+                use_fp16_qk_reduction,
                 mask_mode,
                 dtype_q,  # dtype_q
                 dtype_kv,  # dtype_kv
@@ -184,7 +190,7 @@ def write_if_different(path: Path, content: str) -> None:
                             f"posenc_{pos_encoding_mode}_"
                             f"use_swa_{use_sliding_window}_"
                             f"use_logits_cap_{use_logits_soft_cap}_"
-                            f"f16qk_{bool(allow_fp16_qk_reduction)}"
+                            f"f16qk_{bool(use_fp16_qk_reduction)}"
                         )
             write_if_different(path / fname, content)
 
@@ -193,24 +199,24 @@ def write_if_different(path: Path, content: str) -> None:
     for (
         head_dim,
         pos_encoding_mode,
-        allow_fp16_qk_reduction,
+        use_fp16_qk_reduction,
         mask_mode,
         idtype,
     ) in product(
         head_dims,
         pos_encoding_modes,
-        allow_fp16_qk_reductions,
+        use_fp16_qk_reductions,
         mask_modes,
         idtypes,
     ):
         for dtype_q, dtype_kv in list(zip(prefill_dtypes, prefill_dtypes)) + list(
             product(prefill_dtypes, fp8_dtypes)
         ):
-            fname = f"batch_paged_prefill_head_{head_dim}_posenc_{pos_encoding_mode}_fp16qkred_{allow_fp16_qk_reduction}_mask_{mask_mode}_dtypeq_{dtype_q}_dtypekv_{dtype_kv}_dtypeout_{dtype_q}_idtype_{idtype}.cu"
+            fname = f"batch_paged_prefill_head_{head_dim}_posenc_{pos_encoding_mode}_fp16qkred_{use_fp16_qk_reduction}_mask_{mask_mode}_dtypeq_{dtype_q}_dtypekv_{dtype_kv}_dtypeout_{dtype_q}_idtype_{idtype}.cu"
             content = generate_batch_paged_prefill_inst.get_cu_file_str(
                 head_dim,
                 pos_encoding_mode,
-                allow_fp16_qk_reduction,
+                use_fp16_qk_reduction,
                 mask_mode,
                 dtype_q,  # dtype_q
                 dtype_kv,  # dtype_kv
@@ -219,11 +225,11 @@ def write_if_different(path: Path, content: str) -> None:
             )
             write_if_different(path / fname, content)
 
-            fname = f"batch_ragged_prefill_head_{head_dim}_posenc_{pos_encoding_mode}_fp16qkred_{allow_fp16_qk_reduction}_mask_{mask_mode}_dtypeq_{dtype_q}_dtypekv_{dtype_kv}_dtypeout_{dtype_q}_idtype_{idtype}.cu"
+            fname = f"batch_ragged_prefill_head_{head_dim}_posenc_{pos_encoding_mode}_fp16qkred_{use_fp16_qk_reduction}_mask_{mask_mode}_dtypeq_{dtype_q}_dtypekv_{dtype_kv}_dtypeout_{dtype_q}_idtype_{idtype}.cu"
             content = generate_batch_ragged_prefill_inst.get_cu_file_str(
                 head_dim,
                 pos_encoding_mode,
-                allow_fp16_qk_reduction,
+                use_fp16_qk_reduction,
                 mask_mode,
                 dtype_q,  # dtype_q
                 dtype_kv,  # dtype_kv
@@ -246,7 +252,7 @@ def write_if_different(path: Path, content: str) -> None:
                             f"posenc_{pos_encoding_mode}_"
                             f"use_swa_{sliding_window}_"
                             f"use_logits_cap_{logits_soft_cap}_"
-                            f"f16qk_{bool(allow_fp16_qk_reduction)}"
+                            f"f16qk_{bool(use_fp16_qk_reduction)}"
                         )
 
     return (
@@ -273,7 +279,7 @@ def write_if_different(path: Path, content: str) -> None:
         help="Position encoding modes",
     )
     parser.add_argument(
-        "--allow_fp16_qk_reductions",
+        "--use_fp16_qk_reductions",
         type=lambda x: x if isinstance(x, int) else int(x.lower() == "true"),
         required=True,
         nargs="+",
@@ -287,7 +293,7 @@ def write_if_different(path: Path, content: str) -> None:
         help="Mask modes",
     )
     parser.add_argument(
-        "--enable_fp16",
+        "--enable_f16",
         type=lambda x: x if isinstance(x, int) else x.lower() == "true",
         required=True,
         nargs="+",

aot_build_utils/generate_aot_default_additional_params_header.py

@@ -0,0 +1,148 @@
+"""
+Copyright (c) 2024 by FlashInfer team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+
+def generate_macro_entry(
+    macro_prefix,
+    additional_tensor_names,
+    additional_tensor_dtypes,
+    additional_scalar_names,
+    additional_scalar_dtypes,
+    is_sm90_template: bool = False,
+) -> str:
+    # NOTE(Zihao): mostly copy-paste from generate_additional_params in flashinfer.jit.attention.py
+    additional_func_params = "".join(
+        [
+            (
+                f", std::optional<at::Tensor> {var}"
+                if var.startswith("maybe")
+                else f", at::Tensor {var}"
+            )
+            for var in additional_tensor_names
+        ]
+        + [
+            f", {dtype} {var}"
+            for dtype, var in zip(additional_scalar_dtypes, additional_scalar_names)
+        ]
+    )
+    if is_sm90_template:
+        additional_params_setter = " \\\n".join(
+            [
+                (
+                    f"params.additional_params.{var} = {var} ? static_cast<{dtype}*>({var}->data_ptr()): nullptr;"
+                    if var.startswith("maybe")
+                    else f"params.additional_params.{var} = static_cast<{dtype}*>({var}.data_ptr());"
+                )
+                for dtype, var in zip(additional_tensor_dtypes, additional_tensor_names)
+            ]
+            + [
+                f"params.additional_params.{var} = {var};"
+                for var in additional_scalar_names
+            ]
+        )
+    else:
+        additional_params_setter = " \\\n".join(
+            [
+                (
+                    f"params.{var} = {var} ? static_cast<{dtype}*>({var}->data_ptr()): nullptr;"
+                    if var.startswith("maybe")
+                    else f"params.{var} = static_cast<{dtype}*>({var}.data_ptr());"
+                )
+                for dtype, var in zip(additional_tensor_dtypes, additional_tensor_names)
+            ]
+            + [f"params.{var} = {var};" for var in additional_scalar_names]
+        )
+    return f"""#define {macro_prefix}_ADDITIONAL_FUNC_PARAMS {additional_func_params}
+
+#define {macro_prefix}_ADDITIONAL_PARAMS_SETTER {additional_params_setter}
+
+"""
+
+
+def get_aot_default_additional_params_header_str() -> str:
+    ret = ""
+
+    ret += generate_macro_entry(
+        "SINGLE_DECODE",
+        ["maybe_alibi_slopes"],  # additional_tensor_names
+        ["float"],  # additional_tensor_dtypes
+        [
+            "logits_soft_cap",
+            "sm_scale",
+            "rope_rcp_scale",
+            "rope_rcp_theta",
+        ],  # additional_scalar_names
+        ["float", "float", "float", "float"],  # additional_scalar_dtypes
+    )
+
+    ret += generate_macro_entry(
+        "SINGLE_PREFILL",
+        ["maybe_custom_mask", "maybe_alibi_slopes"],
+        ["uint8_t", "float"],
+        [
+            "logits_soft_cap",
+            "sm_scale",
+            "rope_rcp_scale",
+            "rope_rcp_theta",
+        ],
+        ["float", "float", "float", "float"],
+    )
+
+    ret += generate_macro_entry(
+        "SINGLE_PREFILL_SM90",
+        [],
+        [],
+        ["logits_soft_cap", "sm_scale"],
+        ["float", "float"],
+        is_sm90_template=True,
+    )
+
+    ret += generate_macro_entry(
+        "BATCH_DECODE",
+        ["maybe_alibi_slopes"],  # additional_tensor_names
+        ["float"],  # additional_tensor_dtypes
+        [
+            "logits_soft_cap",
+            "sm_scale",
+            "rope_rcp_scale",
+            "rope_rcp_theta",
+        ],  # additional_scalar_names
+        ["float", "float", "float", "float"],  # additional_scalar_dtypes
+    )
+
+    ret += generate_macro_entry(
+        "BATCH_PREFILL",
+        ["maybe_custom_mask", "maybe_mask_indptr", "maybe_alibi_slopes"],
+        ["uint8_t", "int32_t", "float"],
+        [
+            "logits_soft_cap",
+            "sm_scale",
+            "rope_rcp_scale",
+            "rope_rcp_theta",
+        ],
+        ["float", "float", "float", "float"],
+    )
+
+    ret += generate_macro_entry(
+        "BATCH_PREFILL_SM90",
+        [],
+        [],
+        ["logits_soft_cap", "sm_scale"],
+        ["float", "float"],
+        is_sm90_template=True,
+    )
+
+    return ret