vllm-xpu-kernels/csrc/xpu/attn/xe_2/chunk_prefill_configure.cmake at f7d6a9135ed4c451efed3f1ea782967a84c6059f · vllm-project/vllm-xpu-kernels · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
function(fmha_forward_configure FILENAME_SUFFIX)
  set(GEN_KERNEL_SRCS) # output
  set(L_TYPES "fp16" "bf16")
  set(L_BOOLS "false" "true")
  set(BOOL_FLAG_false "f")
  set(BOOL_FLAG_true "t")
  set(policy_list
      "chunk_policy_head64" "chunk_policy_head96" "chunk_policy_head128"
      "chunk_policy_head192" "chunk_policy_head256")

  set(IMPL_KV_T "fp16")

  foreach(IMPL_POLICY ${policy_list})
    # foreach(IMPL_T ${L_TYPES})
    foreach(IMPL_KISPAGED ${L_BOOLS})
      foreach(IMPL_KISCAUSAL ${L_BOOLS})
        foreach(IMPL_KISLOCAL ${L_BOOLS})
          foreach(IMPL_KISSINK ${L_BOOLS})
            set(FILE_SUFFIX "${IMPL_POLICY}_")
            set(FILE_SUFFIX "${FILE_SUFFIX}${BOOL_FLAG_${IMPL_KISPAGED}}")
            set(FILE_SUFFIX "${FILE_SUFFIX}${BOOL_FLAG_${IMPL_KISCAUSAL}}")
            set(FILE_SUFFIX "${FILE_SUFFIX}${BOOL_FLAG_${IMPL_KISSINK}}")
            set(FILE_SUFFIX "${FILE_SUFFIX}${BOOL_FLAG_${IMPL_KISLOCAL}}")
            configure_file(${FILENAME_SUFFIX}.cpp.in
                           "${FILENAME_SUFFIX}_${FILE_SUFFIX}.cpp")
            list(
              APPEND
              GEN_KERNEL_SRCS
              "${CMAKE_CURRENT_BINARY_DIR}/${FILENAME_SUFFIX}_${FILE_SUFFIX}.cpp"
            )
          endforeach()
        endforeach()
      endforeach()
    endforeach()
  endforeach()

  # Page-size-specific paged-only policies (only Paged=true)
  set(paged_policy_list
      "chunk_policy_head64_p16"
      "chunk_policy_head96_p16"
      "chunk_policy_head128_p16"
      "chunk_policy_head192_p16"
      "chunk_policy_head256_p16"
      "chunk_policy_head96_p32"
      "chunk_policy_head128_p32")

  set(IMPL_KISPAGED "true")
  foreach(IMPL_POLICY ${paged_policy_list})
    foreach(IMPL_KISCAUSAL ${L_BOOLS})
      foreach(IMPL_KISLOCAL ${L_BOOLS})
        foreach(IMPL_KISSINK ${L_BOOLS})
          set(FILE_SUFFIX "${IMPL_POLICY}_")
          set(FILE_SUFFIX "${FILE_SUFFIX}${BOOL_FLAG_${IMPL_KISPAGED}}")
          set(FILE_SUFFIX "${FILE_SUFFIX}${BOOL_FLAG_${IMPL_KISCAUSAL}}")
          set(FILE_SUFFIX "${FILE_SUFFIX}${BOOL_FLAG_${IMPL_KISSINK}}")
          set(FILE_SUFFIX "${FILE_SUFFIX}${BOOL_FLAG_${IMPL_KISLOCAL}}")
          configure_file(${FILENAME_SUFFIX}.cpp.in
                         "${FILENAME_SUFFIX}_${FILE_SUFFIX}.cpp")
          list(
            APPEND GEN_KERNEL_SRCS
            "${CMAKE_CURRENT_BINARY_DIR}/${FILENAME_SUFFIX}_${FILE_SUFFIX}.cpp")
        endforeach()
      endforeach()
    endforeach()
  endforeach()

  list(REMOVE_DUPLICATES GEN_KERNEL_SRCS)
  list(LENGTH GEN_KERNEL_SRCS GEN_KERNEL_SRCS_LENGTH)
  message(
    STATUS
      "Generated ${FILENAME_SUFFIX} kernel sources: ${GEN_KERNEL_SRCS_LENGTH}")
  set(GEN_KERNEL_SRCS
      ${GEN_KERNEL_SRCS}
      PARENT_SCOPE)
  set(GEN_KERNEL_SRCS_LENGTH
      ${GEN_KERNEL_SRCS_LENGTH}
      PARENT_SCOPE)

  list(APPEND ATTN_KERNEL_SRCS_GEN ${GEN_KERNEL_SRCS})
  set(ATTN_KERNEL_SRCS_GEN
      ${ATTN_KERNEL_SRCS_GEN}
      PARENT_SCOPE)

endfunction()