vllm.model_executor.layers.mamba.ops.cpu.gdn_attention ¶

def _conv_buffer_view(layer) -> torch.Tensor:
    """Return the conv-state cache as (num_slots, dim, state_len)."""
    conv_cache = layer.kv_cache[0]
    if is_conv_state_dim_first():
        return conv_cache  # (num_slots, dim, state_len)
    return conv_cache.transpose(-1, -2)  # view -> (num_slots, dim, state_len)

def _spec_aware_nonspec(
    layer,
    attn_metadata_i: GDNAttentionMetadata,
    mixed_qkv: torch.Tensor,
    b: torch.Tensor,
    a: torch.Tensor,
    core_attn_out: torch.Tensor,
    conv_buf: torch.Tensor,
    ssm_state: torch.Tensor,
    width: int,
) -> None:
    """Non-spec prefill/decode with a wide conv buffer."""
    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
    state_indices_tensor = state_indices_tensor.contiguous()

    is_amx = torch.cpu._is_amx_tile_supported()
    if is_amx and is_conv_state_dim_first():
        raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.")

    if not is_amx:
        conv_weights = _unpacked_conv_weight(layer)

    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

    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]
        if is_amx:
            decode_mixed_qkv = ops.causal_conv1d_update_cpu(
                x=decode_mixed_qkv,
                conv_states=conv_buf,
                weight=layer.conv1d.weight,
                bias=layer.conv1d.bias,
                silu_activation=layer.activation == "silu",
                conv_state_indices=decode_state_indices,
                is_vnni=True,
            )
        else:
            # Only the first ``width-1`` columns hold the real conv state.
            conv_state_view = conv_buf[:, :, : width - 1]
            if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
                decode_conv_state = conv_state_view[decode_state_indices].contiguous()
                decode_mixed_qkv = causal_conv1d_update_torch(
                    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_view[decode_state_indices] = decode_conv_state
            else:
                decode_mixed_qkv = causal_conv1d_update_cpu(
                    x=decode_mixed_qkv,
                    conv_state=conv_state_view,
                    weight=conv_weights,
                    bias=layer.conv1d.bias,
                    activation=layer.activation,
                    conv_state_indices=decode_state_indices,
                )

        query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv)
        attn_out = ops.fused_sigmoid_gating_delta_rule_update_cpu(
            A_log=layer.A_log,
            dt_bias=layer.dt_bias,
            q=query,
            k=key,
            v=value,
            a=decode_a,
            b=decode_b,
            initial_state_source=ssm_state,
            initial_state_indices=decode_state_indices,
            cu_seqlens=query_start_loc[: num_decodes + 1],
            use_qk_l2norm_in_kernel=True,
        )
        core_attn_out[:num_decode_tokens] = attn_out.squeeze(1)

    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
        ]
        if is_amx:
            prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu(
                x=prefill_mixed_qkv.transpose(0, 1),
                weight=layer.conv1d.weight,
                bias=layer.conv1d.bias,
                conv_states=conv_buf,
                query_start_loc=prefill_query_start_loc,
                cache_indices=prefill_state_indices,
                has_initial_state=prefill_has_initial_state,
                silu_activation=layer.activation == "silu",
                is_vnni=True,
            ).transpose(0, 1)
        else:
            prefill_mixed_qkv = causal_conv1d_torch(
                x=prefill_mixed_qkv.transpose(0, 1),
                weight=conv_weights,
                bias=layer.conv1d.bias,
                conv_states=conv_buf,
                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 = ops.fused_gdn_gating_cpu(
            A_log=layer.A_log, a=prefill_a, b=prefill_b, dt_bias=layer.dt_bias
        )
        # zero pool slots for sequences without a prior state; the kernel
        # gathers/mutates ssm_state in place via prefill_state_indices, so no
        # separate scatter-back is needed after the call.
        no_initial_state = prefill_state_indices[~prefill_has_initial_state]
        if no_initial_state.numel() > 0:
            ssm_state[no_initial_state] = 0
        attn_out, _ = ops.chunk_gated_delta_rule_cpu(
            query=query,
            key=key,
            value=value,
            g=g,
            beta=beta,
            initial_state=ssm_state,
            output_final_state=True,
            cu_seqlens=prefill_query_start_loc,
            head_first=False,
            use_qk_l2norm_in_kernel=True,
            initial_state_indices=prefill_state_indices,
        )
        core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0)

