| 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 | import vllm._custom_ops as ops |
| 10 | from vllm.forward_context import ForwardContext, get_forward_context |
| 11 | from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first |
| 12 | from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( |
| 13 | causal_conv1d_fn_cpu as causal_conv1d_torch, |
| 14 | ) |
| 15 | from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( |
| 16 | causal_conv1d_update_cpu, |
| 17 | ) |
| 18 | from vllm.platforms import CpuArchEnum, current_platform |
| 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 _stage_ds_conv_states( |
| 30 | conv_states: torch.Tensor, state_indices: torch.Tensor |
| 31 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 32 | storage = conv_states.transpose(1, 2).index_select(0, state_indices) |
| 33 | scratch = storage.transpose(1, 2) |
| 34 | identity_indices = torch.arange( |
| 35 | state_indices.numel(), dtype=torch.int32, device=state_indices.device |
| 36 | ) |
| 37 | return scratch, identity_indices |
| 38 | |
| 39 | |
| 40 | def _scatter_ds_conv_states( |
| 41 | conv_states: torch.Tensor, |
| 42 | state_indices: torch.Tensor, |
| 43 | scratch: torch.Tensor, |
| 44 | ) -> None: |
| 45 | conv_states.index_copy_(0, state_indices.to(torch.int64), scratch) |
| 46 | |
| 47 | |
| 48 | def _causal_conv1d_fwd_cpu( |
| 49 | *, |
| 50 | x: torch.Tensor, |
| 51 | weight: torch.Tensor, |
| 52 | bias: torch.Tensor | None, |
| 53 | conv_states: torch.Tensor, |
| 54 | query_start_loc: torch.Tensor, |
| 55 | cache_indices: torch.Tensor, |
| 56 | has_initial_state: torch.Tensor, |
| 57 | silu_activation: bool, |
| 58 | is_weight_packed: bool, |
| 59 | ) -> torch.Tensor: |
| 60 | if not is_conv_state_dim_first(): |
| 61 | return ops.causal_conv1d_fwd_cpu( |
| 62 | x=x, |
| 63 | weight=weight, |
| 64 | bias=bias, |
| 65 | conv_states=conv_states, |
| 66 | query_start_loc=query_start_loc, |
| 67 | cache_indices=cache_indices, |
| 68 | has_initial_state=has_initial_state, |
| 69 | silu_activation=silu_activation, |
| 70 | is_weight_packed=is_weight_packed, |
| 71 | ) |
| 72 | |
| 73 | scratch, identity_indices = _stage_ds_conv_states(conv_states, cache_indices) |
| 74 | out = ops.causal_conv1d_fwd_cpu( |
| 75 | x=x, |
| 76 | weight=weight, |
| 77 | bias=bias, |
| 78 | conv_states=scratch, |
| 79 | query_start_loc=query_start_loc, |
| 80 | cache_indices=identity_indices, |
| 81 | has_initial_state=has_initial_state, |
| 82 | silu_activation=silu_activation, |
| 83 | is_weight_packed=is_weight_packed, |
| 84 | ) |
| 85 | _scatter_ds_conv_states(conv_states, cache_indices, scratch) |
| 86 | return out |
| 87 | |
| 88 | |
| 89 | def _causal_conv1d_update_cpu( |
| 90 | *, |
| 91 | x: torch.Tensor, |
| 92 | conv_states: torch.Tensor, |
| 93 | weight: torch.Tensor, |
| 94 | bias: torch.Tensor | None, |
| 95 | silu_activation: bool, |
| 96 | conv_state_indices: torch.Tensor, |
| 97 | is_weight_packed: bool, |
| 98 | num_accepted_tokens: torch.Tensor | None = None, |
| 99 | ) -> torch.Tensor: |
| 100 | if not is_conv_state_dim_first(): |
| 101 | return ops.causal_conv1d_update_cpu( |
| 102 | x=x, |
| 103 | conv_states=conv_states, |
| 104 | weight=weight, |
| 105 | bias=bias, |
| 106 | silu_activation=silu_activation, |
| 107 | conv_state_indices=conv_state_indices, |
| 108 | is_weight_packed=is_weight_packed, |
| 109 | num_accepted_tokens=num_accepted_tokens, |
| 110 | ) |
| 111 | |
| 112 | scratch, identity_indices = _stage_ds_conv_states(conv_states, conv_state_indices) |
| 113 | out = ops.causal_conv1d_update_cpu( |
| 114 | x=x, |
| 115 | conv_states=scratch, |
| 116 | weight=weight, |
| 117 | bias=bias, |
| 118 | silu_activation=silu_activation, |
| 119 | conv_state_indices=identity_indices, |
| 120 | is_weight_packed=is_weight_packed, |
| 121 | num_accepted_tokens=num_accepted_tokens, |
| 122 | ) |
| 123 | _scatter_ds_conv_states(conv_states, conv_state_indices, scratch) |
| 124 | return out |
| 125 | |
| 126 | |
| 127 | def is_arm_bf16() -> bool: |
| 128 | return ( |
| 129 | current_platform.get_cpu_architecture() == CpuArchEnum.ARM |
| 130 | and torch.cpu.get_capabilities().get("bf16", False) |
| 131 | ) |
| 132 | |
| 133 | |
| 134 | def cpu_gdn_attention_core( |
| 135 | mixed_qkv: torch.Tensor, |
| 136 | b: torch.Tensor, |
| 137 | a: torch.Tensor, |
| 138 | core_attn_out: torch.Tensor, |
| 139 | layer_name: LayerNameType, |
| 140 | ) -> None: |
| 141 | """CPU custom op for the core GDN attention computation.""" |
| 142 | layer_name = _resolve_layer_name(layer_name) |
| 143 | forward_context: ForwardContext = get_forward_context() |
| 144 | layer = forward_context.no_compile_layers[layer_name] |
| 145 | |
| 146 | attn_metadata = forward_context.attn_metadata |
| 147 | |
| 148 | if attn_metadata is None: |
| 149 | return |
| 150 | |
| 151 | assert isinstance(attn_metadata, dict) |
| 152 | attn_metadata_i = attn_metadata[layer.prefix] |
| 153 | assert isinstance(attn_metadata_i, GDNAttentionMetadata) |
| 154 | |
| 155 | if attn_metadata_i.num_actual_tokens == 0: |
| 156 | return |
| 157 | |
| 158 | assert mixed_qkv.dtype == torch.bfloat16, "CPU GDN attention requires BF16." |
| 159 | |
| 160 | # The conv-state cache temporal dim is ``width - 1`` when speculative |
| 161 | # decoding is disabled, and ``width - 1 + num_spec`` when it is enabled |
| 162 | # (vLLM allocates a wider rolling buffer to support draft-token rollback). |
| 163 | width = layer.conv1d.weight.size(-1) |
| 164 | conv_cache = layer.kv_cache[0] |
| 165 | if is_conv_state_dim_first(): |
| 166 | # DS layout: (num_slots, dim, state_len) |
| 167 | state_len = conv_cache.size(-1) |
| 168 | else: |
| 169 | # SD layout: (num_slots, state_len, dim) |
| 170 | state_len = conv_cache.size(-2) |
| 171 | |
| 172 | spec_decode_cache = state_len > (width - 1) |
| 173 | |
| 174 | if not spec_decode_cache: |
| 175 | # ===== Fast path: speculative decoding NOT configured ===== |
| 176 | # conv-state is exactly (..., dim, width-1); use the original AMX / |
| 177 | # torch implementation unchanged. |
| 178 | _cpu_gdn_attention_nonspec( |
| 179 | layer, attn_metadata_i, mixed_qkv, b, a, core_attn_out |
| 180 | ) |
| 181 | return |
| 182 | |
| 183 | # ===== Speculative-decoding-aware path (wide rolling conv buffer) ===== |
| 184 | _cpu_gdn_attention_spec_aware( |
| 185 | layer, attn_metadata_i, mixed_qkv, b, a, core_attn_out, width, state_len |
| 186 | ) |
| 187 | |
| 188 | |
| 189 | # --------------------------------------------------------------------------- |
| 190 | # Non-speculative implementation: conv-state is exactly width-1. |
| 191 | # --------------------------------------------------------------------------- |
| 192 | def _cpu_gdn_attention_nonspec( |
| 193 | layer, |
| 194 | attn_metadata_i: GDNAttentionMetadata, |
| 195 | mixed_qkv: torch.Tensor, |
| 196 | b: torch.Tensor, |
| 197 | a: torch.Tensor, |
| 198 | core_attn_out: torch.Tensor, |
| 199 | ) -> None: |
| 200 | assert ( |
| 201 | attn_metadata_i.spec_sequence_masks is None |
| 202 | and attn_metadata_i.num_accepted_tokens is None |
| 203 | ), "speculative decode not supported without a wide conv-state cache." |
| 204 | |
| 205 | state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor |
| 206 | query_start_loc = attn_metadata_i.non_spec_query_start_loc |
| 207 | assert state_indices_tensor is not None |
| 208 | assert query_start_loc is not None |
| 209 | |
| 210 | # C++ conv (conv.cpp) runs on aarch64 with bf16 support and on x86 uses VDPBF16PS, |
| 211 | # not AMX tiles, so it runs on any AVX-512BF16 CPU. The weights are packed on this |
| 212 | # same predicate at load time. |
| 213 | use_cpp_conv = torch.cpu._is_avx512_bf16_supported() or is_arm_bf16() |
| 214 | |
| 215 | conv_state = layer.kv_cache[0] |
| 216 | if use_cpp_conv: |
| 217 | if not is_conv_state_dim_first(): |
| 218 | conv_state = conv_state.transpose(1, 2) |
| 219 | else: |
| 220 | if not is_conv_state_dim_first(): |
| 221 | conv_state = conv_state.transpose(-1, -2) |
| 222 | conv_weights = layer.conv1d.weight.view( |
| 223 | layer.conv1d.weight.size(0), layer.conv1d.weight.size(2) |
| 224 | ) |
| 225 | |
| 226 | # [num_allocated_slots, num_v_heads / tp_size, v_dim, k_dim] |
| 227 | ssm_state = layer.kv_cache[1] |
| 228 | mixed_qkv = mixed_qkv.contiguous() |
| 229 | a = a.contiguous() |
| 230 | b = b.contiguous() |
| 231 | |
| 232 | num_allocated_slots, head_num, v_dim, k_dim = ssm_state.size() |
| 233 | ssm_state = ssm_state.view(num_allocated_slots, head_num, k_dim, v_dim) |
| 234 | |
| 235 | num_decodes = attn_metadata_i.num_decodes |
| 236 | num_decode_tokens = attn_metadata_i.num_decode_tokens |
| 237 | num_prefills = attn_metadata_i.num_prefills |
| 238 | num_prefill_tokens = attn_metadata_i.num_prefill_tokens |
| 239 | |
| 240 | # all decode requests (batched) |
| 241 | if num_decodes > 0: |
| 242 | decode_mixed_qkv = mixed_qkv[:num_decode_tokens] |
| 243 | decode_b = b[:num_decode_tokens] |
| 244 | decode_a = a[:num_decode_tokens] |
| 245 | decode_state_indices = state_indices_tensor[:num_decodes] |
| 246 | if use_cpp_conv: |
| 247 | decode_mixed_qkv = _causal_conv1d_update_cpu( |
| 248 | x=decode_mixed_qkv, |
| 249 | conv_states=conv_state, |
| 250 | weight=layer.conv1d.weight, |
| 251 | bias=layer.conv1d.bias, |
| 252 | silu_activation=(layer.activation == "silu"), |
| 253 | conv_state_indices=decode_state_indices, |
| 254 | is_weight_packed=True, |
| 255 | ) |
| 256 | else: |
| 257 | decode_mixed_qkv = causal_conv1d_update_cpu( |
| 258 | x=decode_mixed_qkv, |
| 259 | conv_state=conv_state, |
| 260 | weight=conv_weights, |
| 261 | bias=layer.conv1d.bias, |
| 262 | activation=layer.activation, |
| 263 | conv_state_indices=decode_state_indices, |
| 264 | ) |
| 265 | |
| 266 | query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv) |
| 267 | |
| 268 | attn_out = ops.fused_sigmoid_gating_delta_rule_update_cpu( |
| 269 | A_log=layer.A_log, |
| 270 | dt_bias=layer.dt_bias, |
| 271 | q=query, |
| 272 | k=key, |
| 273 | v=value, |
| 274 | a=decode_a, |
| 275 | b=decode_b, |
| 276 | initial_state_source=ssm_state, |
| 277 | initial_state_indices=decode_state_indices, |
| 278 | cu_seqlens=query_start_loc[: num_decodes + 1], |
| 279 | use_qk_l2norm_in_kernel=True, |
| 280 | ) |
| 281 | core_attn_out[:num_decode_tokens] = attn_out.squeeze(1) |
| 282 | |
| 283 | # all prefill requests: (varlen) currently naively loops over sequences |
| 284 | if num_prefills > 0: |
| 285 | has_initial_state = attn_metadata_i.has_initial_state |
| 286 | assert has_initial_state is not None |
| 287 | |
| 288 | prefill_token_start = num_decode_tokens |
| 289 | prefill_token_end = prefill_token_start + num_prefill_tokens |
| 290 | prefill_mixed_qkv = mixed_qkv[prefill_token_start:prefill_token_end] |
| 291 | prefill_b = b[prefill_token_start:prefill_token_end] |
| 292 | prefill_a = a[prefill_token_start:prefill_token_end] |
| 293 | prefill_state_indices = state_indices_tensor[ |
| 294 | num_decodes : num_decodes + num_prefills |
| 295 | ] |
| 296 | prefill_query_start_loc = ( |
| 297 | query_start_loc[num_decodes : num_decodes + num_prefills + 1] |
| 298 | - num_decode_tokens |
| 299 | ) |
| 300 | prefill_has_initial_state = has_initial_state[ |
| 301 | num_decodes : num_decodes + num_prefills |
| 302 | ] |
| 303 | |
| 304 | if use_cpp_conv: |
| 305 | prefill_mixed_qkv = _causal_conv1d_fwd_cpu( |
| 306 | x=prefill_mixed_qkv.transpose(0, 1), |
| 307 | weight=layer.conv1d.weight, |
| 308 | bias=layer.conv1d.bias, |
| 309 | conv_states=conv_state, |
| 310 | query_start_loc=prefill_query_start_loc, |
| 311 | cache_indices=prefill_state_indices, |
| 312 | has_initial_state=prefill_has_initial_state, |
| 313 | silu_activation=layer.activation == "silu", |
| 314 | is_weight_packed=True, |
| 315 | ).transpose(0, 1) |
| 316 | else: |
| 317 | prefill_mixed_qkv = causal_conv1d_torch( |
| 318 | x=prefill_mixed_qkv.transpose(0, 1), |
| 319 | weight=conv_weights, |
| 320 | bias=layer.conv1d.bias, |
| 321 | conv_states=conv_state, |
| 322 | query_start_loc=prefill_query_start_loc, |
| 323 | cache_indices=prefill_state_indices, |
| 324 | has_initial_state=prefill_has_initial_state, |
| 325 | activation=layer.activation, |
| 326 | ).transpose(0, 1) |
| 327 | |
| 328 | query, key, value = layer.rearrange_mixed_qkv(prefill_mixed_qkv) |
| 329 | g, beta = ops.fused_gdn_gating_cpu( |
| 330 | A_log=layer.A_log, a=prefill_a, b=prefill_b, dt_bias=layer.dt_bias |
| 331 | ) |
| 332 | |
| 333 | # zero pool slots for sequences without a prior state; the kernel |
| 334 | # gathers/mutates ssm_state in place via prefill_state_indices, so no |
| 335 | # separate scatter-back is needed after the call. |
| 336 | no_initial_state = prefill_state_indices[~prefill_has_initial_state] |
| 337 | if no_initial_state.numel() > 0: |
| 338 | ssm_state[no_initial_state] = 0 |
| 339 | attn_out, _ = ops.chunk_gated_delta_rule_cpu( |
| 340 | query=query, |
| 341 | key=key, |
| 342 | value=value, |
| 343 | g=g, |
| 344 | beta=beta, |
| 345 | initial_state=ssm_state, |
| 346 | output_final_state=True, |
| 347 | cu_seqlens=prefill_query_start_loc, |
| 348 | head_first=False, |
| 349 | use_qk_l2norm_in_kernel=True, |
| 350 | initial_state_indices=prefill_state_indices, |
| 351 | ) |
| 352 | core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0) |
| 353 | |
| 354 | |
| 355 | # --------------------------------------------------------------------------- |
| 356 | # Speculative-decoding-aware implementation (wide rolling conv buffer). |
| 357 | # |
| 358 | # Mirrors the GPU `_forward_core` spec path using torch + the existing CPU |
| 359 | # kernels. The conv-state is a per-sequence rolling buffer of length |
| 360 | # ``state_len = width - 1 + num_spec`` (single cache slot, column 0 of the |
| 361 | # spec block table), and the SSM recurrent state is stored multi-slot: the |
| 362 | # state *after* draft token ``t`` lives in block-table column ``t``, so the |
| 363 | # next step can resume from the slot of the last accepted token. |
| 364 | # --------------------------------------------------------------------------- |
| 365 | def _conv_buffer_view(layer) -> torch.Tensor: |
| 366 | """Return the conv-state cache as (num_slots, dim, state_len).""" |
| 367 | conv_cache = layer.kv_cache[0] |
| 368 | if is_conv_state_dim_first(): |
| 369 | return conv_cache # (num_slots, dim, state_len) |
| 370 | return conv_cache.transpose(-1, -2) # view -> (num_slots, dim, state_len) |
| 371 | |
| 372 | |
| 373 | def _ssm_state_view(layer) -> torch.Tensor: |
| 374 | ssm_state = layer.kv_cache[1] |
| 375 | num_slots, head_num, v_dim, k_dim = ssm_state.size() |
| 376 | return ssm_state.view(num_slots, head_num, k_dim, v_dim) |
| 377 | |
| 378 | |
| 379 | def _unpacked_conv_weight(layer) -> torch.Tensor: |
| 380 | """Return the plain (dim, width) conv weight. |
| 381 | |
| 382 | On AMX and aarch64+bf16 the conv1d weight is packed in place at load time, |
| 383 | so the torch spec-decode path relies on the un-packed copy stashed by |
| 384 | ``dispatch_cpu_unquantized_gemm``. |
| 385 | """ |
| 386 | w = getattr(layer.conv1d, "_cpu_unpacked_conv_weight", None) |
| 387 | if w is not None: |
| 388 | return w |
| 389 | w = layer.conv1d.weight |
| 390 | if w.dim() == 3: |
| 391 | return w.view(w.size(0), w.size(2)) |
| 392 | return w |
| 393 | |
| 394 | |
| 395 | def _cpu_gdn_attention_spec_aware( |
| 396 | layer, |
| 397 | attn_metadata_i: GDNAttentionMetadata, |
| 398 | mixed_qkv: torch.Tensor, |
| 399 | b: torch.Tensor, |
| 400 | a: torch.Tensor, |
| 401 | core_attn_out: torch.Tensor, |
| 402 | width: int, |
| 403 | state_len: int, |
| 404 | ) -> None: |
| 405 | spec_sequence_masks = attn_metadata_i.spec_sequence_masks |
| 406 | conv_buf = _conv_buffer_view(layer) # (num_slots, dim, state_len) |
| 407 | ssm_state = _ssm_state_view(layer) |
| 408 | |
| 409 | if spec_sequence_masks is None: |
| 410 | # No spec sequences in this batch (e.g. the prompt prefill step while |
| 411 | # speculative decoding is configured). Process as ordinary |
| 412 | # prefill/decode while touching only the first ``width-1`` columns of |
| 413 | # the wide buffer, leaving the rolling history untouched. |
| 414 | _spec_aware_nonspec( |
| 415 | layer, |
| 416 | attn_metadata_i, |
| 417 | mixed_qkv.contiguous(), |
| 418 | b.contiguous(), |
| 419 | a.contiguous(), |
| 420 | core_attn_out, |
| 421 | conv_buf, |
| 422 | ssm_state, |
| 423 | width, |
| 424 | ) |
| 425 | return |
| 426 | |
| 427 | # Split spec vs non-spec tokens. |
| 428 | spec_token_indx = attn_metadata_i.spec_token_indx |
| 429 | non_spec_token_indx = attn_metadata_i.non_spec_token_indx |
| 430 | num_prefills = attn_metadata_i.num_prefills |
| 431 | num_decodes = attn_metadata_i.num_decodes |
| 432 | |
| 433 | if num_prefills == 0 and num_decodes == 0: |
| 434 | mixed_qkv_spec = mixed_qkv.contiguous() |
| 435 | b_spec = b.contiguous() |
| 436 | a_spec = a.contiguous() |
| 437 | spec_out_indx = None |
| 438 | else: |
| 439 | assert spec_token_indx is not None |
| 440 | mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx) |
| 441 | b_spec = b.index_select(0, spec_token_indx) |
| 442 | a_spec = a.index_select(0, spec_token_indx) |
| 443 | spec_out_indx = spec_token_indx |
| 444 | |
| 445 | spec_out = _spec_forward( |
| 446 | layer, |
| 447 | attn_metadata_i, |
| 448 | mixed_qkv_spec, |
| 449 | b_spec, |
| 450 | a_spec, |
| 451 | conv_buf, |
| 452 | ssm_state, |
| 453 | width, |
| 454 | state_len, |
| 455 | ) |
| 456 | |
| 457 | # Process the (rare) non-spec subset (prefill) tokens, if any. |
| 458 | nonspec_out = None |
| 459 | if (num_prefills > 0 or num_decodes > 0) and non_spec_token_indx is not None: |
| 460 | mixed_qkv_ns = mixed_qkv.index_select(0, non_spec_token_indx) |
| 461 | b_ns = b.index_select(0, non_spec_token_indx) |
| 462 | a_ns = a.index_select(0, non_spec_token_indx) |
| 463 | nonspec_out = _spec_aware_nonspec_subset( |
| 464 | layer, attn_metadata_i, mixed_qkv_ns, b_ns, a_ns, conv_buf, ssm_state, width |
| 465 | ) |
| 466 | |
| 467 | # Scatter outputs back to their token positions. |
| 468 | if spec_out_indx is None: |
| 469 | core_attn_out[: spec_out.size(0)] = spec_out |
| 470 | else: |
| 471 | core_attn_out.index_copy_(0, spec_out_indx, spec_out) |
| 472 | if nonspec_out is not None: |
| 473 | core_attn_out.index_copy_(0, non_spec_token_indx, nonspec_out) |
| 474 | |
| 475 | |
| 476 | def _spec_forward( |
| 477 | layer, |
| 478 | attn_metadata_i: GDNAttentionMetadata, |
| 479 | mixed_qkv_spec: torch.Tensor, |
| 480 | b_spec: torch.Tensor, |
| 481 | a_spec: torch.Tensor, |
| 482 | conv_buf: torch.Tensor, |
| 483 | ssm_state: torch.Tensor, |
| 484 | width: int, |
| 485 | state_len: int, |
| 486 | ) -> torch.Tensor: |
| 487 | """Run the GDN core for the multi-query (speculative) tokens. |
| 488 | |
| 489 | Returns the core attention output for the spec tokens, in the same token |
| 490 | order as ``mixed_qkv_spec`` (i.e. ordered by ``spec_query_start_loc``). |
| 491 | """ |
| 492 | num_spec_decodes = attn_metadata_i.num_spec_decodes |
| 493 | spec_state_indices = attn_metadata_i.spec_state_indices_tensor |
| 494 | spec_qsl = attn_metadata_i.spec_query_start_loc |
| 495 | num_accepted = attn_metadata_i.num_accepted_tokens |
| 496 | assert spec_state_indices is not None |
| 497 | assert spec_qsl is not None |
| 498 | assert num_accepted is not None |
| 499 | |
| 500 | spec_qsl_cpu = spec_qsl[: num_spec_decodes + 1] |
| 501 | seq_starts = spec_qsl_cpu[:-1] |
| 502 | seq_lens = spec_qsl_cpu[1:] - spec_qsl_cpu[:-1] |
| 503 | |
| 504 | is_amx = torch.cpu._is_amx_tile_supported() |
| 505 | |
| 506 | # ---- 1. Convolution (per-sequence rolling buffer) ---- |
| 507 | dim = mixed_qkv_spec.size(-1) |
| 508 | bias = layer.conv1d.bias |
| 509 | silu = layer.activation == "silu" |
| 510 | |
| 511 | can_use_native_conv = ( |
| 512 | (is_amx or is_arm_bf16()) |
| 513 | and width == 4 |
| 514 | and num_spec_decodes > 0 |
| 515 | and bool(torch.all(seq_lens == seq_lens[0]).item()) |
| 516 | and int(seq_lens[0].item()) > 0 |
| 517 | ) |
| 518 | if can_use_native_conv: |
| 519 | q_i = int(seq_lens[0].item()) |
| 520 | conv_out = _causal_conv1d_update_cpu( |
| 521 | x=mixed_qkv_spec.view(num_spec_decodes, q_i, dim), |
| 522 | conv_states=conv_buf, |
| 523 | weight=layer.conv1d.weight, |
| 524 | bias=bias, |
| 525 | silu_activation=silu, |
| 526 | conv_state_indices=spec_state_indices[:num_spec_decodes, 0] |
| 527 | .to("cpu", torch.int32) |
| 528 | .contiguous(), |
| 529 | is_weight_packed=True, |
| 530 | num_accepted_tokens=num_accepted[:num_spec_decodes].to("cpu", torch.int32), |
| 531 | ).view_as(mixed_qkv_spec) |
| 532 | else: |
| 533 | w = _unpacked_conv_weight(layer).unsqueeze(1) |
| 534 | col0 = spec_state_indices[:num_spec_decodes, 0] |
| 535 | num_acc_cpu = num_accepted[:num_spec_decodes] |
| 536 | conv_out = torch.empty_like(mixed_qkv_spec) |
| 537 | for i in range(num_spec_decodes): |
| 538 | q_i = int(seq_lens[i].item()) |
| 539 | if q_i == 0: |
| 540 | continue |
| 541 | start = int(seq_starts[i].item()) |
| 542 | slot0 = int(col0[i].item()) |
| 543 | offset = int(num_acc_cpu[i].item()) - 1 |
| 544 | B = conv_buf[slot0] # (dim, state_len) |
| 545 | x_seq = mixed_qkv_spec[start : start + q_i].transpose(0, 1).to(B.dtype) |
| 546 | prior = B[:, offset : offset + (width - 1)] |
| 547 | conv_in = torch.cat([prior, x_seq], dim=-1).unsqueeze(0) |
| 548 | out = F.conv1d(conv_in, w, bias, groups=dim)[0] # (dim, q_i) |
| 549 | if silu: |
| 550 | out = F.silu(out) |
| 551 | conv_out[start : start + q_i] = out.transpose(0, 1).to(conv_out.dtype) |
| 552 | # Roll the buffer: drop the accepted history from the front, append |
| 553 | # the new draft tokens, keep total length == state_len. |
| 554 | keep = B[:, offset + 1 : offset + 1 + (state_len - q_i)] |
| 555 | new_B = torch.cat([keep, x_seq], dim=-1) |
| 556 | B.copy_(new_B) |
| 557 | |
| 558 | # ---- 2. Recurrent (multi-slot SSM state) ---- |
| 559 | # Single fused kernel call: it runs the recurrence over each sequence's |
| 560 | # draft tokens internally, resumes from slot ``num_accepted-1`` and stores |
| 561 | # the state after token ``t`` into slot ``t`` (rollback for the next step). |
| 562 | query, key, value = layer.rearrange_mixed_qkv(conv_out) |
| 563 | query = query.squeeze(0) |
| 564 | key = key.squeeze(0) |
| 565 | value = value.squeeze(0) |
| 566 | spec_idx = spec_state_indices[:num_spec_decodes].to(torch.int32).contiguous() |
| 567 | num_acc = num_accepted[:num_spec_decodes].to(torch.int32).contiguous() |
| 568 | cu = spec_qsl[: num_spec_decodes + 1].to(torch.int32).contiguous() |
| 569 | out_spec = ops.fused_sigmoid_gating_delta_rule_update_spec_cpu( |
| 570 | A_log=layer.A_log, |
| 571 | dt_bias=layer.dt_bias, |
| 572 | q=query, |
| 573 | k=key, |
| 574 | v=value, |
| 575 | a=a_spec, |
| 576 | b=b_spec, |
| 577 | initial_state_source=ssm_state, |
| 578 | spec_state_indices=spec_idx, |
| 579 | num_accepted_tokens=num_acc, |
| 580 | cu_seqlens=cu, |
| 581 | use_qk_l2norm_in_kernel=True, |
| 582 | ) |
| 583 | return out_spec |
| 584 | |
| 585 | |
| 586 | def _spec_aware_nonspec( |
| 587 | layer, |
| 588 | attn_metadata_i: GDNAttentionMetadata, |
| 589 | mixed_qkv: torch.Tensor, |
| 590 | b: torch.Tensor, |
| 591 | a: torch.Tensor, |
| 592 | core_attn_out: torch.Tensor, |
| 593 | conv_buf: torch.Tensor, |
| 594 | ssm_state: torch.Tensor, |
| 595 | width: int, |
| 596 | ) -> None: |
| 597 | """Non-spec prefill/decode with a wide conv buffer.""" |
| 598 | state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor |
| 599 | query_start_loc = attn_metadata_i.non_spec_query_start_loc |
| 600 | assert state_indices_tensor is not None |
| 601 | assert query_start_loc is not None |
| 602 | state_indices_tensor = state_indices_tensor.contiguous() |
| 603 | |
| 604 | can_use_native_conv = torch.cpu._is_amx_tile_supported() or is_arm_bf16() |
| 605 | |
| 606 | if not can_use_native_conv: |
| 607 | conv_weights = _unpacked_conv_weight(layer) |
| 608 | |
| 609 | num_decodes = attn_metadata_i.num_decodes |
| 610 | num_decode_tokens = attn_metadata_i.num_decode_tokens |
| 611 | num_prefills = attn_metadata_i.num_prefills |
| 612 | num_prefill_tokens = attn_metadata_i.num_prefill_tokens |
| 613 | |
| 614 | if num_decodes > 0: |
| 615 | decode_mixed_qkv = mixed_qkv[:num_decode_tokens] |
| 616 | decode_b = b[:num_decode_tokens] |
| 617 | decode_a = a[:num_decode_tokens] |
| 618 | decode_state_indices = state_indices_tensor[:num_decodes] |
| 619 | |
| 620 | if can_use_native_conv: |
| 621 | decode_mixed_qkv = _causal_conv1d_update_cpu( |
| 622 | x=decode_mixed_qkv, |
| 623 | conv_states=conv_buf, |
| 624 | weight=layer.conv1d.weight, |
| 625 | bias=layer.conv1d.bias, |
| 626 | silu_activation=layer.activation == "silu", |
| 627 | conv_state_indices=decode_state_indices, |
| 628 | is_weight_packed=True, |
| 629 | ) |
| 630 | else: |
| 631 | # Only the first ``width-1`` columns hold the real conv state. |
| 632 | conv_state_view = conv_buf[:, :, : width - 1] |
| 633 | decode_mixed_qkv = causal_conv1d_update_cpu( |
| 634 | x=decode_mixed_qkv, |
| 635 | conv_state=conv_state_view, |
| 636 | weight=conv_weights, |
| 637 | bias=layer.conv1d.bias, |
| 638 | activation=layer.activation, |
| 639 | conv_state_indices=decode_state_indices, |
| 640 | ) |
| 641 | |
| 642 | query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv) |
| 643 | attn_out = ops.fused_sigmoid_gating_delta_rule_update_cpu( |
| 644 | A_log=layer.A_log, |
| 645 | dt_bias=layer.dt_bias, |
| 646 | q=query, |
| 647 | k=key, |
| 648 | v=value, |
| 649 | a=decode_a, |
| 650 | b=decode_b, |
| 651 | initial_state_source=ssm_state, |
| 652 | initial_state_indices=decode_state_indices, |
| 653 | cu_seqlens=query_start_loc[: num_decodes + 1], |
| 654 | use_qk_l2norm_in_kernel=True, |
| 655 | ) |
| 656 | core_attn_out[:num_decode_tokens] = attn_out.squeeze(1) |
| 657 | |
| 658 | if num_prefills > 0: |
| 659 | has_initial_state = attn_metadata_i.has_initial_state |
| 660 | assert has_initial_state is not None |
| 661 | prefill_token_start = num_decode_tokens |
| 662 | prefill_token_end = prefill_token_start + num_prefill_tokens |
| 663 | prefill_mixed_qkv = mixed_qkv[prefill_token_start:prefill_token_end] |
| 664 | prefill_b = b[prefill_token_start:prefill_token_end] |
| 665 | prefill_a = a[prefill_token_start:prefill_token_end] |
| 666 | prefill_state_indices = state_indices_tensor[ |
| 667 | num_decodes : num_decodes + num_prefills |
| 668 | ] |
| 669 | prefill_query_start_loc = ( |
| 670 | query_start_loc[num_decodes : num_decodes + num_prefills + 1] |
| 671 | - num_decode_tokens |
| 672 | ) |
| 673 | prefill_has_initial_state = has_initial_state[ |
| 674 | num_decodes : num_decodes + num_prefills |
| 675 | ] |
| 676 | |
| 677 | if can_use_native_conv: |
| 678 | prefill_mixed_qkv = _causal_conv1d_fwd_cpu( |
| 679 | x=prefill_mixed_qkv.transpose(0, 1), |
| 680 | weight=layer.conv1d.weight, |
| 681 | bias=layer.conv1d.bias, |
| 682 | conv_states=conv_buf, |
| 683 | query_start_loc=prefill_query_start_loc, |
| 684 | cache_indices=prefill_state_indices, |
| 685 | has_initial_state=prefill_has_initial_state, |
| 686 | silu_activation=layer.activation == "silu", |
| 687 | is_weight_packed=True, |
| 688 | ).transpose(0, 1) |
| 689 | else: |
| 690 | prefill_mixed_qkv = causal_conv1d_torch( |
| 691 | x=prefill_mixed_qkv.transpose(0, 1), |
| 692 | weight=conv_weights, |
| 693 | bias=layer.conv1d.bias, |
| 694 | conv_states=conv_buf, |
| 695 | query_start_loc=prefill_query_start_loc, |
| 696 | cache_indices=prefill_state_indices, |
| 697 | has_initial_state=prefill_has_initial_state, |
| 698 | activation=layer.activation, |
| 699 | ).transpose(0, 1) |
| 700 | |
| 701 | query, key, value = layer.rearrange_mixed_qkv(prefill_mixed_qkv) |
| 702 | g, beta = ops.fused_gdn_gating_cpu( |
| 703 | A_log=layer.A_log, a=prefill_a, b=prefill_b, dt_bias=layer.dt_bias |
| 704 | ) |
| 705 | # zero pool slots for sequences without a prior state; the kernel |
| 706 | # gathers/mutates ssm_state in place via prefill_state_indices, so no |
| 707 | # separate scatter-back is needed after the call. |
| 708 | no_initial_state = prefill_state_indices[~prefill_has_initial_state] |
| 709 | if no_initial_state.numel() > 0: |
| 710 | ssm_state[no_initial_state] = 0 |
| 711 | attn_out, _ = ops.chunk_gated_delta_rule_cpu( |
| 712 | query=query, |
| 713 | key=key, |
| 714 | value=value, |
| 715 | g=g, |
| 716 | beta=beta, |
| 717 | initial_state=ssm_state, |
| 718 | output_final_state=True, |
| 719 | cu_seqlens=prefill_query_start_loc, |
| 720 | head_first=False, |
| 721 | use_qk_l2norm_in_kernel=True, |
| 722 | initial_state_indices=prefill_state_indices, |
| 723 | ) |
| 724 | core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0) |
| 725 | |
| 726 | |
| 727 | def _spec_aware_nonspec_subset( |
| 728 | layer, |
| 729 | attn_metadata_i: GDNAttentionMetadata, |
| 730 | mixed_qkv: torch.Tensor, |
| 731 | b: torch.Tensor, |
| 732 | a: torch.Tensor, |
| 733 | conv_buf: torch.Tensor, |
| 734 | ssm_state: torch.Tensor, |
| 735 | width: int, |
| 736 | ) -> torch.Tensor: |
| 737 | """Process non-spec (prefill) tokens that coexist with spec sequences. |
| 738 | |
| 739 | Returns outputs ordered like ``non_spec_token_indx``. |
| 740 | """ |
| 741 | has_initial_state = attn_metadata_i.has_initial_state |
| 742 | prefill_state_indices = attn_metadata_i.prefill_state_indices |
| 743 | prefill_qsl = attn_metadata_i.prefill_query_start_loc |
| 744 | assert prefill_state_indices is not None and prefill_qsl is not None |
| 745 | assert has_initial_state is not None |
| 746 | prefill_state_indices = prefill_state_indices.contiguous() |
| 747 | |
| 748 | is_amx = torch.cpu._is_amx_tile_supported() |
| 749 | can_use_native_conv = is_amx or is_arm_bf16() |
| 750 | |
| 751 | if can_use_native_conv: |
| 752 | conv_out = _causal_conv1d_fwd_cpu( |
| 753 | x=mixed_qkv.transpose(0, 1), |
| 754 | weight=layer.conv1d.weight, |
| 755 | bias=layer.conv1d.bias, |
| 756 | conv_states=conv_buf, |
| 757 | query_start_loc=prefill_qsl, |
| 758 | cache_indices=prefill_state_indices, |
| 759 | has_initial_state=has_initial_state, |
| 760 | silu_activation=layer.activation == "silu", |
| 761 | is_weight_packed=True, |
| 762 | ).transpose(0, 1) |
| 763 | else: |
| 764 | conv_weights = _unpacked_conv_weight(layer) |
| 765 | conv_out = causal_conv1d_torch( |
| 766 | x=mixed_qkv.transpose(0, 1), |
| 767 | weight=conv_weights, |
| 768 | bias=layer.conv1d.bias, |
| 769 | conv_states=conv_buf, |
| 770 | query_start_loc=prefill_qsl, |
| 771 | cache_indices=prefill_state_indices, |
| 772 | has_initial_state=has_initial_state, |
| 773 | activation=layer.activation, |
| 774 | ).transpose(0, 1) |
| 775 | |
| 776 | query, key, value = layer.rearrange_mixed_qkv(conv_out) |
| 777 | g, beta = ops.fused_gdn_gating_cpu( |
| 778 | A_log=layer.A_log, a=a, b=b, dt_bias=layer.dt_bias |
| 779 | ) |
| 780 | # zero pool slots for sequences without a prior state; the kernel |
| 781 | # gathers/mutates ssm_state in place via prefill_state_indices, so no |
| 782 | # separate scatter-back is needed after the call. |
| 783 | no_initial_state = prefill_state_indices[~has_initial_state] |
| 784 | if no_initial_state.numel() > 0: |
| 785 | ssm_state[no_initial_state] = 0 |
| 786 | attn_out, _ = ops.chunk_gated_delta_rule_cpu( |
| 787 | query=query, |
| 788 | key=key, |
| 789 | value=value, |
| 790 | g=g, |
| 791 | beta=beta, |
| 792 | initial_state=ssm_state, |
| 793 | output_final_state=True, |
| 794 | cu_seqlens=prefill_qsl, |
| 795 | head_first=False, |
| 796 | use_qk_l2norm_in_kernel=True, |
| 797 | initial_state_indices=prefill_state_indices, |
| 798 | ) |
| 799 | return attn_out.squeeze(0) |
| 800 | |
| 801 | |
| 802 | def register_cpu_gdn_attention_ops() -> None: |
| 803 | global _CPU_GDN_ATTENTION_OPS_REGISTERED |
| 804 | if _CPU_GDN_ATTENTION_OPS_REGISTERED: |
| 805 | return |
| 806 | |
| 807 | direct_register_custom_op( |
| 808 | op_name="cpu_gdn_attention_core", |
| 809 | op_func=cpu_gdn_attention_core, |
| 810 | mutates_args=["core_attn_out"], |
| 811 | ) |
| 812 | _CPU_GDN_ATTENTION_OPS_REGISTERED = True |