Woosuk KwonandClaude Opus 5.5
e1a5086256
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
559 lines
19 KiB
Python
559 lines
19 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import functools
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, ClassVar
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.config.cache import CacheDType
|
|
|
|
import torch
|
|
|
|
from vllm import _custom_ops as ops
|
|
from vllm import envs
|
|
from vllm.config import (
|
|
VllmConfig,
|
|
get_current_vllm_config,
|
|
get_layers_from_vllm_config,
|
|
)
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.layers.attention import Attention
|
|
from vllm.platforms import CpuArchEnum, current_platform
|
|
from vllm.utils.torch_utils import is_quantized_kv_cache
|
|
from vllm.v1.attention.backend import (
|
|
AttentionBackend,
|
|
AttentionImpl,
|
|
AttentionLayer,
|
|
AttentionMetadataBuilder,
|
|
AttentionType,
|
|
CommonAttentionMetadata,
|
|
MultipleOf,
|
|
)
|
|
from vllm.v1.attention.backends.utils import (
|
|
get_num_attention_heads_from_layers,
|
|
)
|
|
from vllm.v1.attention.backends.zentorch_sdpa import (
|
|
should_use_zentorch_sdpa,
|
|
zentorch_sdpa_attn,
|
|
)
|
|
from vllm.v1.kv_cache_interface import (
|
|
AttentionSpec,
|
|
CrossAttentionSpec,
|
|
EncoderOnlyAttentionSpec,
|
|
KVCacheLayout,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class CPUAttentionBackend(AttentionBackend):
|
|
forward_includes_kv_cache_update: bool = False
|
|
|
|
supported_dtypes: ClassVar[list[torch.dtype]] = [
|
|
torch.float16,
|
|
torch.bfloat16,
|
|
torch.float32,
|
|
]
|
|
supported_kv_cache_dtypes: ClassVar[list["CacheDType"]] = [
|
|
"auto",
|
|
"fp8",
|
|
"fp8_e4m3",
|
|
"fp8_e5m2",
|
|
]
|
|
|
|
@staticmethod
|
|
def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
|
|
return [MultipleOf(32)]
|
|
|
|
@classmethod
|
|
def get_supported_head_sizes(cls) -> list[int]:
|
|
return [32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512]
|
|
|
|
@classmethod
|
|
def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
|
|
# The CPU backend only reads head-major block interiors.
|
|
return (KVCacheLayout.LBHNC,)
|
|
|
|
@staticmethod
|
|
def get_name() -> str:
|
|
return "CPU_ATTN"
|
|
|
|
@classmethod
|
|
def supports_non_causal(cls) -> bool:
|
|
return True
|
|
|
|
@classmethod
|
|
def supports_sliding_window(cls) -> bool:
|
|
return True
|
|
|
|
@classmethod
|
|
def supports_attn_type(cls, attn_type: str) -> bool:
|
|
"""CPU attention supports decoder,
|
|
encoder-only and encoder-decoder attention."""
|
|
return attn_type in (
|
|
AttentionType.DECODER,
|
|
AttentionType.ENCODER,
|
|
AttentionType.ENCODER_ONLY,
|
|
AttentionType.ENCODER_DECODER,
|
|
)
|
|
|
|
@staticmethod
|
|
def get_impl_cls() -> type["CPUAttentionBackendImpl"]:
|
|
return CPUAttentionBackendImpl
|
|
|
|
@staticmethod
|
|
def get_builder_cls() -> type["CPUAttentionMetadataBuilder"]:
|
|
return CPUAttentionMetadataBuilder
|
|
|
|
|
|
@dataclass
|
|
class CPUAttentionMetadata:
|
|
num_actual_tokens: int # Number of tokens excluding padding.
|
|
max_query_len: int
|
|
query_start_loc: torch.Tensor
|
|
max_seq_len: int
|
|
seq_lens: torch.Tensor
|
|
block_table: torch.Tensor
|
|
slot_mapping: torch.Tensor
|
|
scheduler_metadata: torch.Tensor | None
|
|
causal: bool = True
|
|
dynamic_causal: torch.Tensor | None = None
|
|
|
|
# can be removed after deprecate sdpa
|
|
use_sdpa_prefill: bool = False
|
|
num_decode_tokens: int = 0
|
|
sdpa_attn_masks: list[torch.Tensor | None] | None = None
|
|
sdpa_start_loc: torch.Tensor | None = None
|
|
|
|
encoder_cache: torch.Tensor | None = None
|
|
|
|
|
|
class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]):
|
|
def __init__(
|
|
self,
|
|
kv_cache_spec: AttentionSpec,
|
|
layer_names: list[str],
|
|
vllm_config: VllmConfig,
|
|
device: torch.device,
|
|
) -> None:
|
|
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
|
|
|
|
self.kv_cache_spec = kv_cache_spec
|
|
self.vllm_config = vllm_config
|
|
|
|
parallel_config = vllm_config.parallel_config
|
|
self.num_kv_heads = kv_cache_spec.num_kv_heads
|
|
# The scheduler metadata built here sizes a scratchpad from the query
|
|
# head count, so it must come from this group's layers: the model-wide
|
|
# count is wrong for models that vary it per layer (e.g. Laguna).
|
|
self.num_heads = get_num_attention_heads_from_layers(
|
|
vllm_config, layer_names
|
|
) or vllm_config.model_config.get_num_attention_heads(parallel_config)
|
|
self.head_dim = kv_cache_spec.head_size
|
|
self.dtype = vllm_config.model_config.dtype
|
|
self.window_size = self._group_sliding_window()
|
|
self.block_size = kv_cache_spec.block_size
|
|
self.kv_cache_dtype = vllm_config.cache_config.cache_dtype
|
|
self.isa = _get_attn_isa(
|
|
self.dtype,
|
|
self.block_size,
|
|
self.head_dim,
|
|
self.kv_cache_dtype,
|
|
)
|
|
self._set_isa_to_layers(layer_names)
|
|
self.is_cross_attention = isinstance(kv_cache_spec, CrossAttentionSpec)
|
|
self.is_encoder_only_attention = isinstance(
|
|
kv_cache_spec, EncoderOnlyAttentionSpec
|
|
)
|
|
|
|
def _set_isa_to_layers(self, layer_names: list[str]) -> None:
|
|
attn_layers = get_layers_from_vllm_config(
|
|
self.vllm_config,
|
|
Attention,
|
|
layer_names,
|
|
)
|
|
for layer in attn_layers.values():
|
|
layer.isa = self.isa # type: ignore
|
|
|
|
def _group_sliding_window(self) -> int:
|
|
"""The window shared by every layer in this group, else -1 (no window).
|
|
|
|
Taken from the layers rather than the group spec: one KV cache group can
|
|
hold both windowed and global layers (e.g. Gemma-3 with the hybrid KV
|
|
cache manager disabled), and the scheduler metadata built here is shared
|
|
by the whole group, so it may only assume a window all of them agree on.
|
|
"""
|
|
layers = get_layers_from_vllm_config(
|
|
self.vllm_config, Attention, self.layer_names
|
|
)
|
|
windows = {
|
|
layer.impl.sliding_window
|
|
for layer in layers.values()
|
|
if isinstance(layer.impl, CPUAttentionBackendImpl)
|
|
}
|
|
if len(windows) != 1:
|
|
return -1
|
|
window = windows.pop()
|
|
return -1 if window is None else window
|
|
|
|
def build(
|
|
self,
|
|
common_prefix_len: int,
|
|
common_attn_metadata: CommonAttentionMetadata,
|
|
fast_build: bool = False,
|
|
) -> CPUAttentionMetadata:
|
|
num_reqs = common_attn_metadata.num_reqs
|
|
num_actual_tokens = common_attn_metadata.num_actual_tokens
|
|
max_query_len = common_attn_metadata.max_query_len
|
|
max_seq_len = common_attn_metadata.max_seq_len
|
|
query_start_loc = common_attn_metadata.query_start_loc
|
|
seq_lens = common_attn_metadata.seq_lens
|
|
block_table_tensor = common_attn_metadata.block_table_tensor
|
|
slot_mapping = common_attn_metadata.slot_mapping
|
|
is_dynamic_casual = isinstance(common_attn_metadata.causal, torch.Tensor)
|
|
dynamic_casual = None
|
|
if is_dynamic_casual:
|
|
dynamic_casual = common_attn_metadata.causal
|
|
|
|
causal = (
|
|
False
|
|
if self.is_cross_attention or is_dynamic_casual
|
|
else common_attn_metadata.causal
|
|
)
|
|
|
|
encoder_cache_tensor = None
|
|
if self.is_encoder_only_attention:
|
|
block_nums = (seq_lens + self.block_size - 1) // self.block_size
|
|
start_block_ids = torch.zeros_like(seq_lens)
|
|
torch.cumsum(block_nums[:-1], 0, out=start_block_ids[1:])
|
|
total_block_num: int = block_nums.sum().item()
|
|
max_block_num = block_nums.max().item()
|
|
block_offsets = torch.arange(
|
|
0, max_block_num, dtype=block_table_tensor.dtype
|
|
)
|
|
encoder_block_table = start_block_ids[:, None] + block_offsets[None, :]
|
|
torch.ops._C.compute_slot_mapping_kernel_impl(
|
|
query_start_loc,
|
|
common_attn_metadata.positions,
|
|
encoder_block_table,
|
|
slot_mapping,
|
|
self.block_size,
|
|
)
|
|
encoder_cache_tensor = torch.zeros(
|
|
(
|
|
total_block_num,
|
|
self.num_kv_heads,
|
|
self.block_size,
|
|
2 * self.head_dim,
|
|
),
|
|
dtype=self.dtype,
|
|
)
|
|
block_table_tensor = encoder_block_table
|
|
|
|
scheduler_metadata = ops.cpu_attn_get_scheduler_metadata(
|
|
num_reqs=num_reqs,
|
|
num_heads=self.num_heads,
|
|
num_kv_heads=self.num_kv_heads,
|
|
head_dim=self.head_dim,
|
|
seq_lens=seq_lens,
|
|
dtype=self.dtype,
|
|
query_start_loc=query_start_loc,
|
|
causal=causal,
|
|
sliding_window_size=self.window_size,
|
|
isa=self.isa,
|
|
enable_kv_split=envs.VLLM_CPU_ATTN_SPLIT_KV,
|
|
dynamic_causal=dynamic_casual,
|
|
kv_cache_dtype=self.kv_cache_dtype,
|
|
)
|
|
|
|
attn_metadata = CPUAttentionMetadata(
|
|
num_actual_tokens=num_actual_tokens,
|
|
max_query_len=max_query_len,
|
|
query_start_loc=query_start_loc,
|
|
max_seq_len=max_seq_len,
|
|
seq_lens=seq_lens,
|
|
block_table=block_table_tensor,
|
|
slot_mapping=slot_mapping,
|
|
scheduler_metadata=scheduler_metadata,
|
|
causal=causal,
|
|
encoder_cache=encoder_cache_tensor,
|
|
dynamic_causal=dynamic_casual,
|
|
)
|
|
|
|
return attn_metadata
|
|
|
|
|
|
class CPUAttentionBackendImpl(AttentionImpl):
|
|
def __init__(
|
|
self,
|
|
num_heads: int,
|
|
head_size: int,
|
|
scale: float,
|
|
num_kv_heads: int,
|
|
alibi_slopes: list[float] | None,
|
|
sliding_window: int | None,
|
|
kv_cache_dtype: str,
|
|
logits_soft_cap: float | None = None,
|
|
attn_type: str = AttentionType.DECODER,
|
|
kv_sharing_target_layer_name: str | None = None,
|
|
sinks: torch.Tensor | None = None,
|
|
) -> None:
|
|
self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
|
|
self.num_heads = num_heads
|
|
self.head_size = head_size
|
|
self.scale = float(scale)
|
|
if logits_soft_cap is not None and attn_type in (
|
|
AttentionType.ENCODER,
|
|
AttentionType.ENCODER_ONLY,
|
|
):
|
|
logger.warning_once(
|
|
"CPU_ATTN does not support logits softcap for"
|
|
" ENCODER and ENCODER_ONLY, outputs may be slightly off"
|
|
)
|
|
if logits_soft_cap is None:
|
|
logits_soft_cap = 0
|
|
self.logits_soft_cap = logits_soft_cap
|
|
|
|
self.num_kv_heads = num_kv_heads
|
|
if alibi_slopes is not None:
|
|
alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
|
|
self.alibi_slopes = alibi_slopes
|
|
if sliding_window is None:
|
|
self.sliding_window = -1
|
|
else:
|
|
self.sliding_window = sliding_window
|
|
self.kv_cache_dtype = kv_cache_dtype
|
|
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
|
|
|
|
self.is_fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
|
|
self.attn_type = attn_type
|
|
|
|
self.sinks = sinks
|
|
if self.sinks is not None:
|
|
assert self.sinks.shape[0] == num_heads, (
|
|
"Sinks must have the same number of heads as the number of "
|
|
"heads in the layer"
|
|
)
|
|
|
|
vllm_config = get_current_vllm_config()
|
|
self.use_zentorch_sdpa = should_use_zentorch_sdpa(
|
|
attn_type,
|
|
vllm_config.model_config.dtype,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
layer: AttentionLayer,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
kv_cache: torch.Tensor,
|
|
attn_metadata: CPUAttentionMetadata | None,
|
|
output: torch.Tensor,
|
|
output_scale: torch.Tensor | None = None,
|
|
output_block_scale: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
"""Forward pass for CPU attention backend.
|
|
|
|
Args:
|
|
layer: The attention layer, providing the q/k/v quantization scales.
|
|
query: shape = [num_tokens, num_heads, head_size]
|
|
key: shape = [num_tokens, num_kv_heads, head_size]
|
|
value: shape = [num_tokens, num_kv_heads, head_size]
|
|
kv_cache: shape =
|
|
[num_blocks, num_kv_heads, block_size, 2 * head_size]
|
|
attn_metadata: Metadata for attention.
|
|
output: Tensor that the attention result is written into.
|
|
output_scale: Scale for fused output quantization; not supported
|
|
by this backend.
|
|
output_block_scale: Block scale for fused output quantization;
|
|
not supported by this backend.
|
|
|
|
Returns:
|
|
shape = [num_tokens, num_heads * head_size]
|
|
|
|
"""
|
|
if output_scale is not None or output_block_scale is not None:
|
|
raise NotImplementedError(
|
|
"fused output quantization is not yet supported"
|
|
" for CPUAttentionBackendImpl"
|
|
)
|
|
|
|
# For warming-up
|
|
if attn_metadata is None:
|
|
return output
|
|
|
|
num_actual_tokens = attn_metadata.num_actual_tokens
|
|
|
|
is_encoder_attention = self.attn_type in (
|
|
AttentionType.ENCODER_ONLY,
|
|
AttentionType.ENCODER,
|
|
)
|
|
if is_encoder_attention:
|
|
if self.use_zentorch_sdpa:
|
|
# Encoder attention never reads the KV cache back, so the
|
|
# zentorch path attends the packed QKV directly instead of
|
|
# staging it through the scratch encoder cache.
|
|
zentorch_sdpa_attn(
|
|
query[:num_actual_tokens],
|
|
key[:num_actual_tokens],
|
|
value[:num_actual_tokens],
|
|
output[:num_actual_tokens],
|
|
attn_metadata,
|
|
self.scale,
|
|
self.sliding_window,
|
|
self.alibi_slopes,
|
|
)
|
|
return output
|
|
# For encoder attention,
|
|
kv_cache = attn_metadata.encoder_cache
|
|
|
|
# KV cache size are [num_blocks, num_kv_heads, block_size,
|
|
# 2 * head_size]. Make a view [num_blocks, num_kv_heads,
|
|
# block_size * 2, head_size]. Then slice KV at dim 2
|
|
num_blocks, num_kv_heads, block_size, _ = kv_cache.size()
|
|
kv_cache = kv_cache.view((num_blocks, num_kv_heads, block_size * 2, -1))
|
|
key_cache, value_cache = kv_cache.chunk(2, dim=2)
|
|
|
|
if is_encoder_attention:
|
|
ops.cpu_attn_reshape_and_cache(
|
|
key,
|
|
value,
|
|
key_cache,
|
|
value_cache,
|
|
attn_metadata.slot_mapping,
|
|
layer.isa, # type: ignore
|
|
k_scale=layer._k_scale_float,
|
|
v_scale=layer._v_scale_float,
|
|
kv_cache_dtype=self.kv_cache_dtype,
|
|
)
|
|
|
|
# The CPU kernel executes attention sinks natively in bf16. If the
|
|
# sinks tensor is anything other than bf16, cast it to fp32 so it is
|
|
# executed in full float precision (done lazily here, after weights
|
|
# are loaded, rather than at __init__ time).
|
|
if self.sinks is not None and self.sinks.dtype not in [
|
|
torch.bfloat16,
|
|
torch.float32,
|
|
]:
|
|
self.sinks = self.sinks.to(torch.float32)
|
|
ops.cpu_attention_with_kv_cache(
|
|
query=query[:num_actual_tokens],
|
|
key_cache=key_cache,
|
|
value_cache=value_cache,
|
|
output=output[:num_actual_tokens], # type: ignore
|
|
query_start_loc=attn_metadata.query_start_loc,
|
|
seq_lens=attn_metadata.seq_lens,
|
|
scale=self.scale,
|
|
causal=attn_metadata.causal,
|
|
alibi_slopes=self.alibi_slopes, # type: ignore
|
|
sliding_window=self.sliding_window,
|
|
block_table=attn_metadata.block_table,
|
|
softcap=self.logits_soft_cap,
|
|
scheduler_metadata=attn_metadata.scheduler_metadata,
|
|
s_aux=self.sinks,
|
|
dynamic_causal=attn_metadata.dynamic_causal,
|
|
k_scale=layer._k_scale_float,
|
|
v_scale=layer._v_scale_float,
|
|
kv_cache_dtype=self.kv_cache_dtype,
|
|
)
|
|
|
|
return output
|
|
|
|
def do_kv_cache_update(
|
|
self,
|
|
layer: torch.nn.Module,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
kv_cache: torch.Tensor,
|
|
slot_mapping: torch.Tensor,
|
|
) -> None:
|
|
if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
|
|
return
|
|
|
|
num_blocks, num_kv_heads, block_size, _ = kv_cache.size()
|
|
kv_cache = kv_cache.view((num_blocks, num_kv_heads, block_size * 2, -1))
|
|
key_cache, value_cache = kv_cache.chunk(2, dim=2)
|
|
ops.cpu_attn_reshape_and_cache(
|
|
key,
|
|
value,
|
|
key_cache,
|
|
value_cache,
|
|
slot_mapping,
|
|
layer.isa, # type: ignore
|
|
k_scale=layer._k_scale_float,
|
|
v_scale=layer._v_scale_float,
|
|
kv_cache_dtype=self.kv_cache_dtype,
|
|
)
|
|
|
|
|
|
@functools.lru_cache(maxsize=1)
|
|
def _riscv_supports_rvv() -> bool:
|
|
"""Whether the C++ RVV attention path is usable.
|
|
|
|
The kernel in csrc/cpu/cpu_attn_rvv.hpp uses VLEN-agnostic RVVI()
|
|
macros and supports VLEN=128 and VLEN=256. CMake auto-detects the
|
|
largest zvl<N>b from /proc/cpuinfo and passes it via -mrvv-vector-bits.
|
|
The RVV path is compiled whenever __riscv_v_min_vlen is defined, so
|
|
we check that at least one supported zvl<N>b is advertised.
|
|
"""
|
|
# The C++ compile-time check is the ground truth: it knows which
|
|
# VLEN the binary was actually compiled for. The cpuinfo check
|
|
# below is only a fast-path shortcut.
|
|
try:
|
|
import torch
|
|
|
|
if torch.ops._C.cpu_attn_has_isa("rvv"):
|
|
return True
|
|
except Exception:
|
|
pass
|
|
|
|
# Fallback: check /proc/cpuinfo for zvl128b/zvl256b.
|
|
try:
|
|
with open("/proc/cpuinfo") as f:
|
|
cpuinfo = f.read()
|
|
except OSError:
|
|
return False
|
|
return any(f"zvl{n}b" in cpuinfo for n in (128, 256))
|
|
|
|
|
|
def _get_attn_isa(
|
|
dtype: torch.dtype,
|
|
block_size: int,
|
|
head_size: int | None = None,
|
|
kv_cache_dtype: str | None = None,
|
|
) -> str:
|
|
fp8_kv = is_quantized_kv_cache(kv_cache_dtype) if kv_cache_dtype else False
|
|
if head_size is not None and head_size % 32 != 0 and head_size % 16 == 0:
|
|
if fp8_kv:
|
|
raise NotImplementedError(
|
|
"FP8 KV cache requires head_size divisible by 32 on CPU."
|
|
)
|
|
return "vec16"
|
|
supports_amx = torch.cpu._is_amx_tile_supported()
|
|
arch = current_platform.get_cpu_architecture()
|
|
supports_arm = arch == CpuArchEnum.ARM
|
|
supports_vxe = arch == CpuArchEnum.S390X
|
|
supports_riscv = arch == CpuArchEnum.RISCV
|
|
supports_vsx = arch == CpuArchEnum.POWERPC
|
|
supports_avx512 = torch.cpu._is_avx512_supported()
|
|
if fp8_kv and not supports_amx and not supports_avx512:
|
|
raise NotImplementedError(
|
|
"FP8 KV cache on CPU requires x86 with AVX-512 or AMX."
|
|
)
|
|
if supports_amx and dtype in (torch.bfloat16,) and block_size % 32 == 0:
|
|
return "amx"
|
|
elif supports_arm:
|
|
return "neon"
|
|
elif block_size % 32 == 0:
|
|
if supports_riscv and _riscv_supports_rvv():
|
|
return "rvv"
|
|
elif supports_vxe:
|
|
return "vxe"
|
|
elif supports_vsx:
|
|
return "vsx"
|
|
else:
|
|
return "vec"
|
|
else:
|
|
return "vec16"
|