wehub-resource-sync
277 lines
8.1 KiB
Python
277 lines
8.1 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""
|
|
Generate CPU attention dispatch switch cases and kernel instantiations.
|
|
"""
|
|
|
|
import os
|
|
|
|
# Head dimensions divisible by 32 (support all ISAs)
|
|
HEAD_DIMS_32 = [32, 64, 96, 128, 160, 192, 224, 256, 512]
|
|
|
|
# Head dimensions divisible by 16 but not 32 (VEC16 only)
|
|
HEAD_DIMS_16 = [48, 80, 112]
|
|
|
|
# ISA types
|
|
ISA_TYPES = {
|
|
"AMX": 0,
|
|
"VEC": 1,
|
|
"VEC16": 2,
|
|
"NEON": 3,
|
|
"VXE": 4,
|
|
"RVV": 5,
|
|
"VSX": 6,
|
|
}
|
|
|
|
# KV cache index: 0 = auto (same as scalar_t), 1 = fp8_e4m3, 2 = fp8_e5m2
|
|
KV_CACHE_IDX = {
|
|
"auto": 0,
|
|
"fp8_e4m3": 1,
|
|
"fp8_e5m2": 2,
|
|
}
|
|
|
|
# C++ type for each kv_cache index
|
|
KV_CACHE_CPP_TYPES = {
|
|
"auto": "scalar_t",
|
|
"fp8_e4m3": "c10::Float8_e4m3fn",
|
|
"fp8_e5m2": "c10::Float8_e5m2",
|
|
}
|
|
|
|
# ISAs supported for head_dims divisible by 32
|
|
ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE", "RVV", "VSX"]
|
|
|
|
# ISAs supported for head_dims divisible by 16 only
|
|
ISA_FOR_16 = ["VEC16"]
|
|
|
|
# ISAs that support FP8 KV cache (x86 AVX2/AVX-512 required)
|
|
ISA_FOR_FP8 = ["AMX", "VEC"]
|
|
|
|
|
|
def encode_params(head_dim: int, isa_type: str, kv_cache: str = "auto") -> int:
|
|
"""Encode head_dim, ISA type, and KV cache type into a single int64_t."""
|
|
isa_val = ISA_TYPES[isa_type]
|
|
kv_val = KV_CACHE_IDX[kv_cache]
|
|
# Encoding: (head_dim << 16) | (kv_cache_idx << 8) | isa_type
|
|
# This allows head_dim up to 2^48 - 1, 256 KV cache types, and 256 ISA types
|
|
return (head_dim << 16) | (kv_val << 8) | isa_val
|
|
|
|
|
|
def _make_case(
|
|
head_dim: int, isa: str, kv_cache: str = "auto", isa_override: str | None = None
|
|
) -> str:
|
|
"""Generate a single switch case line."""
|
|
encoded = encode_params(head_dim, isa, kv_cache)
|
|
actual_isa = isa_override if isa_override else isa
|
|
cpp_type = KV_CACHE_CPP_TYPES[kv_cache]
|
|
attn_impl = (
|
|
f"cpu_attention::AttentionImpl<"
|
|
f"cpu_attention::ISA::{actual_isa}, \\\n"
|
|
f" "
|
|
f"scalar_t, head_dim, {cpp_type}>"
|
|
)
|
|
comment = (
|
|
f"head_dim={head_dim}, isa={isa}"
|
|
if kv_cache == "auto"
|
|
else f"head_dim={head_dim}, isa={isa}, kv_cache={kv_cache}"
|
|
)
|
|
return (
|
|
f""" case {encoded}LL: {{ """
|
|
f"""/* {comment} */ \\"""
|
|
f"""
|
|
constexpr size_t head_dim = {head_dim}; \\"""
|
|
f"""
|
|
using attn_impl = {attn_impl}; \\"""
|
|
f"""
|
|
return __VA_ARGS__(); \\"""
|
|
f"""
|
|
}} \\"""
|
|
)
|
|
|
|
|
|
def generate_cases_for_isa_group(isa_list: list[str], include_fp8: bool = False) -> str:
|
|
"""Generate switch cases for a specific ISA group."""
|
|
cases = []
|
|
|
|
# Non-FP8 cases for head_dims divisible by 32
|
|
for head_dim in HEAD_DIMS_32:
|
|
for isa in isa_list:
|
|
if isa not in ISA_FOR_32:
|
|
continue
|
|
cases.append(_make_case(head_dim, isa, "auto"))
|
|
|
|
# Non-FP8 cases for head_dims divisible by 16 only
|
|
for head_dim in HEAD_DIMS_16:
|
|
for isa in isa_list:
|
|
cases.append(_make_case(head_dim, isa, "auto", isa_override="VEC16"))
|
|
|
|
# FP8 cases: only AMX and VEC, only head_dims divisible by 32
|
|
if include_fp8:
|
|
for fp8_type in ("fp8_e4m3", "fp8_e5m2"):
|
|
for head_dim in HEAD_DIMS_32:
|
|
for isa in isa_list:
|
|
if isa not in ISA_FOR_FP8:
|
|
continue
|
|
cases.append(_make_case(head_dim, isa, fp8_type))
|
|
|
|
return "\n".join(cases)
|
|
|
|
|
|
def generate_helper_function() -> str:
|
|
"""Generate helper function to encode parameters."""
|
|
return """
|
|
inline int64_t encode_cpu_attn_params(int64_t head_dim, cpu_attention::ISA isa,
|
|
int64_t kv_cache_idx = 0) {
|
|
return (head_dim << 16) | (kv_cache_idx << 8) | static_cast<int64_t>(isa);
|
|
}
|
|
"""
|
|
|
|
|
|
def generate_header_file() -> str:
|
|
"""Generate the complete header file content."""
|
|
header = """// auto generated by generate_cpu_attn_dispatch.py
|
|
// clang-format off
|
|
|
|
#ifndef CPU_ATTN_DISPATCH_GENERATED_H
|
|
#define CPU_ATTN_DISPATCH_GENERATED_H
|
|
|
|
#include "cpu_attn_vec.hpp"
|
|
#include "cpu_attn_vec16.hpp"
|
|
|
|
#ifdef CPU_CAPABILITY_AMXBF16
|
|
#include "cpu_attn_amx.hpp"
|
|
#endif
|
|
|
|
#ifdef __aarch64__
|
|
#include "cpu_attn_neon.hpp"
|
|
#endif
|
|
|
|
#ifdef __s390x__
|
|
#include "cpu_attn_vxe.hpp"
|
|
#endif
|
|
|
|
// cpu_attn_rvv.hpp supports VLEN=128 and VLEN=256 via RVVI() macros.
|
|
// Other VLENs and scalar RISC-V builds skip it entirely.
|
|
#if defined(__riscv) && defined(__riscv_v_min_vlen) && \
|
|
(__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)
|
|
#include "cpu_attn_rvv.hpp"
|
|
#endif
|
|
|
|
#ifdef __powerpc__
|
|
#include "cpu_attn_vsx.hpp"
|
|
#endif
|
|
|
|
"""
|
|
|
|
header += generate_helper_function()
|
|
|
|
# Generate dispatch macro with conditional compilation for different ISA sets
|
|
header += """
|
|
// Dispatch macro using encoded parameters.
|
|
// KV_CACHE_IDX: Fp8KVCacheDataType enum value (kAuto=0, kFp8E4M3=1, kFp8E5M2=2).
|
|
// FP8 cases (kv_cache_idx != 0) are generated on x86 platforms with AVX2 or
|
|
// AVX-512: BF16Vec32 FP8 constructors have both AVX-512 and AVX2 implementations
|
|
// in cpu_types_x86.hpp. Non-x86 platforms (#else fallback) have fp8=False.
|
|
"""
|
|
|
|
def _macro_block(guard: str, isa_list: list[str], fp8: bool) -> str:
|
|
"""Return one CPU_ATTN_DISPATCH macro block for a given guard."""
|
|
enc = (
|
|
" int64_t encoded_params = encode_cpu_attn_params("
|
|
"HEAD_DIM, ISA_TYPE, KV_CACHE_IDX); \\"
|
|
)
|
|
cases = generate_cases_for_isa_group(isa_list, include_fp8=fp8)
|
|
tail = (
|
|
"\n"
|
|
" default: { \\\n"
|
|
" TORCH_CHECK(false, "
|
|
'"Unsupported CPU attention configuration: head_dim=" + \\\n'
|
|
' std::to_string(HEAD_DIM) + " isa=" + \\\n'
|
|
" std::to_string(static_cast<int>(ISA_TYPE))"
|
|
" + \\\n"
|
|
' " kv_cache_idx=" + '
|
|
"std::to_string(KV_CACHE_IDX)); \\\n"
|
|
" } \\\n"
|
|
" } \\\n"
|
|
" }()\n\n"
|
|
)
|
|
return (
|
|
f"{guard}\n"
|
|
"#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, KV_CACHE_IDX, ...) \\\n"
|
|
" [&] { \\\n"
|
|
f"{enc}\n"
|
|
" switch (encoded_params) { \\\n"
|
|
f"{cases}"
|
|
f"{tail}"
|
|
)
|
|
|
|
header += _macro_block(
|
|
"#if defined(CPU_CAPABILITY_AMXBF16)",
|
|
["AMX", "VEC", "VEC16"],
|
|
fp8=True,
|
|
)
|
|
header += _macro_block(
|
|
"#elif defined(__aarch64__)",
|
|
["NEON", "VEC", "VEC16"],
|
|
fp8=False,
|
|
)
|
|
header += _macro_block(
|
|
"#elif defined(__s390x__)",
|
|
["VXE", "VEC", "VEC16"],
|
|
fp8=False,
|
|
)
|
|
# RISC-V with RVV. cpu_attn_rvv.hpp supports VLEN=128 and VLEN=256
|
|
# via RVVI() macros. Builds with a supported VLEN get
|
|
# RVV+VEC+VEC16; other RISC-V builds fall back to VEC/VEC16 only.
|
|
header += _macro_block(
|
|
"#elif defined(__riscv) && defined(__riscv_v_min_vlen) "
|
|
"&& (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)",
|
|
["RVV", "VEC", "VEC16"],
|
|
fp8=False,
|
|
)
|
|
header += _macro_block(
|
|
"#elif defined(__riscv)",
|
|
["VEC", "VEC16"],
|
|
fp8=False,
|
|
)
|
|
header += _macro_block(
|
|
"#elif defined(__powerpc__)",
|
|
["VSX", "VEC", "VEC16"],
|
|
fp8=False,
|
|
)
|
|
header += _macro_block(
|
|
"#elif defined(__AVX512F__)",
|
|
["VEC", "VEC16"],
|
|
fp8=True,
|
|
)
|
|
header += _macro_block(
|
|
"#elif defined(__AVX2__)",
|
|
["VEC", "VEC16"],
|
|
fp8=False,
|
|
)
|
|
header += _macro_block(
|
|
"#else",
|
|
["VEC", "VEC16"],
|
|
fp8=False,
|
|
)
|
|
header += (
|
|
"#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ /"
|
|
" __riscv / __powerpc__ */\n\n"
|
|
"#endif // CPU_ATTN_DISPATCH_GENERATED_H\n"
|
|
)
|
|
|
|
return header
|
|
|
|
|
|
def main():
|
|
output_path = os.path.join(
|
|
os.path.dirname(__file__), "cpu_attn_dispatch_generated.h"
|
|
)
|
|
|
|
with open(output_path, "w") as f:
|
|
f.write(generate_header_file())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|