diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index 7166435ac1..0322397394 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -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. diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 83c5fc1b43..e711694b3b 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -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 | diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 06da426134..765e79331d 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -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) diff --git a/vllm/model_executor/layers/mamba/gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn_linear_attn.py index a621ab962f..b4699d4f00 100644 --- a/vllm/model_executor/layers/mamba/gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn_linear_attn.py @@ -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. diff --git a/vllm/model_executor/layers/mamba/ops/cpu/__init__.py b/vllm/model_executor/layers/mamba/ops/cpu/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py new file mode 100644 index 0000000000..b047ca6d61 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py @@ -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 diff --git a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py new file mode 100644 index 0000000000..8b4122cc48 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py @@ -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 diff --git a/vllm/model_executor/layers/mamba/ops/cpu/recurrent_gated_delta_rule.py b/vllm/model_executor/layers/mamba/ops/cpu/recurrent_gated_delta_rule.py new file mode 100644 index 0000000000..30fca3423a --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/cpu/recurrent_gated_delta_rule.py @@ -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 diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index b99510a664..750b7f2f4b 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -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]]: diff --git a/vllm/triton_utils/importing.py b/vllm/triton_utils/importing.py index f05bc555bf..9b2634c909 100644 --- a/vllm/triton_utils/importing.py +++ b/vllm/triton_utils/importing.py @@ -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): diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index 025be2c51a..51e3e865d5 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -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]