vllm.models.deepseek_v4.cpu.cpu_mla ¶

class DeepseekV4CPUFlashMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):
    """CPU sparse-MLA metadata builder: same fields as the shared base, but
    skips the C128A dense-topk metadata entirely -- unlike CUDA/XPU/ROCm,
    ``DeepseekV4CPUAttention.forward_mqa`` never reads
    ``c128a_global_decode_topk_indices``/``c128a_decode_topk_lens``/
    ``c128a_prefill_topk_indices``; it recomputes the same local top-k
    directly from ``positions`` instead, so building it here (whether
    eagerly or via the shared base's triton kernel) would be wasted work.
    """

    def _build_c128a_metadata(
        self,
        cm: CommonAttentionMetadata,
        req_id_per_token: torch.Tensor,
    ) -> dict[str, torch.Tensor | None]:
        return {}

class DeepseekV4CPUIndexerMetadataBuilder(DeepseekV32IndexerMetadataBuilder):
    """CPU indexer metadata builder: same fields as the shared base, but
    prefill requests are never split into multiple chunks.

    ``DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks`` (the
    shared base's default) bounds two things the CUDA/XPU indexer kernels
    need: the flat-gather workspace size and the dense M*N logits tensor
    those kernels allocate for a chunk. The CPU indexer
    (``sparse_attn_indexer_cpu``) allocates neither -- it reads the paged
    K-cache directly via ``fp8_paged_mqa_logits_cpu``/
    ``topk_transform_512_cpu``, with no chunk-splitting of its own.
    With chunked prefill now enabled for this model on CPU (see
    ``CpuPlatform.check_and_update_config``'s ``amx_mla_or_dsv4_enabled``),
    the total query-token count for one step is already bounded by
    ``max_num_batched_tokens`` -- there is nothing left here to bound, so
    this always returns a single chunk spanning the whole step's prefill
    batch.
    """

    @staticmethod
    def _split_indexer_prefill_chunks(
        compressed_seq_lens_cpu: torch.Tensor,
        prefill_query_lens_cpu: torch.Tensor,
        workspace_size: int,
        max_logits_bytes: int,
        request_offset: int = 0,
    ) -> list[tuple[slice, slice]]:
        # workspace_size/max_logits_bytes are unused: no flat-gather
        # workspace or dense M*N buffer is allocated.
        num_reqs = compressed_seq_lens_cpu.shape[0]
        total_query_len = int(prefill_query_lens_cpu.sum().item())
        return [
            (
                slice(request_offset, request_offset + num_reqs),
                slice(0, total_query_len),
            )
        ]

class DeepseekV4CPUSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder):
    """CPU SWA metadata builder: same fields as the shared base."""

    def build_tile_scheduler(
        self, num_decode_tokens: int
    ) -> dict[str, FlashMLASchedMeta | None]:
        """CPU never runs the FlashMLA tile-scheduler planner (that's a CUDA
        C++ decode-path concern) -- always return the all-``None`` sentinel
        the shared base's own CPU branch returns.
        """
        return {
            _LAYER_TYPE_SWAONLY: None,
            _LAYER_TYPE_C4A: None,
            _LAYER_TYPE_C128A: None,
        }