def _spec_aware_nonspec_subset(
    layer,
    attn_metadata_i: GDNAttentionMetadata,
    mixed_qkv: torch.Tensor,
    b: torch.Tensor,
    a: torch.Tensor,
    conv_buf: torch.Tensor,
    ssm_state: torch.Tensor,
    width: int,
) -> torch.Tensor:
    """Process non-spec (prefill) tokens that coexist with spec sequences.

    Returns outputs ordered like ``non_spec_token_indx``.
    """
    has_initial_state = attn_metadata_i.has_initial_state
    prefill_state_indices = attn_metadata_i.prefill_state_indices
    prefill_qsl = attn_metadata_i.prefill_query_start_loc
    assert prefill_state_indices is not None and prefill_qsl is not None
    assert has_initial_state is not None
    prefill_state_indices = prefill_state_indices.contiguous()

    is_amx = torch.cpu._is_amx_tile_supported()
    if is_amx and is_conv_state_dim_first():
        raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.")

    if is_amx:
        conv_out = ops.causal_conv1d_fwd_cpu(
            x=mixed_qkv.transpose(0, 1),
            weight=layer.conv1d.weight,
            bias=layer.conv1d.bias,
            conv_states=conv_buf,
            query_start_loc=prefill_qsl,
            cache_indices=prefill_state_indices,
            has_initial_state=has_initial_state,
            silu_activation=layer.activation == "silu",
            is_vnni=True,
        ).transpose(0, 1)
    else:
        conv_weights = _unpacked_conv_weight(layer)
        conv_out = causal_conv1d_torch(
            x=mixed_qkv.transpose(0, 1),
            weight=conv_weights,
            bias=layer.conv1d.bias,
            conv_states=conv_buf,
            query_start_loc=prefill_qsl,
            cache_indices=prefill_state_indices,
            has_initial_state=has_initial_state,
            activation=layer.activation,
        ).transpose(0, 1)

    query, key, value = layer.rearrange_mixed_qkv(conv_out)
    g, beta = ops.fused_gdn_gating_cpu(
        A_log=layer.A_log, a=a, b=b, dt_bias=layer.dt_bias
    )
    # zero pool slots for sequences without a prior state; the kernel
    # gathers/mutates ssm_state in place via prefill_state_indices, so no
    # separate scatter-back is needed after the call.
    no_initial_state = prefill_state_indices[~has_initial_state]
    if no_initial_state.numel() > 0:
        ssm_state[no_initial_state] = 0
    attn_out, _ = ops.chunk_gated_delta_rule_cpu(
        query=query,
        key=key,
        value=value,
        g=g,
        beta=beta,
        initial_state=ssm_state,
        output_final_state=True,
        cu_seqlens=prefill_qsl,
        head_first=False,
        use_qk_l2norm_in_kernel=True,
        initial_state_indices=prefill_state_indices,
    )
    return attn_out.squeeze(0)

