vllm.models.deepseek_v4.cpu.cpu_sparse ¶

class DeepseekV4CPUAttention(DeepseekV4Attention):
    """CPU sparse MLA attention layer for DeepSeek V4."""

    backend_cls = DeepseekV4CPUSparseBackend

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

        # Swap classes post-construction (no class-swap hook on the shared
        # base) so forward()/get_attn_backend() dispatch to CPU kernels.
        # SWA-only layers (compress_ratio <= 1) never build a
        # compressor/indexer.
        if self.compressor is not None:
            self.compressor.__class__ = DeepseekV4CPUCompressor
        if self.indexer is not None:
            self.indexer.__class__ = DeepseekV4CPUIndexer
            self.indexer.k_cache.__class__ = DeepseekV4CPUIndexerCache
            self.indexer.compressor.__class__ = DeepseekV4CPUCompressor

        self._wo_a_packed: torch.Tensor | None = None
        self._wrap_wo_a_process_weights_after_loading()

    def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
        """Cache fp32-contiguous copies of tensors CPU kernels want that way
        but that arrive bf16 (rotary cos/sin table, compressor RMSNorm
        weights) -- cast once here instead of per forward call.

        Runs after every quantized layer's own
        ``process_weights_after_loading`` (see ``is_deferred_attention_layer``);
        ``wo_a``'s packing must happen *before* that phase instead, so it's a
        separate monkeypatch in ``__init__``.
        """
        self.rotary_emb.cos_sin_cache = self.rotary_emb.cos_sin_cache.to(
            torch.float32
        ).contiguous()
        if self.compressor is not None:
            self.compressor.cache_norm_weight_fp32()
        if self.indexer is not None:
            self.indexer.compressor.cache_norm_weight_fp32()

    def _wrap_wo_a_process_weights_after_loading(self) -> None:
        """Snapshot and pack ``wo_a``'s weight into a bf16 copy for
        ``bmm_cpu`` before the FP8 kernel's own
        ``process_weights_after_loading`` VNNI-repacks it in place for
        row-major reads ``_o_proj`` never does."""
        wo_a = self.wo_a
        quant_method = wo_a.quant_method
        orig_pwal = quant_method.process_weights_after_loading
        heads_per_group = self.n_local_heads // self.n_local_groups
        k_dim = heads_per_group * self.head_dim

        def _capture_and_pack(layer: torch.nn.Module) -> None:
            dequant = _dequant_linear_weight(layer).to(torch.bfloat16)
            dequant = dequant.view(self.n_local_groups, self.o_lora_rank, k_dim)
            self._wo_a_packed = torch.ops._C.convert_weight_packed(dequant.contiguous())
            orig_pwal(layer)
            # orig_pwal VNNI-repacks layer.weight/scale for ordinary-linear
            # use that _o_proj never needs -- free them, same pattern as
            # dispatch_cpu_unquantized_gemm's remove_weight=True.
            layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
            scale_attr = (
                "weight_scale_inv"
                if hasattr(layer, "weight_scale_inv")
                else "weight_scale"
            )
            if hasattr(layer, scale_attr):
                setattr(
                    layer,
                    scale_attr,
                    torch.nn.Parameter(torch.empty(0), requires_grad=False),
                )

        quant_method.process_weights_after_loading = (  # type: ignore[method-assign]
            _capture_and_pack
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        llama_4_scaling: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """CPU override: wraps attention-prep + MLA + output projection in the
        CPU-only opaque custom op (see ``_deepseek_v4_cpu_prepare_and_attn``).
        """
        num_tokens = hidden_states.shape[0]
        o_padded = torch.empty(
            (num_tokens, self.padded_heads, self.head_dim),
            dtype=hidden_states.dtype,
            device=hidden_states.device,
        )

        qr_kv = self._fused_wqa_wkv_gemm(hidden_states)

        kv_score: torch.Tensor | None = None
        if self.compressor is not None:
            # Must go through the module's own forward (not a raw matmul on
            # its weight): the packed-GEMM kernel frees the raw weight after
            # packing, same as wo_a's.
            kv_score = self.compressor.fused_wkv_wgate(hidden_states).to(torch.float32)

        indexer_weights: torch.Tensor | None = None
        indexer_kv_score: torch.Tensor | None = None
        if self.indexer is not None:
            indexer_weights, _ = self.indexer.weights_proj(hidden_states)
            indexer_kv_score = self.indexer.compressor.fused_wkv_wgate(
                hidden_states
            ).to(torch.float32)

        qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv)

        return torch.ops.vllm.deepseek_v4_cpu_prepare_and_attn(
            hidden_states,
            qr,
            kv,
            qr_scale,
            kv_score,
            indexer_kv_score,
            indexer_weights,
            positions,
            o_padded,
            _encode_layer_name(self.prefix),
        )

    def _fused_qnorm_rope_kv_insert(
        self,
        q: torch.Tensor,
        kv: torch.Tensor,
        positions: torch.Tensor,
        attn_metadata: (
            dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None
        ),
    ) -> torch.Tensor:
        """CPU override: only the fp8_ds_mla (uint8) SWA cache layout is
        supported here, so the base method's bf16/per-tensor-fp8 branches
        are dropped."""
        if not isinstance(attn_metadata, dict):
            # Profile run: kernel doesn't fire; produce a padded tensor so
            # downstream FlashMLA gets the right shape.
            if self.n_local_heads < self.padded_heads:
                return F.pad(
                    q,
                    (0, 0, 0, self.padded_heads - self.n_local_heads),
                    value=0.0,
                )
            return q

        swa_metadata = cast(
            "DeepseekSparseSWAMetadata | None",
            attn_metadata.get(self.swa_cache_layer.prefix),
        )
        assert swa_metadata is not None

        swa_kv_cache = self.swa_cache_layer.kv_cache
        # positions is already int64 (the runner's buffer), no cast needed.
        assert positions.dtype == torch.int64
        assert swa_kv_cache.dtype == torch.uint8, (
            "DeepseekV4CPUAttention only supports the fp8_ds_mla cache layout"
        )
        # Fused: Q side does per-head weight-free RMSNorm + GPT-J RoPE
        # (zero-filling padded head slots); KV side does GPT-J RoPE + UE8M0
        # FP8 quant + paged insert.
        swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
        return ops.fused_qnorm_rope_kv_insert_cpu(
            q,
            kv,
            positions,
            swa_kv_cache_2d,
            swa_metadata.slot_mapping,
            self.rotary_emb.cos_sin_cache,
            self.padded_heads,
            self.eps,
            swa_metadata.block_size,
        )

    def _split_qkv_and_norm(
        self, qr_kv: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
        """CPU override: two ``RMSNorm`` calls instead of the shared
        ``fused_q_kv_rmsnorm``, whose raw ``@triton.jit`` kernel fails to
        link under triton-cpu (``undefined symbol: __truncdfbf2``)."""
        qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
        return self.q_norm(qr), None, self.kv_norm(kv)

    def _prepare_and_attn(
        self,
        hidden_states: torch.Tensor,
        qr: torch.Tensor,
        kv: torch.Tensor,
        qr_scale: torch.Tensor | None,
        kv_score: torch.Tensor,
        indexer_kv_score: torch.Tensor,
        indexer_weights: torch.Tensor,
        positions: torch.Tensor,
        o_padded: torch.Tensor,
    ) -> None:
        """CPU override: no aux streams, so query-projection+KV-insert, the
        indexer, and the compressor run straight-line instead of through
        execute_in_parallel/maybe_execute_in_parallel."""
        attn_metadata = get_forward_context().attn_metadata
        indexer = self.indexer
        compressor = self.compressor

        q = self._wq_b_proj(qr, qr_scale).view(-1, self.n_local_heads, self.head_dim)
        q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)

        index_q: torch.Tensor | None = None
        index_q_scale: torch.Tensor | None = None
        index_weights_out: torch.Tensor | None = None
        if indexer is not None:
            assert compressor is not None
            index_q, index_q_scale, index_weights_out = indexer(
                hidden_states,
                qr,
                indexer_kv_score,
                indexer_weights,
                positions,
                self.indexer_rotary_emb,
                qr_scale,
            )
            compressor(kv_score, positions, self.rotary_emb)
        elif compressor is not None:
            compressor(kv_score, positions, self.rotary_emb)

        self._sparse_indexer_and_attn(
            hidden_states,
            index_q,
            index_q_scale,
            index_weights_out,
            q,
            kv,
            positions,
            o_padded,
        )

    def _sparse_indexer_and_attn(
        self,
        hidden_states: torch.Tensor,
        index_q: torch.Tensor | None,
        index_q_scale: torch.Tensor | None,
        index_weights: torch.Tensor | None,
        q: torch.Tensor,
        kv: torch.Tensor,
        positions: torch.Tensor,
        out: torch.Tensor,
    ) -> None:
        """CPU override: identical body, minus ``@eager_break_during_capture``
        (this platform never captures a CUDA graph, so it was a no-op)."""
        if self.indexer is not None and index_q is not None:
            assert index_weights is not None
            # index_q_scale is always None on CPU (FP8-only path).
            assert index_q_scale is None
            self.indexer.indexer_op(hidden_states, index_q, None, index_weights)

        self.forward_mqa(q, kv, positions, out)

    @classmethod
    def get_padded_num_q_heads(cls, num_heads: int) -> int:
        # No head-count padding constraint for the (not yet ported) CPU kernel.
        return num_heads

    def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
        if get_forward_context().attn_metadata is None:
            # Warmup/profile run: no real metadata, o is a zeroed placeholder.
            return torch.zeros(
                (o.shape[0], self.hidden_size), dtype=o.dtype, device=o.device
            )

        num_tokens, num_heads, _ = o.shape
        cos_sin_cache = self.rotary_emb.cos_sin_cache
        # Undo the query's own GPT-J RoPE rotation: the attention output is
        # a weighted sum of KV rows at different positions, so R(-pos_q)
        # recovers the translation-invariant quantity wo_a/wo_b expect.
        o_derot = ops.inverse_gptj_rope_o_proj_cpu(
            o, positions, cos_sin_cache, _ROPE_DIM
        )

        heads_per_group = num_heads // self.n_local_groups
        o_grouped = o_derot.reshape(
            num_tokens, self.n_local_groups, heads_per_group * self.head_dim
        ).to(torch.bfloat16)
        assert self._wo_a_packed is not None
        z = o_grouped.new_empty((self.n_local_groups, num_tokens, self.o_lora_rank))
        ops.bmm_cpu(
            z,
            o_grouped.transpose(0, 1).contiguous(),
            self._wo_a_packed,
            True,
            None,
        )
        z = z.transpose(0, 1).to(o.dtype)
        return self.wo_b(z.flatten(1))

    def forward_mqa(
        self,
        q: torch.Tensor,
        kv: torch.Tensor,
        positions: torch.Tensor,
        output: torch.Tensor,
    ) -> None:
        forward_context = get_forward_context()
        attn_metadata = forward_context.attn_metadata
        if attn_metadata is None:
            # Warmup/profile run: reserve nothing extra, skip real kernels.
            output.zero_()
            return

        assert isinstance(attn_metadata, dict)
        swa_metadata = cast(
            DeepseekSparseSWAMetadata,
            attn_metadata.get(self.swa_cache_layer.prefix),
        )
        assert swa_metadata is not None
        assert swa_metadata.token_to_req_indices is not None
        req_idx = swa_metadata.token_to_req_indices

        swa_only = self.compress_ratio <= 1
        flashmla_metadata = cast(
            DeepseekV4FlashMLAMetadata | None, attn_metadata.get(self.prefix)
        )

        num_tokens = q.shape[0]
        device = q.device

        # ---- SWA window: causal sliding window over the raw per-token cache
        window = self.window_size
        win_offsets = torch.arange(window, device=device)
        start_pos = (positions - window + 1).clamp(min=0)
        local_pos = start_pos.unsqueeze(-1) + win_offsets
        valid_win = local_pos <= positions.unsqueeze(-1)
        win_local = torch.where(valid_win, local_pos, torch.full_like(local_pos, -1))
        window_slots = map_local_to_global_slots_cpu(
            win_local, req_idx, swa_metadata.block_table, swa_metadata.block_size
        )
        window_cache_2d = self.swa_cache_layer.kv_cache.view(
            self.swa_cache_layer.kv_cache.shape[0], -1
        )

        # ---- Compressed KV: C4A's learned top-k, or C128A's dense top-k ----
        if swa_only:
            compressed_slots = window_slots.new_full((num_tokens, 0), -1)
            # Unused when num_compressed == 0 (only dtype/2D-ness matter);
            # reuse the window cache to avoid an allocation.
            compressed_cache_2d = window_cache_2d
            compressed_block_size = 1
        else:
            assert flashmla_metadata is not None
            main_block_size = flashmla_metadata.block_size // self.compress_ratio
            if self.compress_ratio == 4:
                assert self.topk_indices_buffer is not None
                local_topk = self.topk_indices_buffer[:num_tokens]
            else:
                width = -(-self.max_model_len // self.compress_ratio)
                num_compressed = torch.div(
                    positions + 1, self.compress_ratio, rounding_mode="floor"
                ).clamp(max=width)
                arange_w = torch.arange(width, device=device)
                local_topk = torch.where(
                    arange_w.unsqueeze(0) < num_compressed.unsqueeze(-1),
                    arange_w.unsqueeze(0).expand(num_tokens, width),
                    torch.full(
                        (num_tokens, width), -1, device=device, dtype=torch.int64
                    ),
                )
            compressed_slots = map_local_to_global_slots_cpu(
                local_topk, req_idx, flashmla_metadata.block_table, main_block_size
            )
            compressed_cache_2d = self.kv_cache.view(self.kv_cache.shape[0], -1)
            compressed_block_size = main_block_size

        torch.ops._C.flash_mla_with_kvcache_cpu(
            output,
            q,
            window_cache_2d,
            window_slots,
            swa_metadata.block_size,
            compressed_cache_2d,
            compressed_slots,
            compressed_block_size,
            self.attn_sink,
            self.scale,
        )

class DeepseekV4CPUIndexer(DeepseekV4Indexer):
    """CPU indexer: the C4A short-context fallback runs as eager PyTorch
    instead of Triton. Never constructed directly -- ``__class__`` is
    swapped onto an existing ``DeepseekV4Indexer`` instance post-construction.
    """

    def forward(
        self,
        hidden_states: torch.Tensor,
        qr: torch.Tensor,
        compressed_kv_score: torch.Tensor,
        indexer_weights: torch.Tensor,
        positions: torch.Tensor,
        rotary_emb: nn.Module,
        qr_scale: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
        """CPU override: no aux streams, so wq_b_and_q_quant and the
        compressor run straight-line instead of through
        ``maybe_execute_in_parallel``; the short-context check drops the
        CUDA-only ``is_current_stream_capturing()`` guard (always False
        here)."""
        compressor = self.compressor

        attn_metadata = get_forward_context().attn_metadata
        if isinstance(attn_metadata, dict):
            indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix])
            if indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens:
                # Fewer candidates than topk: all are selected, but the K
                # cache still needs to be built.
                compressor(compressed_kv_score, positions, rotary_emb)
                assert self.topk_indices_buffer is not None
                num_tokens = (
                    indexer_metadata.num_decode_tokens
                    + indexer_metadata.num_prefill_tokens
                )
                if num_tokens > 0:
                    self._fill_short_context_topk(positions, num_tokens)
                return None, None, None

        def wq_b_and_q_quant():
            q = self._wq_b_proj(qr, qr_scale)
            q = q.view(-1, self.n_head, self.head_dim)
            return _fused_indexer_q_rope_quant_cpu(
                positions,
                q,
                rotary_emb.cos_sin_cache,
                indexer_weights,
                self.softmax_scale,
                self.n_head**-0.5,
            )

        # Order matches the base class's parallel-join contract: compressor
        # must complete before indexer_op runs (skip_k_cache_insert=True).
        q, weights = wq_b_and_q_quant()
        compressor(compressed_kv_score, positions, rotary_emb)
        # No q_scale: FP8-only path (MXFP4 needs Blackwell).
        return q, None, weights

    def _fill_short_context_topk(
        self, positions: torch.Tensor, num_tokens: int
    ) -> None:
        assert self.topk_indices_buffer is not None
        device = positions.device
        arange = torch.arange(self.topk_tokens, device=device)
        num_compressed = torch.div(
            positions[:num_tokens] + 1, self.compress_ratio, rounding_mode="floor"
        )
        valid = arange.unsqueeze(0) < num_compressed.unsqueeze(-1)
        filled = torch.where(
            valid, arange.unsqueeze(0), torch.full_like(arange.unsqueeze(0), -1)
        )
        self.topk_indices_buffer[:num_tokens, : self.topk_tokens] = filled.to(
            self.topk_indices_buffer.dtype
        )

class DeepseekV4CPUIndexerCache(DeepseekV4IndexerCache):
    """CPU indexer K-cache descriptor: same fields as the shared base, just
    pointing ``get_attn_backend()`` at ``DeepseekV4CPUIndexerBackend``
    instead of the shared, CUDA/XPU-oriented ``DeepseekV4IndexerBackend``."""

    def get_attn_backend(self) -> type[AttentionBackend]:
        return DeepseekV4CPUIndexerBackend

class _NoOpEvent:
    """Stand-in for ``torch.cuda.Event`` on CPU: since CPU never sets an aux
    stream list, these events are constructed but never actually used."""