@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4dffc5e044
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 8, 22:45 UTC

[CPU] Split attention dispatch by head_dim alignment (#32161)

Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com>

R3hankhan committed Feb 4, 2026 at 09:07 UTC 4dffc5e044317326b9e2b2fd2a019c499d63c427
6 files changed +241 -107
cmake/cpu_extension.cmake
+13
@@ -359,6 +359,19 @@ else()
359 add_compile_definitions(-DVLLM_NUMA_DISABLED)
360 endif()
361
362 +#
363 +# Generate CPU attention dispatch header
364 +#
365 +message(STATUS "Generating CPU attention dispatch header")
366 +execute_process(
367 + COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/csrc/cpu/generate_cpu_attn_dispatch.py
368 + WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}/csrc/cpu
369 + RESULT_VARIABLE GEN_RESULT
370 +)
371 +if(NOT GEN_RESULT EQUAL 0)
372 + message(FATAL_ERROR "Failed to generate CPU attention dispatch header")
373 +endif()
374 +
375 #
376 # _C extension
377 #
csrc/cpu/cpu_attn.cpp
+21 -104
@@ -1,79 +1,4 @@
1 -#include "cpu_attn_vec.hpp"
2 -#include "cpu_attn_vec16.hpp"
3 -
4 -#ifdef CPU_CAPABILITY_AMXBF16
5 - #include "cpu_attn_amx.hpp"
6 - #define AMX_DISPATCH(...) \
7 - case cpu_attention::ISA::AMX: { \
8 - using attn_impl = cpu_attention::AttentionImpl<cpu_attention::ISA::AMX, \
9 - scalar_t, head_dim>; \
10 - return __VA_ARGS__(); \
11 - }
12 -#else
13 - #define AMX_DISPATCH(...) case cpu_attention::ISA::AMX:
14 -#endif
15 -
16 -#ifdef __aarch64__
17 - #include "cpu_attn_neon.hpp"
18 - // NEON requires head_dim to be a multiple of 32
19 - #define NEON_DISPATCH(...) \
20 - case cpu_attention::ISA::NEON: { \
21 - using attn_impl = cpu_attention::AttentionImpl<cpu_attention::ISA::NEON, \
22 - scalar_t, head_dim>; \
23 - return __VA_ARGS__(); \
24 - }
25 -#else
26 - #define NEON_DISPATCH(...) case cpu_attention::ISA::NEON:
27 -#endif // #ifdef __aarch64__
28 -
29 -#define CPU_ATTN_DISPATCH_CASE(HEAD_DIM, ...) \
30 - case HEAD_DIM: { \
31 - constexpr size_t head_dim = HEAD_DIM; \
32 - return __VA_ARGS__(); \
33 - }
34 -
35 -#define CPU_ATTN_DISPATCH_CASE_HEADDIM(HEAD_DIM, ...) \
36 - [&] { \
37 - switch (HEAD_DIM) { \
38 - CPU_ATTN_DISPATCH_CASE(32, __VA_ARGS__) \
39 - CPU_ATTN_DISPATCH_CASE(64, __VA_ARGS__) \
40 - CPU_ATTN_DISPATCH_CASE(80, __VA_ARGS__) \
41 - CPU_ATTN_DISPATCH_CASE(96, __VA_ARGS__) \
42 - CPU_ATTN_DISPATCH_CASE(112, __VA_ARGS__) \
43 - CPU_ATTN_DISPATCH_CASE(128, __VA_ARGS__) \
44 - CPU_ATTN_DISPATCH_CASE(160, __VA_ARGS__) \
45 - CPU_ATTN_DISPATCH_CASE(192, __VA_ARGS__) \
46 - CPU_ATTN_DISPATCH_CASE(224, __VA_ARGS__) \
47 - CPU_ATTN_DISPATCH_CASE(256, __VA_ARGS__) \
48 - default: { \
49 - TORCH_CHECK(false, "Invalid CPU attention head_dim: " + \
50 - std::to_string(HEAD_DIM)); \
51 - } \
52 - } \
53 - }()
54 -
55 -#define CPU_ATTN_DISPATCH_IMPL(ISA_TYPE, ...) \
56 - [&] { \
57 - switch (ISA_TYPE) { \
58 - AMX_DISPATCH(__VA_ARGS__) \
59 - NEON_DISPATCH(__VA_ARGS__) \
60 - case cpu_attention::ISA::VEC: { \
61 - using attn_impl = \
62 - cpu_attention::AttentionImpl<cpu_attention::ISA::VEC, scalar_t, \
63 - head_dim>; \
64 - return __VA_ARGS__(); \
65 - } \
66 - case cpu_attention::ISA::VEC16: { \
67 - using attn_impl = \
68 - cpu_attention::AttentionImpl<cpu_attention::ISA::VEC16, scalar_t, \
69 - head_dim>; \
70 - return __VA_ARGS__(); \
71 - } \
72 - default: { \
73 - TORCH_CHECK(false, "Invalid CPU attention ISA type."); \
74 - } \
75 - } \
76 - }()
1 +#include "cpu_attn_dispatch_generated.h"
2
3 torch::Tensor get_scheduler_metadata(
4 const int64_t num_req, const int64_t num_heads_q,
@@ -122,16 +47,14 @@ torch::Tensor get_scheduler_metadata(
47 input.enable_kv_split = enable_kv_split;
48
49 VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
125 - CPU_ATTN_DISPATCH_CASE_HEADDIM(head_dim, [&] {
126 - CPU_ATTN_DISPATCH_IMPL(isa, [&]() {
127 - input.elem_size = sizeof(scalar_t);
128 - input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t);
129 - input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t);
130 - input.output_buffer_elem_size =
131 - sizeof(attn_impl::partial_output_buffer_t);
132 - input.max_num_q_per_iter = attn_impl::MaxQHeadNumPerIteration;
133 - input.kv_block_alignment = attn_impl::BlockSizeAlignment;
134 - });
50 + CPU_ATTN_DISPATCH(head_dim, isa, [&]() {
51 + input.elem_size = sizeof(scalar_t);
52 + input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t);
53 + input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t);
54 + input.output_buffer_elem_size =
55 + sizeof(attn_impl::partial_output_buffer_t);
56 + input.max_num_q_per_iter = attn_impl::MaxQHeadNumPerIteration;
57 + input.kv_block_alignment = attn_impl::BlockSizeAlignment;
58 });
59 });
60
@@ -184,18 +107,14 @@ void cpu_attn_reshape_and_cache(
107
108 VLLM_DISPATCH_FLOATING_TYPES(
109 key.scalar_type(), "cpu_attn_reshape_and_cache", [&]() {
187 - CPU_ATTN_DISPATCH_CASE_HEADDIM(head_dim, [&] {
188 - CPU_ATTN_DISPATCH_IMPL(isa_tag, [&]() {
189 - attn_impl::reshape_and_cache(
190 - key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
191 - key_cache.data_ptr<scalar_t>(),
192 - value_cache.data_ptr<scalar_t>(),
193 - slot_mapping.data_ptr<int64_t>(), token_num,
194 - key_token_num_stride, value_token_num_stride, head_num,
195 - key_head_num_stride, value_head_num_stride, num_blocks,
196 - num_blocks_stride, cache_head_num_stride, block_size,
197 - block_size_stride);
198 - });
110 + CPU_ATTN_DISPATCH(head_dim, isa_tag, [&]() {
111 + attn_impl::reshape_and_cache(
112 + key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
113 + key_cache.data_ptr<scalar_t>(), value_cache.data_ptr<scalar_t>(),
114 + slot_mapping.data_ptr<int64_t>(), token_num, key_token_num_stride,
115 + value_token_num_stride, head_num, key_head_num_stride,
116 + value_head_num_stride, num_blocks, num_blocks_stride,
117 + cache_head_num_stride, block_size, block_size_stride);
118 });
119 });
120 }
@@ -257,12 +176,10 @@ void cpu_attention_with_kv_cache(
176
177 VLLM_DISPATCH_FLOATING_TYPES(
178 query.scalar_type(), "cpu_attention_with_kv_cache", [&]() {
260 - CPU_ATTN_DISPATCH_CASE_HEADDIM(query.size(2), [&] {
261 - CPU_ATTN_DISPATCH_IMPL(input.metadata->isa, [&]() {
262 - TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment, 0);
263 - cpu_attention::AttentionMainLoop<attn_impl> mainloop;
264 - mainloop(&input);
265 - });
179 + CPU_ATTN_DISPATCH(query.size(2), input.metadata->isa, [&]() {
180 + TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment, 0);
181 + cpu_attention::AttentionMainLoop<attn_impl> mainloop;
182 + mainloop(&input);
183 });
184 });
185 }
csrc/cpu/cpu_attn_amx.hpp
+1 -1
@@ -377,7 +377,7 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
377 const int32_t q_heads_per_kv, const int64_t q_num_stride,
378 const int64_t q_head_stride, const float scale) {
379 constexpr int64_t bytes_per_head = head_dim * sizeof(scalar_t);
380 - // static_assert(bytes_per_head % AMX_TILE_ROW_BYTES == 0);
380 + static_assert(bytes_per_head % AMX_TILE_ROW_BYTES == 0);
381 constexpr int64_t head_size_block_num = bytes_per_head / AMX_TILE_ROW_BYTES;
382 constexpr int64_t head_elem_num_pre_block =
383 AMX_TILE_ROW_BYTES / sizeof(scalar_t);
csrc/cpu/cpu_attn_neon.hpp
+1 -1
@@ -264,7 +264,7 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
264 constexpr static ISA ISAType = ISA::NEON;
265 constexpr static bool scale_on_logits = false; // apply scale on q_buffer
266
267 - // static_assert(HeadDim % HeadDimAlignment == 0);
267 + static_assert(HeadDim % HeadDimAlignment == 0);
268 // the gemm micro kernel is Mx8
269 static_assert(HeadDimAlignment % 8 == 0);
270 static_assert(BlockSizeAlignment % 8 == 0);
csrc/cpu/generate_cpu_attn_dispatch.py new
+203
@@ -0,0 +1,203 @@
1 +#!/usr/bin/env python3
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +"""
5 +Generate CPU attention dispatch switch cases and kernel instantiations.
6 +"""
7 +
8 +import os
9 +
10 +# Head dimensions divisible by 32 (support all ISAs)
11 +HEAD_DIMS_32 = [32, 64, 96, 128, 160, 192, 224, 256]
12 +
13 +# Head dimensions divisible by 16 but not 32 (VEC16 only)
14 +HEAD_DIMS_16 = [80, 112]
15 +
16 +# ISA types
17 +ISA_TYPES = {
18 + "AMX": 0,
19 + "VEC": 1,
20 + "VEC16": 2,
21 + "NEON": 3,
22 +}
23 +
24 +# ISAs supported for head_dims divisible by 32
25 +ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16"]
26 +
27 +# ISAs supported for head_dims divisible by 16 only
28 +ISA_FOR_16 = ["VEC16"]
29 +
30 +
31 +def encode_params(head_dim: int, isa_type: str) -> int:
32 + """Encode head_dim and ISA type into a single int64_t."""
33 + isa_val = ISA_TYPES[isa_type]
34 + # Encoding: (head_dim << 8) | isa_type
35 + # This allows head_dim up to 2^56 - 1 and 256 ISA types
36 + return (head_dim << 8) | isa_val
37 +
38 +
39 +def generate_cases_for_isa_group(isa_list: list[str]) -> str:
40 + """Generate switch cases for a specific ISA group."""
41 + cases = []
42 +
43 + # Generate cases for head_dims divisible by 32
44 + for head_dim in HEAD_DIMS_32:
45 + for isa in isa_list:
46 + if isa not in ISA_FOR_32:
47 + continue
48 + encoded = encode_params(head_dim, isa)
49 + case_str = (
50 + f""" case {encoded}LL: {{ """
51 + f"""/* head_dim={head_dim}, isa={isa} */ \\"""
52 + f"""
53 + constexpr size_t head_dim = {head_dim}; \\"""
54 + f"""
55 + using attn_impl = cpu_attention::AttentionImpl<"""
56 + f"""cpu_attention::ISA::{isa}, \\"""
57 + f"""
58 + """
59 + f"""scalar_t, head_dim>; \\"""
60 + f"""
61 + return __VA_ARGS__(); \\"""
62 + f"""
63 + }} \\"""
64 + )
65 + cases.append(case_str)
66 +
67 + # Generate cases for head_dims divisible by 16 only
68 + for head_dim in HEAD_DIMS_16:
69 + for isa in isa_list:
70 + encoded = encode_params(head_dim, isa)
71 + case_str = (
72 + f""" case {encoded}LL: {{ """
73 + f"""/* head_dim={head_dim}, isa={isa} """
74 + f"""(using VEC16) */ \\"""
75 + f"""
76 + constexpr size_t head_dim = {head_dim}; \\"""
77 + f"""
78 + using attn_impl = cpu_attention::AttentionImpl<"""
79 + f"""cpu_attention::ISA::VEC16, \\"""
80 + f"""
81 + """
82 + f"""scalar_t, head_dim>; \\"""
83 + f"""
84 + return __VA_ARGS__(); \\"""
85 + f"""
86 + }} \\"""
87 + )
88 + cases.append(case_str)
89 +
90 + return "\n".join(cases)
91 +
92 +
93 +def generate_helper_function() -> str:
94 + """Generate helper function to encode parameters."""
95 + return """
96 +inline int64_t encode_cpu_attn_params(int64_t head_dim, cpu_attention::ISA isa) {
97 + return (head_dim << 8) | static_cast<int64_t>(isa);
98 +}
99 +"""
100 +
101 +
102 +def generate_header_file() -> str:
103 + """Generate the complete header file content."""
104 + header = """// auto generated by generate_cpu_attn_dispatch.py
105 +// clang-format off
106 +
107 +#ifndef CPU_ATTN_DISPATCH_GENERATED_H
108 +#define CPU_ATTN_DISPATCH_GENERATED_H
109 +
110 +#include "cpu_attn_vec.hpp"
111 +#include "cpu_attn_vec16.hpp"
112 +
113 +#ifdef CPU_CAPABILITY_AMXBF16
114 + #include "cpu_attn_amx.hpp"
115 +#endif
116 +
117 +#ifdef __aarch64__
118 + #include "cpu_attn_neon.hpp"
119 +#endif
120 +
121 +"""
122 +
123 + header += generate_helper_function()
124 +
125 + # Generate dispatch macro with conditional compilation for different ISA sets
126 + header += """
127 +// Dispatch macro using encoded parameters
128 +"""
129 +
130 + # x86_64 with AMX
131 + header += """#if defined(CPU_CAPABILITY_AMXBF16)
132 +#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
133 + [&] { \\
134 + int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
135 + switch (encoded_params) { \\
136 +"""
137 + header += generate_cases_for_isa_group(["AMX", "VEC", "VEC16"])
138 + header += """
139 + default: { \\
140 + TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
141 + std::to_string(HEAD_DIM) + " isa=" + \\
142 + std::to_string(static_cast<int>(ISA_TYPE))); \\
143 + } \\
144 + } \\
145 + }()
146 +
147 +"""
148 +
149 + # ARM64 with NEON
150 + header += """#elif defined(__aarch64__)
151 +#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
152 + [&] { \\
153 + int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
154 + switch (encoded_params) { \\
155 +"""
156 + header += generate_cases_for_isa_group(["NEON", "VEC", "VEC16"])
157 + header += """
158 + default: { \\
159 + TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
160 + std::to_string(HEAD_DIM) + " isa=" + \\
161 + std::to_string(static_cast<int>(ISA_TYPE))); \\
162 + } \\
163 + } \\
164 + }()
165 +
166 +"""
167 +
168 + # Fallback: VEC and VEC16 only
169 + header += """#else
170 +#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
171 + [&] { \\
172 + int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
173 + switch (encoded_params) { \\
174 +"""
175 + header += generate_cases_for_isa_group(["VEC", "VEC16"])
176 + header += """
177 + default: { \\
178 + TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
179 + std::to_string(HEAD_DIM) + " isa=" + \\
180 + std::to_string(static_cast<int>(ISA_TYPE))); \\
181 + } \\
182 + } \\
183 + }()
184 +
185 +#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ */
186 +
187 +#endif // CPU_ATTN_DISPATCH_GENERATED_H
188 +"""
189 +
190 + return header
191 +
192 +
193 +def main():
194 + output_path = os.path.join(
195 + os.path.dirname(__file__), "cpu_attn_dispatch_generated.h"
196 + )
197 +
198 + with open(output_path, "w") as f:
199 + f.write(generate_header_file())
200 +
201 +
202 +if __name__ == "__main__":
203 + main()
tests/kernels/attention/test_cpu_attn.py
+2 -1
@@ -26,6 +26,7 @@ NUM_HEADS = [
26 (9, 3),
27 ]
28 HEAD_SIZES = [96, 128]
29 +HEAD_SIZES_VEC16 = [96, 80, 112, 128]
30 QTYPES = [torch.bfloat16, torch.half, torch.float32]
31 SLIDING_WINDOWS = [None, 256]
32 NUM_BLOCKS = [
@@ -432,7 +433,7 @@ def test_varlen_with_paged_kv_normal_amx(
433
434 @pytest.mark.parametrize("seq_lens", SEQ_LENS)
435 @pytest.mark.parametrize("num_heads", NUM_HEADS)
435 -@pytest.mark.parametrize("head_size", HEAD_SIZES)
436 +@pytest.mark.parametrize("head_size", HEAD_SIZES_VEC16)
437 @pytest.mark.parametrize("block_size", [48])
438 @pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
439 @pytest.mark.parametrize("dtype", [torch.bfloat16])