def _spec_forward(
    layer,
    attn_metadata_i: GDNAttentionMetadata,
    mixed_qkv_spec: torch.Tensor,
    b_spec: torch.Tensor,
    a_spec: torch.Tensor,
    conv_buf: torch.Tensor,
    ssm_state: torch.Tensor,
    width: int,
    state_len: int,
) -> torch.Tensor:
    """Run the GDN core for the multi-query (speculative) tokens.

    Returns the core attention output for the spec tokens, in the same token
    order as ``mixed_qkv_spec`` (i.e. ordered by ``spec_query_start_loc``).
    """
    num_spec_decodes = attn_metadata_i.num_spec_decodes
    spec_state_indices = attn_metadata_i.spec_state_indices_tensor
    spec_qsl = attn_metadata_i.spec_query_start_loc
    num_accepted = attn_metadata_i.num_accepted_tokens
    assert spec_state_indices is not None
    assert spec_qsl is not None
    assert num_accepted is not None

    spec_qsl_cpu = spec_qsl[: num_spec_decodes + 1]
    seq_starts = spec_qsl_cpu[:-1]
    seq_lens = spec_qsl_cpu[1:] - spec_qsl_cpu[:-1]

    # ---- 1. Convolution (per-sequence rolling buffer) ----
    dim = mixed_qkv_spec.size(-1)
    bias = layer.conv1d.bias
    silu = layer.activation == "silu"

    can_use_native_conv = (
        torch.cpu._is_amx_tile_supported()
        and not is_conv_state_dim_first()
        and width == 4
        and num_spec_decodes > 0
        and bool(torch.all(seq_lens == seq_lens[0]).item())
        and int(seq_lens[0].item()) > 0
    )
    if can_use_native_conv:
        q_i = int(seq_lens[0].item())
        conv_out = ops.causal_conv1d_update_cpu(
            x=mixed_qkv_spec.view(num_spec_decodes, q_i, dim),
            conv_states=conv_buf,
            weight=layer.conv1d.weight,
            bias=bias,
            silu_activation=silu,
            conv_state_indices=spec_state_indices[:num_spec_decodes, 0]
            .to("cpu", torch.int32)
            .contiguous(),
            is_vnni=True,
            num_accepted_tokens=num_accepted[:num_spec_decodes].to("cpu", torch.int32),
        ).view_as(mixed_qkv_spec)
    else:
        w = _unpacked_conv_weight(layer).unsqueeze(1)
        col0 = spec_state_indices[:num_spec_decodes, 0]
        num_acc_cpu = num_accepted[:num_spec_decodes]
        conv_out = torch.empty_like(mixed_qkv_spec)
        for i in range(num_spec_decodes):
            q_i = int(seq_lens[i].item())
            if q_i == 0:
                continue
            start = int(seq_starts[i].item())
            slot0 = int(col0[i].item())
            offset = int(num_acc_cpu[i].item()) - 1
            B = conv_buf[slot0]  # (dim, state_len)
            x_seq = mixed_qkv_spec[start : start + q_i].transpose(0, 1).to(B.dtype)
            prior = B[:, offset : offset + (width - 1)]
            conv_in = torch.cat([prior, x_seq], dim=-1).unsqueeze(0)
            out = F.conv1d(conv_in, w, bias, groups=dim)[0]  # (dim, q_i)
            if silu:
                out = F.silu(out)
            conv_out[start : start + q_i] = out.transpose(0, 1).to(conv_out.dtype)
            # Roll the buffer: drop the accepted history from the front, append
            # the new draft tokens, keep total length == state_len.
            keep = B[:, offset + 1 : offset + 1 + (state_len - q_i)]
            new_B = torch.cat([keep, x_seq], dim=-1)
            B.copy_(new_B)

    # ---- 2. Recurrent (multi-slot SSM state) ----
    # Single fused kernel call: it runs the recurrence over each sequence's
    # draft tokens internally, resumes from slot ``num_accepted-1`` and stores
    # the state after token ``t`` into slot ``t`` (rollback for the next step).
    query, key, value = layer.rearrange_mixed_qkv(conv_out)
    query = query.squeeze(0)
    key = key.squeeze(0)
    value = value.squeeze(0)
    spec_idx = spec_state_indices[:num_spec_decodes].to(torch.int32).contiguous()
    num_acc = num_accepted[:num_spec_decodes].to(torch.int32).contiguous()
    cu = spec_qsl[: num_spec_decodes + 1].to(torch.int32).contiguous()
    out_spec = ops.fused_sigmoid_gating_delta_rule_update_spec_cpu(
        A_log=layer.A_log,
        dt_bias=layer.dt_bias,
        q=query,
        k=key,
        v=value,
        a=a_spec,
        b=b_spec,
        initial_state_source=ssm_state,
        spec_state_indices=spec_idx,
        num_accepted_tokens=num_acc,
        cu_seqlens=cu,
        use_qk_l2norm_in_kernel=True,
    )
    return out_spec

def _unpacked_conv_weight(layer) -> torch.Tensor:
    """Return the plain (dim, width) conv weight.

    On AMX the conv1d weight is VNNI-packed in place at load time (only usable
    by the AMX C++ kernel), so the torch spec-decode path relies on the
    un-packed copy stashed by ``dispatch_cpu_unquantized_gemm``.
    """
    w = getattr(layer.conv1d, "_cpu_unpacked_conv_weight", None)
    if w is not None:
        return w
    w = layer.conv1d.weight
    if w.dim() == 3:
        return w.view(w.size(0), w.size(2))
    return w

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 mixed_qkv.dtype == torch.bfloat16, "CPU GDN attention requires BF16."

    # The conv-state cache temporal dim is ``width - 1`` when speculative
    # decoding is disabled, and ``width - 1 + num_spec`` when it is enabled
    # (vLLM allocates a wider rolling buffer to support draft-token rollback).
    width = layer.conv1d.weight.size(-1)
    conv_cache = layer.kv_cache[0]
    if is_conv_state_dim_first():
        # DS layout: (num_slots, dim, state_len)
        state_len = conv_cache.size(-1)
    else:
        # SD layout: (num_slots, state_len, dim)
        state_len = conv_cache.size(-2)

    spec_decode_cache = state_len > (width - 1)

    if not spec_decode_cache:
        # ===== Fast path: speculative decoding NOT configured =====
        # conv-state is exactly (..., dim, width-1); use the original AMX /
        # torch implementation unchanged.
        _cpu_gdn_attention_nonspec(
            layer, attn_metadata_i, mixed_qkv, b, a, core_attn_out
        )
        return

    # ===== Speculative-decoding-aware path (wide rolling conv buffer) =====
    _cpu_gdn_attention_spec_aware(
        layer, attn_metadata_i, mixed_qkv, b, a, core_attn_out, width, state_len
    )

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