| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """CPU DeepSeek-V4 sparse-MLA backend descriptor. |
| 4 | |
| 5 | DeepSeek-V4 attention runs entirely through |
| 6 | ``DeepseekV4Attention``/``DeepseekV4CPUAttention``, never through a generic |
| 7 | ``AttentionImpl``. ``DeepseekV4CPUSparseBackend`` exists so the CPU attention |
| 8 | layer has its own backend name and its own metadata-builder class: CPU is not |
| 9 | allowed to run triton-cpu inside the model's forward path, so |
| 10 | ``DeepseekV4CPUFlashMLAMetadataBuilder`` overrides the shared base's |
| 11 | ``_build_c128a_metadata`` (triton-backed in |
| 12 | ``DeepseekV4SparseMLAMetadataBuilder``) with a plain-eager-PyTorch |
| 13 | reimplementation of the same per-token block-table-resolution math. |
| 14 | """ |
| 15 | |
| 16 | import torch |
| 17 | |
| 18 | from vllm.models.deepseek_v4.sparse_mla import ( |
| 19 | DeepseekV4SparseMLABackend, |
| 20 | DeepseekV4SparseMLAMetadataBuilder, |
| 21 | ) |
| 22 | from vllm.v1.attention.backend import CommonAttentionMetadata |
| 23 | from vllm.v1.attention.backends.mla.indexer import ( |
| 24 | DeepseekV4IndexerBackend, |
| 25 | DeepseekV32IndexerMetadataBuilder, |
| 26 | ) |
| 27 | from vllm.v1.attention.backends.mla.sparse_swa import ( |
| 28 | _LAYER_TYPE_C4A, |
| 29 | _LAYER_TYPE_C128A, |
| 30 | _LAYER_TYPE_SWAONLY, |
| 31 | DeepseekSparseSWAMetadataBuilder, |
| 32 | ) |
| 33 | from vllm.v1.attention.ops.flashmla import FlashMLASchedMeta |
| 34 | |
| 35 | |
| 36 | class DeepseekV4CPUFlashMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder): |
| 37 | """CPU sparse-MLA metadata builder: same fields as the shared base, but |
| 38 | skips the C128A dense-topk metadata entirely -- unlike CUDA/XPU/ROCm, |
| 39 | ``DeepseekV4CPUAttention.forward_mqa`` never reads |
| 40 | ``c128a_global_decode_topk_indices``/``c128a_decode_topk_lens``/ |
| 41 | ``c128a_prefill_topk_indices``; it recomputes the same local top-k |
| 42 | directly from ``positions`` instead, so building it here (whether |
| 43 | eagerly or via the shared base's triton kernel) would be wasted work. |
| 44 | """ |
| 45 | |
| 46 | def _build_c128a_metadata( |
| 47 | self, |
| 48 | cm: CommonAttentionMetadata, |
| 49 | req_id_per_token: torch.Tensor, |
| 50 | ) -> dict[str, torch.Tensor | None]: |
| 51 | return {} |
| 52 | |
| 53 | |
| 54 | class DeepseekV4CPUSparseBackend(DeepseekV4SparseMLABackend): |
| 55 | @staticmethod |
| 56 | def get_name() -> str: |
| 57 | return "CPU_V4_MLA_SPARSE" |
| 58 | |
| 59 | @staticmethod |
| 60 | def get_builder_cls() -> type[DeepseekV4CPUFlashMLAMetadataBuilder]: |
| 61 | return DeepseekV4CPUFlashMLAMetadataBuilder |
| 62 | |
| 63 | |
| 64 | class DeepseekV4CPUIndexerMetadataBuilder(DeepseekV32IndexerMetadataBuilder): |
| 65 | """CPU indexer metadata builder: same fields as the shared base, but |
| 66 | prefill requests are never split into multiple chunks. |
| 67 | |
| 68 | ``DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks`` (the |
| 69 | shared base's default) bounds two things the CUDA/XPU indexer kernels |
| 70 | need: the flat-gather workspace size and the dense M*N logits tensor |
| 71 | those kernels allocate for a chunk. The CPU indexer |
| 72 | (``sparse_attn_indexer_cpu``) allocates neither -- it reads the paged |
| 73 | K-cache directly via ``fp8_paged_mqa_logits_cpu``/ |
| 74 | ``topk_transform_512_cpu``, with no chunk-splitting of its own. |
| 75 | With chunked prefill now enabled for this model on CPU (see |
| 76 | ``CpuPlatform.check_and_update_config``'s ``amx_mla_or_dsv4_enabled``), |
| 77 | the total query-token count for one step is already bounded by |
| 78 | ``max_num_batched_tokens`` -- there is nothing left here to bound, so |
| 79 | this always returns a single chunk spanning the whole step's prefill |
| 80 | batch. |
| 81 | """ |
| 82 | |
| 83 | @staticmethod |
| 84 | def _split_indexer_prefill_chunks( |
| 85 | compressed_seq_lens_cpu: torch.Tensor, |
| 86 | prefill_query_lens_cpu: torch.Tensor, |
| 87 | workspace_size: int, |
| 88 | max_logits_bytes: int, |
| 89 | request_offset: int = 0, |
| 90 | ) -> list[tuple[slice, slice]]: |
| 91 | # workspace_size/max_logits_bytes are unused: no flat-gather |
| 92 | # workspace or dense M*N buffer is allocated. |
| 93 | num_reqs = compressed_seq_lens_cpu.shape[0] |
| 94 | total_query_len = int(prefill_query_lens_cpu.sum().item()) |
| 95 | return [ |
| 96 | ( |
| 97 | slice(request_offset, request_offset + num_reqs), |
| 98 | slice(0, total_query_len), |
| 99 | ) |
| 100 | ] |
| 101 | |
| 102 | |
| 103 | class DeepseekV4CPUIndexerBackend(DeepseekV4IndexerBackend): |
| 104 | @staticmethod |
| 105 | def get_name() -> str: |
| 106 | return "CPU_V4_INDEXER" |
| 107 | |
| 108 | @staticmethod |
| 109 | def get_builder_cls() -> type[DeepseekV4CPUIndexerMetadataBuilder]: |
| 110 | return DeepseekV4CPUIndexerMetadataBuilder |
| 111 | |
| 112 | |
| 113 | class DeepseekV4CPUSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder): |
| 114 | """CPU SWA metadata builder: same fields as the shared base.""" |
| 115 | |
| 116 | def build_tile_scheduler( |
| 117 | self, num_decode_tokens: int |
| 118 | ) -> dict[str, FlashMLASchedMeta | None]: |
| 119 | """CPU never runs the FlashMLA tile-scheduler planner (that's a CUDA |
| 120 | C++ decode-path concern) -- always return the all-``None`` sentinel |
| 121 | the shared base's own CPU branch returns.""" |
| 122 | return { |
| 123 | _LAYER_TYPE_SWAONLY: None, |
| 124 | _LAYER_TYPE_C4A: None, |
| 125 | _LAYER_TYPE_C128A: None, |
| 126 | } |