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 5, 16:00 UTC
[CPU] Support head_size 512 in cpu_attn (#38676)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Apr 1, 2026 at 13:42 UTC
36d7f19897843c9cbdb701ba88d0f2c29954fe44
4 files changed
+4
-4
csrc/cpu/generate_cpu_attn_dispatch.py
+1
-1
@@ -8,7 +8,7 @@ Generate CPU attention dispatch switch cases and kernel instantiations.
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]
11
+HEAD_DIMS_32 = [32, 64, 96, 128, 160, 192, 224, 256, 512]
12
13
# Head dimensions divisible by 16 but not 32 (VEC16 only)
14
HEAD_DIMS_16 = [80, 112]
docs/design/attention_backends.md
+1
-1
@@ -165,7 +165,7 @@ Priority is **1 = highest** (tried first).
165
166
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
167
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
168
-| `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | All | N/A |
168
+| `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
169
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
170
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
171
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
tests/kernels/attention/test_cpu_attn.py
+1
-1
@@ -25,7 +25,7 @@ NUM_HEADS = [
25
(8, 2),
26
(9, 3),
27
]
28
-HEAD_SIZES = [96, 128]
28
+HEAD_SIZES = [96, 128, 512]
29
HEAD_SIZES_VEC16 = [96, 80, 112, 128]
30
QTYPES = [torch.bfloat16, torch.half, torch.float32]
31
SLIDING_WINDOWS = [None, 256]
vllm/v1/attention/backends/cpu_attn.py
+1
-1
@@ -38,7 +38,7 @@ class CPUAttentionBackend(AttentionBackend):
38
39
@classmethod
40
def get_supported_head_sizes(cls) -> list[int]:
41
- return [32, 64, 80, 96, 112, 128, 160, 192, 224, 256]
41
+ return [32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512]
42
43
@staticmethod
44
def get_name() -> str: