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 10, 10:15 UTC
[CPU] Add head sizes 80 and 112 with vec16 fallback (#31968)
Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com>
R3hankhan committed
Jan 9, 2026 at 19:44 UTC
8e27663b6a2ba9d4d0c35e592c118aaf7b50fd13
4 files changed
+12
-5
csrc/cpu/cpu_attn.cpp
+3
@@ -15,6 +15,7 @@
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, \
@@ -36,7 +37,9 @@
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__) \
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);
vllm/v1/attention/backends/cpu_attn.py
+7
-3
@@ -42,7 +42,7 @@ class CPUAttentionBackend(AttentionBackend):
42
43
@classmethod
44
def get_supported_head_sizes(cls) -> list[int]:
45
- return [32, 64, 96, 128, 160, 192, 224, 256]
45
+ return [32, 64, 80, 96, 112, 128, 160, 192, 224, 256]
46
47
@staticmethod
48
def get_name() -> str:
@@ -137,7 +137,7 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
137
if self.window_size is None:
138
self.window_size = -1
139
self.block_size = vllm_config.cache_config.block_size
140
- self.isa = _get_attn_isa(self.dtype, self.block_size)
140
+ self.isa = _get_attn_isa(self.dtype, self.block_size, self.head_dim)
141
self.is_cross_attention = isinstance(kv_cache_spec, CrossAttentionSpec)
142
143
def build(
@@ -484,7 +484,11 @@ def _make_sliding_window_bias(
484
return attn_biases
485
486
487
-def _get_attn_isa(dtype: torch.dtype, block_size: int) -> str:
487
+def _get_attn_isa(
488
+ dtype: torch.dtype, block_size: int, head_size: int | None = None
489
+) -> str:
490
+ if head_size is not None and head_size % 32 != 0 and head_size % 16 == 0:
491
+ return "vec16"
492
supports_amx = torch._C._cpu._is_amx_tile_supported()
493
if supports_amx and dtype in (torch.bfloat16,) and block_size % 32 == 0:
494
return "amx"