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 4, 05:00 UTC
[Feat][CPU] Enable Gated DeltaNet Attention (Qwen 3.5 / 3.6) (#41025)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com> Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com>
Fadi Arafeh committed
May 7, 2026 at 05:57 UTC
51f22dcfd068fe8f1e3192da2a1e825b930223cf
11 files changed
+615
-7
.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
+15
@@ -67,6 +67,21 @@ function cpu_tests() {
67
--num-prompts 20 \
68
--endpoint /v1/completions
69
kill -s SIGTERM $server_pid &'
70
+
71
+ # smoke test for Gated DeltaNet
72
+ docker exec cpu-test bash -c '
73
+ set -e
74
+ VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS vllm serve Qwen/Qwen3.5-0.8B --max-model-len 2048 &
75
+ server_pid=$!
76
+ timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
77
+ vllm bench serve \
78
+ --backend vllm \
79
+ --dataset-name random \
80
+ --model Qwen/Qwen3.5-0.8B \
81
+ --num-prompts 20 \
82
+ --endpoint /v1/completions
83
+ kill -s SIGTERM $server_pid &'
84
+
85
}
86
87
# All of CPU tests are expected to be finished less than 40 mins.
docs/design/attention_backends.md
+1
-1
@@ -168,7 +168,7 @@ Priority is **1 = highest** (tried first).
168
169
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
170
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
171
-| `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 |
171
+| `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 |
172
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
173
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
174
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
vllm/model_executor/layers/linear.py
+3
@@ -213,6 +213,9 @@ class UnquantizedLinearMethod(LinearMethodBase):
213
214
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
215
if current_platform.is_cpu():
216
+ if layer.weight.ndim != 2:
217
+ # this is not a linear layer
218
+ return
219
from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm
220
221
dispatch_cpu_unquantized_gemm(layer, remove_weight=True)
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+60
-3
@@ -271,9 +271,16 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
271
else 0
272
)
273
self.gqa_interleaved_layout = gqa_interleaved_layout
274
- self._forward_method = (
275
- self.forward_xpu if current_platform.is_xpu() else self.forward_cuda
276
- )
274
+ self._forward_method = self.forward_cuda
275
+ if current_platform.is_xpu():
276
+ self._forward_method = self.forward_xpu
277
+ elif current_platform.is_cpu():
278
+ from vllm.model_executor.layers.mamba.ops.cpu.gdn_attention import (
279
+ register_cpu_gdn_attention_ops,
280
+ )
281
+
282
+ register_cpu_gdn_attention_ops()
283
+ self._forward_method = self.forward_cpu
284
285
# QKV
286
self.conv_dim = self.key_dim * 2 + self.value_dim
@@ -645,6 +652,56 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
652
core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
653
output[:num_tokens], _ = self.out_proj(core_attn_out)
654
655
+ def forward_cpu(
656
+ self,
657
+ hidden_states: torch.Tensor,
658
+ output: torch.Tensor,
659
+ ):
660
+ assert not hasattr(self, "in_proj_qkv"), "lora isn't supported on CPU."
661
+
662
+ mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
663
+ ba, _ = self.in_proj_ba(hidden_states)
664
+
665
+ if self.gqa_interleaved_layout:
666
+ # Qwen3-Next: unpack the interleaved GQA layout
667
+ query, key, value, z, b, a = self.fix_query_key_value_ordering(
668
+ mixed_qkvz, ba
669
+ )
670
+ query, key, value = map(
671
+ lambda x: rearrange(x, "l p d -> l (p d)"), (query, key, value)
672
+ )
673
+ mixed_qkv = torch.cat((query, key, value), dim=-1)
674
+ else:
675
+ # Qwen3.5: weights are already in [q, k, v, z] and [b, a] order
676
+ qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size
677
+ z_size = self.value_dim // self.tp_size
678
+ mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1)
679
+ z = z.reshape(z.size(0), -1, self.head_v_dim)
680
+ b, a = ba.chunk(2, dim=-1)
681
+
682
+ num_tokens = hidden_states.size(0)
683
+ core_attn_out = torch.zeros(
684
+ (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
685
+ dtype=hidden_states.dtype,
686
+ device=hidden_states.device,
687
+ )
688
+
689
+ torch.ops.vllm.cpu_gdn_attention_core(
690
+ mixed_qkv,
691
+ b,
692
+ a,
693
+ core_attn_out,
694
+ _encode_layer_name(self.prefix),
695
+ )
696
+
697
+ z_shape_og = z.shape
698
+ core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
699
+ z = z.reshape(-1, z.shape[-1])
700
+ core_attn_out = self.norm(core_attn_out, z)
701
+ core_attn_out = core_attn_out.reshape(z_shape_og)
702
+ core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
703
+ output[:num_tokens], _ = self.out_proj(core_attn_out)
704
+
705
def _warmup_prefill_kernels(self, mixed_qkv: torch.Tensor) -> None:
706
"""Warm up GDN prefill kernels during V1 profiling.
707
vllm/model_executor/layers/mamba/ops/cpu/__init__.py
vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py
new
+88
@@ -0,0 +1,88 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from __future__ import annotations
5
+
6
+import torch
7
+import torch.nn.functional as F
8
+
9
+
10
+# for prefill
11
+def causal_conv1d_torch(
12
+ x: torch.Tensor,
13
+ weight: torch.Tensor,
14
+ bias: torch.Tensor | None,
15
+ conv_states: torch.Tensor,
16
+ query_start_loc: torch.Tensor,
17
+ cache_indices: torch.Tensor,
18
+ has_initial_state: torch.Tensor,
19
+ activation: str | None = "silu",
20
+) -> torch.Tensor:
21
+ out = torch.empty_like(x)
22
+ state_len = weight.shape[1] - 1
23
+ assert activation in {None, "silu", "swish"}
24
+
25
+ seq_begin_end_idx = [
26
+ (int(query_start_loc[idx].item()), int(query_start_loc[idx + 1].item()))
27
+ for idx in range(query_start_loc.shape[0] - 1)
28
+ ]
29
+ weight = weight.unsqueeze(1)
30
+ for seq_idx, (bos, eos) in enumerate(seq_begin_end_idx):
31
+ slot = int(cache_indices[seq_idx].item())
32
+
33
+ seq_x = x[:, bos:eos].unsqueeze(0)
34
+ if bool(has_initial_state[seq_idx].item()):
35
+ initial_state = conv_states[slot, :, :state_len].unsqueeze(0)
36
+ else:
37
+ initial_state = torch.zeros(
38
+ 1,
39
+ weight.shape[0],
40
+ state_len,
41
+ device=seq_x.device,
42
+ dtype=seq_x.dtype,
43
+ )
44
+
45
+ conv_input = torch.cat([initial_state, seq_x], dim=-1).to(weight.dtype)
46
+ seq_out = F.conv1d(
47
+ conv_input,
48
+ weight,
49
+ bias,
50
+ padding=0,
51
+ groups=weight.shape[0],
52
+ )
53
+ seq_out = seq_out[..., -seq_x.shape[-1] :].to(dtype=x.dtype)
54
+ if activation in ("silu", "swish"):
55
+ seq_out = F.silu(seq_out)
56
+
57
+ out[:, bos:eos] = seq_out.squeeze(0)
58
+ conv_states[slot, :, :state_len].copy_(conv_input[..., -state_len:].squeeze(0))
59
+
60
+ return out
61
+
62
+
63
+# for decode
64
+def causal_conv1d_update_torch(
65
+ x: torch.Tensor,
66
+ conv_state: torch.Tensor,
67
+ weight: torch.Tensor,
68
+ bias: torch.Tensor | None = None,
69
+ activation: str | None = None,
70
+) -> torch.Tensor:
71
+ assert activation in {None, "silu", "swish"}
72
+
73
+ _, dim, seq_len = x.shape
74
+ state_len = conv_state.shape[-1]
75
+
76
+ x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype)
77
+ conv_state.copy_(x_new[:, :, -state_len:])
78
+
79
+ out = F.conv1d(
80
+ x_new,
81
+ weight.unsqueeze(1),
82
+ bias,
83
+ padding=0,
84
+ groups=dim,
85
+ )[:, :, -seq_len:]
86
+ if activation in ("silu", "swish"):
87
+ out = F.silu(out)
88
+ return out
vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py
new
+208
@@ -0,0 +1,208 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from __future__ import annotations
5
+
6
+import torch
7
+
8
+from vllm.forward_context import ForwardContext, get_forward_context
9
+from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
10
+from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
11
+ causal_conv1d_torch,
12
+ causal_conv1d_update_torch,
13
+)
14
+from vllm.model_executor.layers.mamba.ops.cpu.recurrent_gated_delta_rule import (
15
+ chunk_gated_delta_rule,
16
+ gdn_gating,
17
+ recurrent_gated_delta_rule,
18
+)
19
+from vllm.utils.torch_utils import (
20
+ LayerNameType,
21
+ _resolve_layer_name,
22
+ direct_register_custom_op,
23
+)
24
+from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
25
+
26
+_CPU_GDN_ATTENTION_OPS_REGISTERED = False
27
+
28
+
29
+def cpu_gdn_attention_core(
30
+ mixed_qkv: torch.Tensor,
31
+ b: torch.Tensor,
32
+ a: torch.Tensor,
33
+ core_attn_out: torch.Tensor,
34
+ layer_name: LayerNameType,
35
+) -> None:
36
+ """CPU custom op for the core GDN attention computation."""
37
+ layer_name = _resolve_layer_name(layer_name)
38
+ forward_context: ForwardContext = get_forward_context()
39
+ layer = forward_context.no_compile_layers[layer_name]
40
+
41
+ attn_metadata = forward_context.attn_metadata
42
+
43
+ if attn_metadata is None:
44
+ return
45
+
46
+ assert isinstance(attn_metadata, dict)
47
+ attn_metadata_i = attn_metadata[layer.prefix]
48
+ assert isinstance(attn_metadata_i, GDNAttentionMetadata)
49
+
50
+ if attn_metadata_i.num_actual_tokens == 0:
51
+ return
52
+
53
+ assert (
54
+ attn_metadata_i.spec_sequence_masks is None
55
+ and attn_metadata_i.num_accepted_tokens is None
56
+ ), "speculative decode not supported in CPU GDN attention."
57
+
58
+ state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor
59
+ query_start_loc = attn_metadata_i.non_spec_query_start_loc
60
+ assert state_indices_tensor is not None
61
+ assert query_start_loc is not None
62
+
63
+ # [num_allocated_slots, conv_dim, kernel - 1]
64
+ conv_state = layer.kv_cache[0]
65
+ if not is_conv_state_dim_first():
66
+ conv_state = conv_state.transpose(-1, -2)
67
+
68
+ # [num_allocated_slots, num_v_heads / tp_size, v_dim, k_dim]
69
+ ssm_state = layer.kv_cache[1]
70
+
71
+ num_decodes = attn_metadata_i.num_decodes
72
+ num_decode_tokens = attn_metadata_i.num_decode_tokens
73
+ num_prefills = attn_metadata_i.num_prefills
74
+ num_prefill_tokens = attn_metadata_i.num_prefill_tokens
75
+
76
+ conv_weights = layer.conv1d.weight.view(
77
+ layer.conv1d.weight.size(0), layer.conv1d.weight.size(2)
78
+ )
79
+
80
+ # all decode requests (batched)
81
+ if num_decodes > 0:
82
+ decode_mixed_qkv = mixed_qkv[:num_decode_tokens]
83
+ decode_b = b[:num_decode_tokens]
84
+ decode_a = a[:num_decode_tokens]
85
+ decode_state_indices = state_indices_tensor[:num_decodes]
86
+ decode_conv_state = conv_state[decode_state_indices].contiguous()
87
+
88
+ decode_mixed_qkv = causal_conv1d_update_torch(
89
+ # [B, dim] -> [B, dim, 1]
90
+ x=decode_mixed_qkv.unsqueeze(-1),
91
+ conv_state=decode_conv_state,
92
+ weight=conv_weights,
93
+ bias=layer.conv1d.bias,
94
+ activation=layer.activation,
95
+ ).squeeze(-1)
96
+ conv_state[decode_state_indices] = decode_conv_state
97
+
98
+ query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv)
99
+
100
+ # [1, L, H, D] -> [B, 1, H, D] for batched decode
101
+ query = query.transpose(0, 1).contiguous()
102
+ key = key.transpose(0, 1).contiguous()
103
+ value = value.transpose(0, 1).contiguous()
104
+
105
+ g, beta_output = gdn_gating(
106
+ A_log=layer.A_log,
107
+ a=decode_a,
108
+ b=decode_b,
109
+ dt_bias=layer.dt_bias,
110
+ )
111
+ if g.ndim == 2:
112
+ g = g.unsqueeze(1)
113
+ beta_output = beta_output.unsqueeze(1)
114
+
115
+ initial_state = ssm_state[decode_state_indices].contiguous()
116
+ attn_out, last_recurrent_state = recurrent_gated_delta_rule(
117
+ query=query,
118
+ key=key,
119
+ value=value,
120
+ g=g,
121
+ beta=beta_output,
122
+ initial_state=initial_state,
123
+ scale=None,
124
+ use_qk_l2norm_in_kernel=True,
125
+ )
126
+ ssm_state[decode_state_indices] = last_recurrent_state.to(
127
+ ssm_state.dtype
128
+ ).contiguous()
129
+ core_attn_out[:num_decode_tokens] = attn_out.squeeze(1)
130
+
131
+ # all prefill requests: (varlen) currently naively loops over sequences
132
+ if num_prefills > 0:
133
+ has_initial_state = attn_metadata_i.has_initial_state
134
+ assert has_initial_state is not None
135
+
136
+ prefill_token_start = num_decode_tokens
137
+ prefill_token_end = prefill_token_start + num_prefill_tokens
138
+ prefill_mixed_qkv = mixed_qkv[prefill_token_start:prefill_token_end]
139
+ prefill_b = b[prefill_token_start:prefill_token_end]
140
+ prefill_a = a[prefill_token_start:prefill_token_end]
141
+ prefill_state_indices = state_indices_tensor[
142
+ num_decodes : num_decodes + num_prefills
143
+ ]
144
+ prefill_query_start_loc = (
145
+ query_start_loc[num_decodes : num_decodes + num_prefills + 1]
146
+ - num_decode_tokens
147
+ )
148
+ prefill_has_initial_state = has_initial_state[
149
+ num_decodes : num_decodes + num_prefills
150
+ ]
151
+
152
+ prefill_mixed_qkv = causal_conv1d_torch(
153
+ x=prefill_mixed_qkv.transpose(0, 1),
154
+ weight=conv_weights,
155
+ bias=layer.conv1d.bias,
156
+ conv_states=conv_state,
157
+ query_start_loc=prefill_query_start_loc,
158
+ cache_indices=prefill_state_indices,
159
+ has_initial_state=prefill_has_initial_state,
160
+ activation=layer.activation,
161
+ ).transpose(0, 1)
162
+
163
+ query, key, value = layer.rearrange_mixed_qkv(prefill_mixed_qkv)
164
+ g, beta = gdn_gating(layer.A_log, prefill_a, prefill_b, layer.dt_bias)
165
+ if g.ndim == 2:
166
+ g = g.unsqueeze(0)
167
+ beta = beta.unsqueeze(0)
168
+
169
+ initial_state = ssm_state[prefill_state_indices].contiguous()
170
+ initial_state[~prefill_has_initial_state, ...] = 0
171
+ attn_out, last_recurrent_state = chunk_gated_delta_rule(
172
+ q=query,
173
+ k=key,
174
+ v=value,
175
+ g=g,
176
+ beta=beta,
177
+ scale=None,
178
+ initial_state=initial_state,
179
+ cu_seqlens=prefill_query_start_loc,
180
+ use_qk_l2norm_in_kernel=True,
181
+ )
182
+ ssm_state[prefill_state_indices] = last_recurrent_state.to(ssm_state.dtype)
183
+ core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0)
184
+
185
+
186
+def cpu_gdn_attention_core_fake(
187
+ mixed_qkv: torch.Tensor,
188
+ b: torch.Tensor,
189
+ a: torch.Tensor,
190
+ core_attn_out: torch.Tensor,
191
+ layer_name: LayerNameType,
192
+) -> None:
193
+ """Fake implementation for torch.compile."""
194
+ return
195
+
196
+
197
+def register_cpu_gdn_attention_ops() -> None:
198
+ global _CPU_GDN_ATTENTION_OPS_REGISTERED
199
+ if _CPU_GDN_ATTENTION_OPS_REGISTERED:
200
+ return
201
+
202
+ direct_register_custom_op(
203
+ op_name="cpu_gdn_attention_core",
204
+ op_func=cpu_gdn_attention_core,
205
+ mutates_args=["core_attn_out"],
206
+ fake_impl=cpu_gdn_attention_core_fake,
207
+ )
208
+ _CPU_GDN_ATTENTION_OPS_REGISTERED = True
vllm/model_executor/layers/mamba/ops/cpu/recurrent_gated_delta_rule.py
new
+223
@@ -0,0 +1,223 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+
5
+import torch
6
+import torch.nn.functional as F
7
+
8
+
9
+def l2norm(
10
+ x: torch.Tensor,
11
+ dim: int = -1,
12
+ eps: float = 1e-6,
13
+) -> torch.Tensor:
14
+ inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
15
+ return x * inv_norm
16
+
17
+
18
+def recurrent_gated_delta_rule(
19
+ query: torch.Tensor,
20
+ key: torch.Tensor,
21
+ value: torch.Tensor,
22
+ g: torch.Tensor,
23
+ beta: torch.Tensor,
24
+ initial_state: torch.Tensor,
25
+ scale: float | None = None,
26
+ use_qk_l2norm_in_kernel: bool = False,
27
+) -> tuple[torch.Tensor, torch.Tensor]:
28
+ initial_dtype = query.dtype
29
+ if use_qk_l2norm_in_kernel:
30
+ query = l2norm(query, dim=-1, eps=1e-6)
31
+ key = l2norm(key, dim=-1, eps=1e-6)
32
+
33
+ if query.shape[2] != value.shape[2]:
34
+ repeat_factor = value.shape[2] // query.shape[2]
35
+ query = query.repeat_interleave(repeat_factor, dim=2)
36
+ key = key.repeat_interleave(repeat_factor, dim=2)
37
+
38
+ query, key, value, beta, g = [
39
+ x.transpose(1, 2).contiguous().to(torch.float32)
40
+ for x in (query, key, value, beta, g)
41
+ ]
42
+
43
+ batch_size, num_heads, sequence_length, _ = key.shape
44
+ v_head_dim = value.shape[-1]
45
+ if scale is None:
46
+ scale = 1 / (query.shape[-1] ** 0.5)
47
+ query = query * scale
48
+
49
+ core_attn_out = torch.empty(
50
+ batch_size,
51
+ num_heads,
52
+ sequence_length,
53
+ v_head_dim,
54
+ dtype=value.dtype,
55
+ )
56
+ last_recurrent_state = initial_state.to(value)
57
+
58
+ for token_idx in range(sequence_length):
59
+ q_t = query[:, :, token_idx]
60
+ k_t = key[:, :, token_idx]
61
+ v_t = value[:, :, token_idx]
62
+ g_t = g[:, :, token_idx].exp().unsqueeze(-1).unsqueeze(-1)
63
+ beta_t = beta[:, :, token_idx].unsqueeze(-1)
64
+
65
+ last_recurrent_state = last_recurrent_state * g_t
66
+ kv_mem = (last_recurrent_state * k_t.unsqueeze(-2)).sum(dim=-1)
67
+ delta = (v_t - kv_mem) * beta_t
68
+ last_recurrent_state = last_recurrent_state + delta.unsqueeze(
69
+ -1
70
+ ) * k_t.unsqueeze(-2)
71
+ core_attn_out[:, :, token_idx] = (last_recurrent_state * q_t.unsqueeze(-2)).sum(
72
+ dim=-1
73
+ )
74
+
75
+ core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
76
+ return core_attn_out, last_recurrent_state
77
+
78
+
79
+def gdn_gating(
80
+ A_log: torch.Tensor,
81
+ a: torch.Tensor,
82
+ b: torch.Tensor,
83
+ dt_bias: torch.Tensor,
84
+ beta: float = 1.0,
85
+ threshold: float = 20.0,
86
+) -> tuple[torch.Tensor, torch.Tensor]:
87
+ softplus_x = F.softplus(a.float() + dt_bias.float(), beta=beta, threshold=threshold)
88
+ g = -torch.exp(A_log.float()) * softplus_x
89
+ beta_output = torch.sigmoid(b.float()).to(dtype=b.dtype)
90
+ return g, beta_output
91
+
92
+
93
+def chunk_gated_delta_rule(
94
+ q: torch.Tensor,
95
+ k: torch.Tensor,
96
+ v: torch.Tensor,
97
+ g: torch.Tensor,
98
+ beta: torch.Tensor,
99
+ *,
100
+ initial_state: torch.Tensor,
101
+ scale: float | None = None,
102
+ cu_seqlens: torch.Tensor,
103
+ use_qk_l2norm_in_kernel: bool = False,
104
+) -> tuple[torch.Tensor, torch.Tensor]:
105
+ output = torch.empty_like(v)
106
+ state_dtype = initial_state.dtype
107
+ chunk_size = 128
108
+ sequence_bounds = [
109
+ (
110
+ seq_idx,
111
+ int(cu_seqlens[seq_idx].item()),
112
+ int(cu_seqlens[seq_idx + 1].item()),
113
+ )
114
+ for seq_idx in range(len(cu_seqlens) - 1)
115
+ ]
116
+ chunk_eye = torch.eye(chunk_size, dtype=torch.float32)
117
+ num_sequences = len(sequence_bounds)
118
+ num_value_heads = v.shape[2]
119
+ value_head_dim = v.shape[3]
120
+ key_head_dim = k.shape[3]
121
+ final_state = torch.empty(
122
+ (num_sequences, num_value_heads, value_head_dim, key_head_dim),
123
+ dtype=state_dtype,
124
+ )
125
+
126
+ for seq_idx, begin, end in sequence_bounds:
127
+ q_seq = q[:, begin:end]
128
+ k_seq = k[:, begin:end]
129
+ v_seq = v[:, begin:end]
130
+ g_seq = g[:, begin:end]
131
+ beta_seq = beta[:, begin:end]
132
+
133
+ initial_dtype = q_seq.dtype
134
+ if use_qk_l2norm_in_kernel:
135
+ q_seq = l2norm(q_seq, dim=-1, eps=1e-6)
136
+ k_seq = l2norm(k_seq, dim=-1, eps=1e-6)
137
+
138
+ num_qk_heads = q_seq.shape[2]
139
+ num_value_heads = v_seq.shape[2]
140
+ if num_qk_heads != num_value_heads:
141
+ repeat_factor = num_value_heads // num_qk_heads
142
+ q_seq = q_seq.repeat_interleave(repeat_factor, dim=2)
143
+ k_seq = k_seq.repeat_interleave(repeat_factor, dim=2)
144
+
145
+ q_seq, k_seq, v_seq, beta_seq, g_seq = [
146
+ x.transpose(1, 2).contiguous().to(torch.float32)
147
+ for x in (q_seq, k_seq, v_seq, beta_seq, g_seq)
148
+ ]
149
+ seq_batch_size, num_heads, seq_len, qk_head_dim = q_seq.shape
150
+ value_head_dim = v_seq.shape[-1]
151
+
152
+ if scale is None:
153
+ scale = 1 / (qk_head_dim**0.5)
154
+
155
+ q_seq = q_seq * scale
156
+
157
+ seq_state = initial_state[seq_idx : seq_idx + 1].to(v_seq)
158
+ seq_output = torch.empty(
159
+ seq_batch_size,
160
+ num_heads,
161
+ seq_len,
162
+ value_head_dim,
163
+ dtype=v_seq.dtype,
164
+ )
165
+
166
+ for chunk_start in range(0, seq_len, chunk_size):
167
+ chunk_end = min(chunk_start + chunk_size, seq_len)
168
+ q_chunk = q_seq[:, :, chunk_start:chunk_end]
169
+ k_chunk = k_seq[:, :, chunk_start:chunk_end]
170
+ v_chunk = v_seq[:, :, chunk_start:chunk_end]
171
+ beta_chunk = beta_seq[:, :, chunk_start:chunk_end]
172
+ g_chunk = g_seq[:, :, chunk_start:chunk_end]
173
+ chunk_len = chunk_end - chunk_start
174
+
175
+ cum_g = g_chunk.cumsum(dim=-1)
176
+ exp_cum_g = cum_g.exp()
177
+ decay = (cum_g.unsqueeze(-1) - cum_g.unsqueeze(-2)).exp()
178
+
179
+ interaction = (k_chunk * beta_chunk.unsqueeze(-1)) @ k_chunk.transpose(
180
+ -1, -2
181
+ )
182
+ interaction = torch.tril(interaction * decay, diagonal=-1)
183
+ system = interaction + chunk_eye[:chunk_len, :chunk_len]
184
+
185
+ solved_values = torch.linalg.solve_triangular(
186
+ system,
187
+ v_chunk * beta_chunk.unsqueeze(-1),
188
+ upper=False,
189
+ )
190
+ solved_keys = torch.linalg.solve_triangular(
191
+ system,
192
+ (k_chunk * beta_chunk.unsqueeze(-1)) * exp_cum_g.unsqueeze(-1),
193
+ upper=False,
194
+ )
195
+
196
+ incoming_memory = torch.einsum("bhvk,bhck->bhcv", seq_state, solved_keys)
197
+ transformed_values = solved_values - incoming_memory
198
+
199
+ # Each chunk contributes both from the incoming recurrent state and
200
+ # from its own in-chunk interactions.
201
+ inter_chunk = torch.einsum(
202
+ "bhvk,bhck->bhcv",
203
+ seq_state,
204
+ q_chunk * exp_cum_g.unsqueeze(-1),
205
+ )
206
+ intra_chunk = torch.tril((q_chunk @ k_chunk.transpose(-1, -2)) * decay)
207
+ seq_output[:, :, chunk_start:chunk_end] = (
208
+ inter_chunk + intra_chunk @ transformed_values
209
+ )
210
+
211
+ # Carry the recurrent state forward to the next chunk boundary.
212
+ end_decay = (cum_g[:, :, -1:] - cum_g).exp().unsqueeze(-1)
213
+ decayed_keys = k_chunk * end_decay
214
+ seq_state = seq_state * exp_cum_g[:, :, -1, None, None] + torch.einsum(
215
+ "bhcv,bhck->bhvk", transformed_values, decayed_keys
216
+ )
217
+
218
+ output[0, begin:end].copy_(
219
+ seq_output.transpose(1, 2).contiguous().to(initial_dtype).squeeze(0)
220
+ )
221
+ final_state[seq_idx].copy_(seq_state.squeeze(0).to(state_dtype).contiguous())
222
+
223
+ return output, final_state
vllm/platforms/cpu.py
+10
-3
@@ -292,9 +292,16 @@ class CpuPlatform(Platform):
292
293
@classmethod
294
def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
295
- # TODO: CPU still sets block_size in check_and_update_config.
296
- # Move that logic here so block_size is chosen by the backend.
297
- pass
295
+ model_config = vllm_config.model_config
296
+ if model_config is None or not model_config.is_hybrid:
297
+ return
298
+
299
+ # reconcile attention and mamba page sizes
300
+ backend_cls = cls._find_non_ssm_backend(vllm_config)
301
+ if backend_cls is None:
302
+ return
303
+
304
+ cls._align_hybrid_block_size(vllm_config, backend_cls)
305
306
@classmethod
307
def discover_numa_topology(cls) -> list[list[int]]:
vllm/triton_utils/importing.py
+2
@@ -6,6 +6,7 @@ import types
6
from importlib.util import find_spec
7
8
from vllm.logger import init_logger
9
+from vllm.utils.math_utils import cdiv
10
11
logger = init_logger(__name__)
12
@@ -79,6 +80,7 @@ class TritonPlaceholder(types.ModuleType):
80
self.autotune = self._dummy_decorator("autotune")
81
self.heuristics = self._dummy_decorator("heuristics")
82
self.Config = self._dummy_decorator("Config")
83
+ self.cdiv = cdiv
84
self.language = TritonLanguagePlaceholder()
85
86
def _dummy_decorator(self, name):
vllm/v1/attention/backends/cpu_attn.py
+5
@@ -21,6 +21,7 @@ from vllm.v1.attention.backend import (
21
AttentionMetadataBuilder,
22
AttentionType,
23
CommonAttentionMetadata,
24
+ MultipleOf,
25
)
26
from vllm.v1.attention.backends.utils import (
27
split_decodes_and_prefills,
@@ -50,6 +51,10 @@ class CPUAttentionBackend(AttentionBackend):
51
"fp8_e5m2",
52
]
53
54
+ @staticmethod
55
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
56
+ return [MultipleOf(16)]
57
+
58
@classmethod
59
def get_supported_head_sizes(cls) -> list[int]:
60
return [32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512]