| 1 | #!/usr/bin/env python3 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | """Generate CPU attention dispatch switch cases and kernel instantiations.""" |
| 5 | |
| 6 | import os |
| 7 | |
| 8 | # Head dimensions divisible by 32 (support all ISAs) |
| 9 | HEAD_DIMS_32 = [32, 64, 96, 128, 160, 192, 224, 256, 512] |
| 10 | |
| 11 | # Head dimensions divisible by 16 but not 32 (VEC16 only) |
| 12 | HEAD_DIMS_16 = [48, 80, 112] |
| 13 | |
| 14 | # ISA types |
| 15 | ISA_TYPES = { |
| 16 | "AMX": 0, |
| 17 | "VEC": 1, |
| 18 | "VEC16": 2, |
| 19 | "NEON": 3, |
| 20 | "VXE": 4, |
| 21 | "RVV": 5, |
| 22 | "VSX": 6, |
| 23 | "AMX_FP8": 7, |
| 24 | } |
| 25 | |
| 26 | # KV cache index: 0 = auto (same as scalar_t), 1 = fp8_e4m3, 2 = fp8_e5m2 |
| 27 | KV_CACHE_IDX = { |
| 28 | "auto": 0, |
| 29 | "fp8_e4m3": 1, |
| 30 | "fp8_e5m2": 2, |
| 31 | } |
| 32 | |
| 33 | # C++ type for each kv_cache index |
| 34 | KV_CACHE_CPP_TYPES = { |
| 35 | "auto": "scalar_t", |
| 36 | "fp8_e4m3": "c10::Float8_e4m3fn", |
| 37 | "fp8_e5m2": "c10::Float8_e5m2", |
| 38 | } |
| 39 | |
| 40 | # ISAs supported for head_dims divisible by 32 |
| 41 | ISA_FOR_32 = ["AMX", "AMX_FP8", "NEON", "VEC", "VEC16", "VXE", "RVV", "VSX"] |
| 42 | |
| 43 | # ISAs supported for head_dims divisible by 16 only |
| 44 | ISA_FOR_16 = ["VEC16"] |
| 45 | |
| 46 | # ISAs that support FP8 KV cache (x86 AVX2/AVX-512 required) |
| 47 | ISA_FOR_FP8 = ["AMX_FP8", "AMX", "VEC"] |
| 48 | |
| 49 | |
| 50 | def encode_params(head_dim: int, isa_type: str, kv_cache: str = "auto") -> int: |
| 51 | """Encode head_dim, ISA type, and KV cache type into a single int64_t.""" |
| 52 | isa_val = ISA_TYPES[isa_type] |
| 53 | kv_val = KV_CACHE_IDX[kv_cache] |
| 54 | # Encoding: (head_dim << 16) | (kv_cache_idx << 8) | isa_type |
| 55 | # This allows head_dim up to 2^48 - 1, 256 KV cache types, and 256 ISA types |
| 56 | return (head_dim << 16) | (kv_val << 8) | isa_val |
| 57 | |
| 58 | |
| 59 | def _make_case( |
| 60 | head_dim: int, isa: str, kv_cache: str = "auto", isa_override: str | None = None |
| 61 | ) -> str: |
| 62 | """Generate a single switch case line.""" |
| 63 | encoded = encode_params(head_dim, isa, kv_cache) |
| 64 | actual_isa = isa_override if isa_override else isa |
| 65 | cpp_type = KV_CACHE_CPP_TYPES[kv_cache] |
| 66 | native_impl = ( |
| 67 | f"cpu_attention::AttentionImpl<cpu_attention::ISA::{actual_isa}, " |
| 68 | f"scalar_t, head_dim, {cpp_type}>" |
| 69 | ) |
| 70 | if actual_isa == "AMX_FP8": |
| 71 | fallback_impl = ( |
| 72 | f"cpu_attention::AttentionImpl<cpu_attention::ISA::AMX, " |
| 73 | f"scalar_t, head_dim, {cpp_type}>" |
| 74 | ) |
| 75 | attn_impl = ( |
| 76 | "std::conditional_t<std::is_same_v<scalar_t, c10::BFloat16>, " |
| 77 | f"{native_impl}, {fallback_impl}>" |
| 78 | ) |
| 79 | else: |
| 80 | attn_impl = native_impl |
| 81 | comment = ( |
| 82 | f"head_dim={head_dim}, isa={isa}" |
| 83 | if kv_cache == "auto" |
| 84 | else f"head_dim={head_dim}, isa={isa}, kv_cache={kv_cache}" |
| 85 | ) |
| 86 | return ( |
| 87 | f""" case {encoded}LL: {{ """ |
| 88 | f"""/* {comment} */ \\""" |
| 89 | f""" |
| 90 | constexpr size_t head_dim = {head_dim}; \\""" |
| 91 | f""" |
| 92 | using attn_impl = {attn_impl}; \\""" |
| 93 | f""" |
| 94 | return __VA_ARGS__(); \\""" |
| 95 | f""" |
| 96 | }} \\""" |
| 97 | ) |
| 98 | |
| 99 | |
| 100 | def generate_cases_for_isa_group(isa_list: list[str], include_fp8: bool = False) -> str: |
| 101 | """Generate switch cases for a specific ISA group.""" |
| 102 | cases = [] |
| 103 | |
| 104 | # Non-FP8 cases for head_dims divisible by 32 |
| 105 | for head_dim in HEAD_DIMS_32: |
| 106 | for isa in isa_list: |
| 107 | if isa not in ISA_FOR_32: |
| 108 | continue |
| 109 | isa_override = "AMX" if isa == "AMX_FP8" else None |
| 110 | cases.append(_make_case(head_dim, isa, "auto", isa_override)) |
| 111 | |
| 112 | # Non-FP8 cases for head_dims divisible by 16 only |
| 113 | for head_dim in HEAD_DIMS_16: |
| 114 | for isa in isa_list: |
| 115 | cases.append(_make_case(head_dim, isa, "auto", isa_override="VEC16")) |
| 116 | |
| 117 | # FP8 cases: only AMX and VEC, only head_dims divisible by 32 |
| 118 | if include_fp8: |
| 119 | for fp8_type in ("fp8_e4m3", "fp8_e5m2"): |
| 120 | for head_dim in HEAD_DIMS_32: |
| 121 | for isa in isa_list: |
| 122 | if isa not in ISA_FOR_FP8: |
| 123 | continue |
| 124 | isa_override = "AMX" if isa == "AMX_FP8" and head_dim % 64 else None |
| 125 | cases.append(_make_case(head_dim, isa, fp8_type, isa_override)) |
| 126 | |
| 127 | return "\n".join(cases) |
| 128 | |
| 129 | |
| 130 | def generate_helper_function() -> str: |
| 131 | """Generate helper function to encode parameters.""" |
| 132 | return """ |
| 133 | inline int64_t encode_cpu_attn_params(int64_t head_dim, cpu_attention::ISA isa, |
| 134 | int64_t kv_cache_idx = 0) { |
| 135 | return (head_dim << 16) | (kv_cache_idx << 8) | static_cast<int64_t>(isa); |
| 136 | } |
| 137 | """ |
| 138 | |
| 139 | |
| 140 | def generate_header_file() -> str: |
| 141 | """Generate the complete header file content.""" |
| 142 | header = """// auto generated by generate_cpu_attn_dispatch.py |
| 143 | // clang-format off |
| 144 | |
| 145 | #ifndef CPU_ATTN_DISPATCH_GENERATED_H |
| 146 | #define CPU_ATTN_DISPATCH_GENERATED_H |
| 147 | |
| 148 | #include "cpu_attn_vec.hpp" |
| 149 | #include "cpu_attn_vec16.hpp" |
| 150 | |
| 151 | #ifdef CPU_CAPABILITY_AMXFP8 |
| 152 | #include "cpu_attn_amx_fp8.hpp" |
| 153 | #endif |
| 154 | |
| 155 | #ifdef CPU_CAPABILITY_AMXBF16 |
| 156 | #include "cpu_attn_amx.hpp" |
| 157 | #endif |
| 158 | |
| 159 | #ifdef __aarch64__ |
| 160 | #include "cpu_attn_neon.hpp" |
| 161 | #endif |
| 162 | |
| 163 | #ifdef __s390x__ |
| 164 | #include "cpu_attn_vxe.hpp" |
| 165 | #endif |
| 166 | |
| 167 | // cpu_attn_rvv.hpp supports VLEN=128 and VLEN=256 via RVVI() macros. |
| 168 | // Other VLENs and scalar RISC-V builds skip it entirely. |
| 169 | #if defined(__riscv) && defined(__riscv_v_min_vlen) && \ |
| 170 | (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256) |
| 171 | #include "cpu_attn_rvv.hpp" |
| 172 | #endif |
| 173 | |
| 174 | #ifdef __powerpc__ |
| 175 | #include "cpu_attn_vsx.hpp" |
| 176 | #endif |
| 177 | |
| 178 | """ |
| 179 | |
| 180 | header += generate_helper_function() |
| 181 | |
| 182 | # Generate dispatch macro with conditional compilation for different ISA sets |
| 183 | header += """ |
| 184 | // Dispatch macro using encoded parameters. |
| 185 | // KV_CACHE_IDX: Fp8KVCacheDataType enum value (kAuto=0, kFp8E4M3=1, kFp8E5M2=2). |
| 186 | // FP8 cases (kv_cache_idx != 0) are generated on x86 platforms with AVX2 or |
| 187 | // AVX-512: BF16Vec32 FP8 constructors have both AVX-512 and AVX2 implementations |
| 188 | // in cpu_types_x86.hpp. Non-x86 platforms (#else fallback) have fp8=False. |
| 189 | """ |
| 190 | |
| 191 | def _macro_block(guard: str, isa_list: list[str], fp8: bool) -> str: |
| 192 | """Return one CPU_ATTN_DISPATCH macro block for a given guard.""" |
| 193 | enc = ( |
| 194 | " int64_t encoded_params = encode_cpu_attn_params(" |
| 195 | "HEAD_DIM, ISA_TYPE, KV_CACHE_IDX); \\" |
| 196 | ) |
| 197 | cases = generate_cases_for_isa_group(isa_list, include_fp8=fp8) |
| 198 | tail = ( |
| 199 | "\n" |
| 200 | " default: { \\\n" |
| 201 | " TORCH_CHECK(false, " |
| 202 | '"Unsupported CPU attention configuration: head_dim=" + \\\n' |
| 203 | ' std::to_string(HEAD_DIM) + " isa=" + \\\n' |
| 204 | " std::to_string(static_cast<int>(ISA_TYPE))" |
| 205 | " + \\\n" |
| 206 | ' " kv_cache_idx=" + ' |
| 207 | "std::to_string(KV_CACHE_IDX)); \\\n" |
| 208 | " } \\\n" |
| 209 | " } \\\n" |
| 210 | " }()\n\n" |
| 211 | ) |
| 212 | return ( |
| 213 | f"{guard}\n" |
| 214 | "#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, KV_CACHE_IDX, ...) \\\n" |
| 215 | " [&] { \\\n" |
| 216 | f"{enc}\n" |
| 217 | " switch (encoded_params) { \\\n" |
| 218 | f"{cases}" |
| 219 | f"{tail}" |
| 220 | ) |
| 221 | |
| 222 | header += _macro_block( |
| 223 | "#if defined(CPU_CAPABILITY_AMXFP8)", |
| 224 | ["AMX_FP8", "AMX", "VEC", "VEC16"], |
| 225 | fp8=True, |
| 226 | ) |
| 227 | header += _macro_block( |
| 228 | "#elif defined(CPU_CAPABILITY_AMXBF16)", |
| 229 | ["AMX", "VEC", "VEC16"], |
| 230 | fp8=True, |
| 231 | ) |
| 232 | header += _macro_block( |
| 233 | "#elif defined(__aarch64__)", |
| 234 | ["NEON", "VEC", "VEC16"], |
| 235 | fp8=False, |
| 236 | ) |
| 237 | header += _macro_block( |
| 238 | "#elif defined(__s390x__)", |
| 239 | ["VXE", "VEC", "VEC16"], |
| 240 | fp8=False, |
| 241 | ) |
| 242 | # RISC-V with RVV. cpu_attn_rvv.hpp supports VLEN=128 and VLEN=256 |
| 243 | # via RVVI() macros. Builds with a supported VLEN get |
| 244 | # RVV+VEC+VEC16; other RISC-V builds fall back to VEC/VEC16 only. |
| 245 | header += _macro_block( |
| 246 | "#elif defined(__riscv) && defined(__riscv_v_min_vlen) " |
| 247 | "&& (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)", |
| 248 | ["RVV", "VEC", "VEC16"], |
| 249 | fp8=False, |
| 250 | ) |
| 251 | header += _macro_block( |
| 252 | "#elif defined(__riscv)", |
| 253 | ["VEC", "VEC16"], |
| 254 | fp8=False, |
| 255 | ) |
| 256 | header += _macro_block( |
| 257 | "#elif defined(__powerpc__)", |
| 258 | ["VSX", "VEC", "VEC16"], |
| 259 | fp8=False, |
| 260 | ) |
| 261 | header += _macro_block( |
| 262 | "#elif defined(__AVX512F__)", |
| 263 | ["VEC", "VEC16"], |
| 264 | fp8=True, |
| 265 | ) |
| 266 | header += _macro_block( |
| 267 | "#elif defined(__AVX2__)", |
| 268 | ["VEC", "VEC16"], |
| 269 | fp8=False, |
| 270 | ) |
| 271 | header += _macro_block( |
| 272 | "#else", |
| 273 | ["VEC", "VEC16"], |
| 274 | fp8=False, |
| 275 | ) |
| 276 | header += ( |
| 277 | "#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ /" |
| 278 | " __riscv / __powerpc__ */\n\n" |
| 279 | "#endif // CPU_ATTN_DISPATCH_GENERATED_H\n" |
| 280 | ) |
| 281 | |
| 282 | return header |
| 283 | |
| 284 | |
| 285 | def main(): |
| 286 | output_path = os.path.join( |
| 287 | os.path.dirname(__file__), "cpu_attn_dispatch_generated.h" |
| 288 | ) |
| 289 | |
| 290 | with open(output_path, "w") as f: |
| 291 | f.write(generate_header_file()) |
| 292 | |
| 293 | |
| 294 | if __name__ == "__main__": |
| 295 | main() |