iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Feat][CPU] Enable Gated DeltaNet Attention (Qwen 3.5 / 3.6) (#41025)

Browse Source 
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>
Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com>
This commit is contained in:
Fadi Arafeh
2026-05-07 12:57:09 +08:00
committed by GitHub
co-authored by Vadim Gimpelson
parent 20cac26b19
commit 51f22dcfd0
11 changed files with 615 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
+15
View File
@@ -67,6 +67,21 @@ function cpu_tests() {
--num-prompts 20 \
--endpoint /v1/completions
kill -s SIGTERM $server_pid &'
# smoke test for Gated DeltaNet
docker exec cpu-test bash -c '
set -e
VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS vllm serve Qwen/Qwen3.5-0.8B --max-model-len 2048 &
server_pid=$!
timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
vllm bench serve \
--backend vllm \
--dataset-name random \
--model Qwen/Qwen3.5-0.8B \
--num-prompts 20 \
--endpoint /v1/completions
kill -s SIGTERM $server_pid &'
}
# All of CPU tests are expected to be finished less than 40 mins.
docs/design/attention_backends.md
+1 -1
View File
@@ -168,7 +168,7 @@ Priority is **1 = highest** (tried first).
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
vllm/model_executor/layers/linear.py
+3
View File
@@ -213,6 +213,9 @@ class UnquantizedLinearMethod(LinearMethodBase):
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
if current_platform.is_cpu():
if layer.weight.ndim != 2:
# this is not a linear layer
return
from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm
dispatch_cpu_unquantized_gemm(layer, remove_weight=True)
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+60 -3
View File
@@ -271,9 +271,16 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
else 0
)
self.gqa_interleaved_layout = gqa_interleaved_layout
self._forward_method = (
self.forward_xpu if current_platform.is_xpu() else self.forward_cuda
)
self._forward_method = self.forward_cuda
if current_platform.is_xpu():
self._forward_method = self.forward_xpu
elif current_platform.is_cpu():
from vllm.model_executor.layers.mamba.ops.cpu.gdn_attention import (
register_cpu_gdn_attention_ops,
)
register_cpu_gdn_attention_ops()
self._forward_method = self.forward_cpu
# QKV
self.conv_dim = self.key_dim * 2 + self.value_dim
@@ -645,6 +652,56 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
output[:num_tokens], _ = self.out_proj(core_attn_out)
def forward_cpu(
self,
hidden_states: torch.Tensor,
output: torch.Tensor,
):
assert not hasattr(self, "in_proj_qkv"), "lora isn't supported on CPU."
mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
ba, _ = self.in_proj_ba(hidden_states)
if self.gqa_interleaved_layout:
# Qwen3-Next: unpack the interleaved GQA layout
query, key, value, z, b, a = self.fix_query_key_value_ordering(
mixed_qkvz, ba
)
query, key, value = map(
lambda x: rearrange(x, "l p d -> l (p d)"), (query, key, value)
)
mixed_qkv = torch.cat((query, key, value), dim=-1)
else:
# Qwen3.5: weights are already in [q, k, v, z] and [b, a] order
qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size
z_size = self.value_dim // self.tp_size
mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1)
z = z.reshape(z.size(0), -1, self.head_v_dim)
b, a = ba.chunk(2, dim=-1)
num_tokens = hidden_states.size(0)
core_attn_out = torch.zeros(
(num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
dtype=hidden_states.dtype,
device=hidden_states.device,
)
torch.ops.vllm.cpu_gdn_attention_core(
mixed_qkv,
b,
a,
core_attn_out,
_encode_layer_name(self.prefix),
)
z_shape_og = z.shape
core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
z = z.reshape(-1, z.shape[-1])
core_attn_out = self.norm(core_attn_out, z)
core_attn_out = core_attn_out.reshape(z_shape_og)
core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
output[:num_tokens], _ = self.out_proj(core_attn_out)
def _warmup_prefill_kernels(self, mixed_qkv: torch.Tensor) -> None:
"""Warm up GDN prefill kernels during V1 profiling.
vllm/model_executor/layers/mamba/ops/cpu/__init__.py
View File
vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py
+88
View File
@@ -0,0 +1,88 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import torch
import torch.nn.functional as F
# for prefill
def causal_conv1d_torch(
x: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor | None,
conv_states: torch.Tensor,
query_start_loc: torch.Tensor,
cache_indices: torch.Tensor,
has_initial_state: torch.Tensor,
activation: str | None = "silu",
) -> torch.Tensor:
out = torch.empty_like(x)
state_len = weight.shape[1] - 1
assert activation in {None, "silu", "swish"}
seq_begin_end_idx = [
(int(query_start_loc[idx].item()), int(query_start_loc[idx + 1].item()))
for idx in range(query_start_loc.shape[0] - 1)
]
weight = weight.unsqueeze(1)
for seq_idx, (bos, eos) in enumerate(seq_begin_end_idx):
slot = int(cache_indices[seq_idx].item())
seq_x = x[:, bos:eos].unsqueeze(0)
if bool(has_initial_state[seq_idx].item()):
initial_state = conv_states[slot, :, :state_len].unsqueeze(0)
else:
initial_state = torch.zeros(
1,
weight.shape[0],
state_len,
device=seq_x.device,
dtype=seq_x.dtype,
)
conv_input = torch.cat([initial_state, seq_x], dim=-1).to(weight.dtype)
seq_out = F.conv1d(
conv_input,
weight,
bias,
padding=0,
groups=weight.shape[0],
)
seq_out = seq_out[..., -seq_x.shape[-1] :].to(dtype=x.dtype)
if activation in ("silu", "swish"):
seq_out = F.silu(seq_out)
out[:, bos:eos] = seq_out.squeeze(0)
conv_states[slot, :, :state_len].copy_(conv_input[..., -state_len:].squeeze(0))
return out
# for decode
def causal_conv1d_update_torch(
x: torch.Tensor,
conv_state: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor | None = None,
activation: str | None = None,
) -> torch.Tensor:
assert activation in {None, "silu", "swish"}
_, dim, seq_len = x.shape
state_len = conv_state.shape[-1]
x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype)
conv_state.copy_(x_new[:, :, -state_len:])
out = F.conv1d(
x_new,
weight.unsqueeze(1),
bias,
padding=0,
groups=dim,
)[:, :, -seq_len:]
if activation in ("silu", "swish"):
out = F.silu(out)
return out
vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py
+208
View File
@@ -0,0 +1,208 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import torch
from vllm.forward_context import ForwardContext, get_forward_context
from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
causal_conv1d_torch,
causal_conv1d_update_torch,
)
from vllm.model_executor.layers.mamba.ops.cpu.recurrent_gated_delta_rule import (
chunk_gated_delta_rule,
gdn_gating,
recurrent_gated_delta_rule,
)
from vllm.utils.torch_utils import (
LayerNameType,
_resolve_layer_name,
direct_register_custom_op,
)
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
_CPU_GDN_ATTENTION_OPS_REGISTERED = False
def cpu_gdn_attention_core(
mixed_qkv: torch.Tensor,
b: torch.Tensor,
a: torch.Tensor,
core_attn_out: torch.Tensor,
layer_name: LayerNameType,
) -> None:
"""CPU custom op for the core GDN attention computation."""
layer_name = _resolve_layer_name(layer_name)
forward_context: ForwardContext = get_forward_context()
layer = forward_context.no_compile_layers[layer_name]
attn_metadata = forward_context.attn_metadata
if attn_metadata is None:
return
assert isinstance(attn_metadata, dict)
attn_metadata_i = attn_metadata[layer.prefix]
assert isinstance(attn_metadata_i, GDNAttentionMetadata)
if attn_metadata_i.num_actual_tokens == 0:
return
assert (
attn_metadata_i.spec_sequence_masks is None
and attn_metadata_i.num_accepted_tokens is None
), "speculative decode not supported in CPU GDN attention."
state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor
query_start_loc = attn_metadata_i.non_spec_query_start_loc
assert state_indices_tensor is not None
assert query_start_loc is not None
# [num_allocated_slots, conv_dim, kernel - 1]
conv_state = layer.kv_cache[0]
if not is_conv_state_dim_first():
conv_state = conv_state.transpose(-1, -2)
# [num_allocated_slots, num_v_heads / tp_size, v_dim, k_dim]
ssm_state = layer.kv_cache[1]
num_decodes = attn_metadata_i.num_decodes
num_decode_tokens = attn_metadata_i.num_decode_tokens
num_prefills = attn_metadata_i.num_prefills
num_prefill_tokens = attn_metadata_i.num_prefill_tokens
conv_weights = layer.conv1d.weight.view(
layer.conv1d.weight.size(0), layer.conv1d.weight.size(2)
)
# all decode requests (batched)
if num_decodes > 0:
decode_mixed_qkv = mixed_qkv[:num_decode_tokens]
decode_b = b[:num_decode_tokens]
decode_a = a[:num_decode_tokens]
decode_state_indices = state_indices_tensor[:num_decodes]
decode_conv_state = conv_state[decode_state_indices].contiguous()
decode_mixed_qkv = causal_conv1d_update_torch(
# [B, dim] -> [B, dim, 1]
x=decode_mixed_qkv.unsqueeze(-1),
conv_state=decode_conv_state,
weight=conv_weights,
bias=layer.conv1d.bias,
activation=layer.activation,
).squeeze(-1)
conv_state[decode_state_indices] = decode_conv_state
query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv)
# [1, L, H, D] -> [B, 1, H, D] for batched decode
query = query.transpose(0, 1).contiguous()
key = key.transpose(0, 1).contiguous()
value = value.transpose(0, 1).contiguous()
g, beta_output = gdn_gating(
A_log=layer.A_log,
a=decode_a,
b=decode_b,
dt_bias=layer.dt_bias,
)
if g.ndim == 2:
g = g.unsqueeze(1)
beta_output = beta_output.unsqueeze(1)
initial_state = ssm_state[decode_state_indices].contiguous()
attn_out, last_recurrent_state = recurrent_gated_delta_rule(
query=query,
key=key,
value=value,
g=g,
beta=beta_output,
initial_state=initial_state,
scale=None,
use_qk_l2norm_in_kernel=True,
)
ssm_state[decode_state_indices] = last_recurrent_state.to(
ssm_state.dtype
).contiguous()
core_attn_out[:num_decode_tokens] = attn_out.squeeze(1)
# all prefill requests: (varlen) currently naively loops over sequences
if num_prefills > 0:
has_initial_state = attn_metadata_i.has_initial_state
assert has_initial_state is not None
prefill_token_start = num_decode_tokens
prefill_token_end = prefill_token_start + num_prefill_tokens
prefill_mixed_qkv = mixed_qkv[prefill_token_start:prefill_token_end]
prefill_b = b[prefill_token_start:prefill_token_end]
prefill_a = a[prefill_token_start:prefill_token_end]
prefill_state_indices = state_indices_tensor[
num_decodes : num_decodes + num_prefills
]
prefill_query_start_loc = (
query_start_loc[num_decodes : num_decodes + num_prefills + 1]
- num_decode_tokens
)
prefill_has_initial_state = has_initial_state[
num_decodes : num_decodes + num_prefills
]
prefill_mixed_qkv = causal_conv1d_torch(
x=prefill_mixed_qkv.transpose(0, 1),
weight=conv_weights,
bias=layer.conv1d.bias,
conv_states=conv_state,
query_start_loc=prefill_query_start_loc,
cache_indices=prefill_state_indices,
has_initial_state=prefill_has_initial_state,
activation=layer.activation,
).transpose(0, 1)
query, key, value = layer.rearrange_mixed_qkv(prefill_mixed_qkv)
g, beta = gdn_gating(layer.A_log, prefill_a, prefill_b, layer.dt_bias)
if g.ndim == 2:
g = g.unsqueeze(0)
beta = beta.unsqueeze(0)
initial_state = ssm_state[prefill_state_indices].contiguous()
initial_state[~prefill_has_initial_state, ...] = 0
attn_out, last_recurrent_state = chunk_gated_delta_rule(
q=query,
k=key,
v=value,
g=g,
beta=beta,
scale=None,
initial_state=initial_state,
cu_seqlens=prefill_query_start_loc,
use_qk_l2norm_in_kernel=True,
)
ssm_state[prefill_state_indices] = last_recurrent_state.to(ssm_state.dtype)
core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0)
def cpu_gdn_attention_core_fake(
mixed_qkv: torch.Tensor,
b: torch.Tensor,
a: torch.Tensor,
core_attn_out: torch.Tensor,
layer_name: LayerNameType,
) -> None:
"""Fake implementation for torch.compile."""
return
def register_cpu_gdn_attention_ops() -> None:
global _CPU_GDN_ATTENTION_OPS_REGISTERED
if _CPU_GDN_ATTENTION_OPS_REGISTERED:
return
direct_register_custom_op(
op_name="cpu_gdn_attention_core",
op_func=cpu_gdn_attention_core,
mutates_args=["core_attn_out"],
fake_impl=cpu_gdn_attention_core_fake,
)
_CPU_GDN_ATTENTION_OPS_REGISTERED = True
vllm/model_executor/layers/mamba/ops/cpu/recurrent_gated_delta_rule.py
+223
View File
@@ -0,0 +1,223 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
import torch.nn.functional as F
def l2norm(
x: torch.Tensor,
dim: int = -1,
eps: float = 1e-6,
) -> torch.Tensor:
inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
return x * inv_norm
def recurrent_gated_delta_rule(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
initial_state: torch.Tensor,
scale: float | None = None,
use_qk_l2norm_in_kernel: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
initial_dtype = query.dtype
if use_qk_l2norm_in_kernel:
query = l2norm(query, dim=-1, eps=1e-6)
key = l2norm(key, dim=-1, eps=1e-6)
if query.shape[2] != value.shape[2]:
repeat_factor = value.shape[2] // query.shape[2]
query = query.repeat_interleave(repeat_factor, dim=2)
key = key.repeat_interleave(repeat_factor, dim=2)
query, key, value, beta, g = [
x.transpose(1, 2).contiguous().to(torch.float32)
for x in (query, key, value, beta, g)
]
batch_size, num_heads, sequence_length, _ = key.shape
v_head_dim = value.shape[-1]
if scale is None:
scale = 1 / (query.shape[-1] ** 0.5)
query = query * scale
core_attn_out = torch.empty(
batch_size,
num_heads,
sequence_length,
v_head_dim,
dtype=value.dtype,
)
last_recurrent_state = initial_state.to(value)
for token_idx in range(sequence_length):
q_t = query[:, :, token_idx]
k_t = key[:, :, token_idx]
v_t = value[:, :, token_idx]
g_t = g[:, :, token_idx].exp().unsqueeze(-1).unsqueeze(-1)
beta_t = beta[:, :, token_idx].unsqueeze(-1)
last_recurrent_state = last_recurrent_state * g_t
kv_mem = (last_recurrent_state * k_t.unsqueeze(-2)).sum(dim=-1)
delta = (v_t - kv_mem) * beta_t
last_recurrent_state = last_recurrent_state + delta.unsqueeze(
-1
) * k_t.unsqueeze(-2)
core_attn_out[:, :, token_idx] = (last_recurrent_state * q_t.unsqueeze(-2)).sum(
dim=-1
)
core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
return core_attn_out, last_recurrent_state
def gdn_gating(
A_log: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
dt_bias: torch.Tensor,
beta: float = 1.0,
threshold: float = 20.0,
) -> tuple[torch.Tensor, torch.Tensor]:
softplus_x = F.softplus(a.float() + dt_bias.float(), beta=beta, threshold=threshold)
g = -torch.exp(A_log.float()) * softplus_x
beta_output = torch.sigmoid(b.float()).to(dtype=b.dtype)
return g, beta_output
def chunk_gated_delta_rule(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
*,
initial_state: torch.Tensor,
scale: float | None = None,
cu_seqlens: torch.Tensor,
use_qk_l2norm_in_kernel: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
output = torch.empty_like(v)
state_dtype = initial_state.dtype
chunk_size = 128
sequence_bounds = [
(
seq_idx,
int(cu_seqlens[seq_idx].item()),
int(cu_seqlens[seq_idx + 1].item()),
)
for seq_idx in range(len(cu_seqlens) - 1)
]
chunk_eye = torch.eye(chunk_size, dtype=torch.float32)
num_sequences = len(sequence_bounds)
num_value_heads = v.shape[2]
value_head_dim = v.shape[3]
key_head_dim = k.shape[3]
final_state = torch.empty(
(num_sequences, num_value_heads, value_head_dim, key_head_dim),
dtype=state_dtype,
)
for seq_idx, begin, end in sequence_bounds:
q_seq = q[:, begin:end]
k_seq = k[:, begin:end]
v_seq = v[:, begin:end]
g_seq = g[:, begin:end]
beta_seq = beta[:, begin:end]
initial_dtype = q_seq.dtype
if use_qk_l2norm_in_kernel:
q_seq = l2norm(q_seq, dim=-1, eps=1e-6)
k_seq = l2norm(k_seq, dim=-1, eps=1e-6)
num_qk_heads = q_seq.shape[2]
num_value_heads = v_seq.shape[2]
if num_qk_heads != num_value_heads:
repeat_factor = num_value_heads // num_qk_heads
q_seq = q_seq.repeat_interleave(repeat_factor, dim=2)
k_seq = k_seq.repeat_interleave(repeat_factor, dim=2)
q_seq, k_seq, v_seq, beta_seq, g_seq = [
x.transpose(1, 2).contiguous().to(torch.float32)
for x in (q_seq, k_seq, v_seq, beta_seq, g_seq)
]
seq_batch_size, num_heads, seq_len, qk_head_dim = q_seq.shape
value_head_dim = v_seq.shape[-1]
if scale is None:
scale = 1 / (qk_head_dim**0.5)
q_seq = q_seq * scale
seq_state = initial_state[seq_idx : seq_idx + 1].to(v_seq)
seq_output = torch.empty(
seq_batch_size,
num_heads,
seq_len,
value_head_dim,
dtype=v_seq.dtype,
)
for chunk_start in range(0, seq_len, chunk_size):
chunk_end = min(chunk_start + chunk_size, seq_len)
q_chunk = q_seq[:, :, chunk_start:chunk_end]
k_chunk = k_seq[:, :, chunk_start:chunk_end]
v_chunk = v_seq[:, :, chunk_start:chunk_end]
beta_chunk = beta_seq[:, :, chunk_start:chunk_end]
g_chunk = g_seq[:, :, chunk_start:chunk_end]
chunk_len = chunk_end - chunk_start
cum_g = g_chunk.cumsum(dim=-1)
exp_cum_g = cum_g.exp()
decay = (cum_g.unsqueeze(-1) - cum_g.unsqueeze(-2)).exp()
interaction = (k_chunk * beta_chunk.unsqueeze(-1)) @ k_chunk.transpose(
-1, -2
)
interaction = torch.tril(interaction * decay, diagonal=-1)
system = interaction + chunk_eye[:chunk_len, :chunk_len]
solved_values = torch.linalg.solve_triangular(
system,
v_chunk * beta_chunk.unsqueeze(-1),
upper=False,
)
solved_keys = torch.linalg.solve_triangular(
system,
(k_chunk * beta_chunk.unsqueeze(-1)) * exp_cum_g.unsqueeze(-1),
upper=False,
)
incoming_memory = torch.einsum("bhvk,bhck->bhcv", seq_state, solved_keys)
transformed_values = solved_values - incoming_memory
# Each chunk contributes both from the incoming recurrent state and
# from its own in-chunk interactions.
inter_chunk = torch.einsum(
"bhvk,bhck->bhcv",
seq_state,
q_chunk * exp_cum_g.unsqueeze(-1),
)
intra_chunk = torch.tril((q_chunk @ k_chunk.transpose(-1, -2)) * decay)
seq_output[:, :, chunk_start:chunk_end] = (
inter_chunk + intra_chunk @ transformed_values
)
# Carry the recurrent state forward to the next chunk boundary.
end_decay = (cum_g[:, :, -1:] - cum_g).exp().unsqueeze(-1)
decayed_keys = k_chunk * end_decay
seq_state = seq_state * exp_cum_g[:, :, -1, None, None] + torch.einsum(
"bhcv,bhck->bhvk", transformed_values, decayed_keys
)
output[0, begin:end].copy_(
seq_output.transpose(1, 2).contiguous().to(initial_dtype).squeeze(0)
)
final_state[seq_idx].copy_(seq_state.squeeze(0).to(state_dtype).contiguous())
return output, final_state
vllm/platforms/cpu.py
+10 -3
View File
@@ -292,9 +292,16 @@ class CpuPlatform(Platform):
@classmethod
def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
# TODO: CPU still sets block_size in check_and_update_config.
# Move that logic here so block_size is chosen by the backend.
pass
model_config = vllm_config.model_config
if model_config is None or not model_config.is_hybrid:
return
# reconcile attention and mamba page sizes
backend_cls = cls._find_non_ssm_backend(vllm_config)
if backend_cls is None:
return
cls._align_hybrid_block_size(vllm_config, backend_cls)
@classmethod
def discover_numa_topology(cls) -> list[list[int]]:
vllm/triton_utils/importing.py
+2
View File
@@ -6,6 +6,7 @@ import types
from importlib.util import find_spec
from vllm.logger import init_logger
from vllm.utils.math_utils import cdiv
logger = init_logger(__name__)
@@ -79,6 +80,7 @@ class TritonPlaceholder(types.ModuleType):
self.autotune = self._dummy_decorator("autotune")
self.heuristics = self._dummy_decorator("heuristics")
self.Config = self._dummy_decorator("Config")
self.cdiv = cdiv
self.language = TritonLanguagePlaceholder()
def _dummy_decorator(self, name):
vllm/v1/attention/backends/cpu_attn.py
+5
View File
@@ -21,6 +21,7 @@ from vllm.v1.attention.backend import (
AttentionMetadataBuilder,
AttentionType,
CommonAttentionMetadata,
MultipleOf,
)
from vllm.v1.attention.backends.utils import (
split_decodes_and_prefills,
@@ -50,6 +51,10 @@ class CPUAttentionBackend(AttentionBackend):
"fp8_e5m2",
]
@staticmethod
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
return [MultipleOf(16)]
@classmethod
def get_supported_head_sizes(cls) -> list[int]:
return [32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512]