| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import field |
| 5 | from typing import Any, Literal |
| 6 | |
| 7 | from pydantic import Field, field_validator |
| 8 | |
| 9 | from vllm.config.utils import config |
| 10 | from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum |
| 11 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 12 | |
| 13 | IndexerKVDType = Literal["auto", "bf16", "fp8", "mxfp4", "nvfp4"] |
| 14 | MiniMaxM3MSADecodeBackend = Literal["triton", "cutlass"] |
| 15 | |
| 16 | |
| 17 | @config |
| 18 | class HiSparseConfig: |
| 19 | """Configuration for HiSparse sparse-MLA KV offloading.""" |
| 20 | |
| 21 | device_buffer_size: int | None = Field(default=None, gt=0) |
| 22 | """Total per-request GPU hot-buffer rows, including the newest-token slot. |
| 23 | |
| 24 | Defaults to one top-k per decode query plus one top-k of LRU slack. The |
| 25 | physical allocation is rounded up to the GPU cache block size selected |
| 26 | from the active backends. |
| 27 | """ |
| 28 | |
| 29 | eager_host_mirror: bool = True |
| 30 | """Mirror decode-written KV rows to the host pool during the forward so |
| 31 | page spills complete without moving data. When disabled, decode rows stay |
| 32 | resident-only and evicted pages are copied to host at spill time. Prefill |
| 33 | rows are always mirrored during the forward.""" |
| 34 | |
| 35 | |
| 36 | @config |
| 37 | class AttentionConfig: |
| 38 | """Configuration for attention mechanisms in vLLM.""" |
| 39 | |
| 40 | backend: AttentionBackendEnum | None = None |
| 41 | """Attention backend to use. Use "auto" or None for automatic selection.""" |
| 42 | |
| 43 | minimax_m3_msa_decode_backend: MiniMaxM3MSADecodeBackend = "triton" |
| 44 | """Sparse decode kernel used by the MiniMax M3 MSA backend.""" |
| 45 | |
| 46 | backend_per_kind: dict[str, AttentionBackendEnum] = field(default_factory=dict) |
| 47 | """Per-KV-cache-group attention backend overrides, keyed by |
| 48 | `KVCacheSpecKind` (e.g. `{"mla_attention": "FLASHINFER_MLA", |
| 49 | "sliding_window_mla": "TRITON_MLA"}`). This lets a model that splits its |
| 50 | layers across multiple KV-cache groups (e.g. interleaved full and |
| 51 | sliding-window attention) use a different backend per group. |
| 52 | |
| 53 | An entry overrides `backend` for layers of the matching kind; kinds not |
| 54 | listed fall back to `backend` (or automatic selection). A selected backend |
| 55 | that is invalid for that kind raises at startup.""" |
| 56 | |
| 57 | flash_attn_version: Literal[2, 3, 4] | None = None |
| 58 | """Force vllm to use a specific flash-attention version (2, 3, or 4). |
| 59 | Only valid when using the flash-attention backend.""" |
| 60 | |
| 61 | flash_attn_max_num_splits_for_cuda_graph: int = 32 |
| 62 | """Flash Attention max number splits for cuda graph decode.""" |
| 63 | |
| 64 | tq_max_kv_splits_for_cuda_graph: int = 32 |
| 65 | """TurboQuant max NUM_KV_SPLITS for cuda graph decode. |
| 66 | Fixes the split count so grid dimensions are constant across captures, |
| 67 | and buffers can be pre-allocated to avoid inflating the memory estimate.""" |
| 68 | |
| 69 | use_trtllm_attention: bool | None = None |
| 70 | """If set to True/False, use or don't use the TRTLLM attention backend |
| 71 | in flashinfer. If None, auto-detect the attention backend in flashinfer.""" |
| 72 | |
| 73 | disable_flashinfer_q_quantization: bool = False |
| 74 | """If set, when using fp8 kv, do not quantize Q to fp8.""" |
| 75 | |
| 76 | mla_prefill_backend: MLAPrefillBackendEnum | None = None |
| 77 | """MLA prefill backend to use. If None, will be selected automatically. |
| 78 | Valid options: FLASH_ATTN (FA3/FA4), FLASHINFER, TRTLLM_RAGGED.""" |
| 79 | |
| 80 | use_prefill_query_quantization: bool = False |
| 81 | """If set, quantize query for attention in prefill.""" |
| 82 | |
| 83 | indexer_kv_dtype: IndexerKVDType = "auto" |
| 84 | """Data type for the sparse-attention indexer K cache. "auto" picks the |
| 85 | model's default (bf16 for MiniMax M3, fp8 for the DeepSeek sparse |
| 86 | indexer). Quantized formats (fp8, mxfp4, nvfp4) require indexer kernel |
| 87 | support in the backend.""" |
| 88 | |
| 89 | indexer_sparse_logits: bool = False |
| 90 | """DeepSeek V4.1 two-level indexer: score only the candidate blocks with |
| 91 | DeepGEMM's sparse MQA-logits kernels instead of computing dense logits over |
| 92 | the whole context and masking them. Requires `indexer_kv_dtype="mxfp4"`, |
| 93 | an SM100-class GPU, DeepGEMM >= 2.8 and the DeepSelect top-k extension |
| 94 | (the top-k runs on the kernels' bf16 logits). The sparse path costs |
| 95 | O(candidate blocks) per query regardless of context length, so it pays off |
| 96 | for long contexts (roughly 32K tokens and beyond) and is slower below. |
| 97 | On ROCm gfx950 it runs aiter's paged MXFP4 MQA-logits kernel over the |
| 98 | candidate pool instead, with fp32 logits, and keeps the masked dense walk |
| 99 | for steps whose contexts are too short for the pool to pay.""" |
| 100 | |
| 101 | hisparse_config: HiSparseConfig | None = None |
| 102 | """HiSparse host-resident KV configuration. Setting this enables experimental |
| 103 | Model Runner V2-only HiSparse sparse-MLA decode hot-buffering. It is inferred |
| 104 | with defaults when HiSparseConnector is configured (directly or via |
| 105 | MultiConnector); set it explicitly only to tune its fields.""" |
| 106 | |
| 107 | use_non_causal: bool = False |
| 108 | """Whether to use non-causal (bidirectional) attention.""" |
| 109 | |
| 110 | sparse_mla_force_mqa: bool = False |
| 111 | """Force sparse MLA to use forward_mqa for all requests, including prefill. |
| 112 | When False (default), pure prefill batches use forward_mha when implemented. |
| 113 | Set to True to always use the MQA path.""" |
| 114 | |
| 115 | flex_attn_block_m: int | None = None |
| 116 | """Triton kernel BLOCK_M tile size for flex attention. |
| 117 | Must be a power of 2 >= 16. If None and VLLM_BATCH_INVARIANT=1, |
| 118 | defaults to 16.""" |
| 119 | |
| 120 | flex_attn_block_n: int | None = None |
| 121 | """Triton kernel BLOCK_N tile size for flex attention. |
| 122 | Must be a power of 2 >= 16. If None and VLLM_BATCH_INVARIANT=1, |
| 123 | defaults to 16.""" |
| 124 | |
| 125 | flex_attn_q_block_size: int | None = None |
| 126 | """Logical Q block size for the flex attention block mask. |
| 127 | Must be a power of 2 and divisible by flex_attn_block_m. |
| 128 | If None, uses 16 for paged KV attention on PyTorch >= 2.9, and 128 |
| 129 | for encoder-only attention or older PyTorch versions.""" |
| 130 | |
| 131 | flex_attn_kv_block_size: int | None = None |
| 132 | """Logical KV block size for the flex attention block mask. |
| 133 | Must be a power of 2 and divisible by flex_attn_block_n. |
| 134 | If None, uses the KV cache block size for paged KV attention on |
| 135 | PyTorch >= 2.9, and 128 for encoder-only attention or older PyTorch |
| 136 | versions.""" |
| 137 | |
| 138 | def __post_init__(self) -> None: |
| 139 | msa_aliases: dict[AttentionBackendEnum, MiniMaxM3MSADecodeBackend] = { |
| 140 | AttentionBackendEnum.CUTLASS_MSA: "cutlass", |
| 141 | AttentionBackendEnum.TRITON_MSA: "triton", |
| 142 | } |
| 143 | if self.backend in msa_aliases: |
| 144 | self.minimax_m3_msa_decode_backend = msa_aliases[self.backend] |
| 145 | # The alias selects only MiniMax's sparse decode kernel. Dense |
| 146 | # layers still use the platform's normal automatic backend. |
| 147 | self.backend = None |
| 148 | |
| 149 | def resolve_indexer_kv_dtype(self, default: IndexerKVDType) -> IndexerKVDType: |
| 150 | """Resolve `indexer_kv_dtype`, substituting `default` for "auto".""" |
| 151 | if self.indexer_kv_dtype == "auto": |
| 152 | return default |
| 153 | return self.indexer_kv_dtype |
| 154 | |
| 155 | def compute_hash(self) -> str: |
| 156 | """Provide a hash that uniquely identifies all the configs |
| 157 | that affect the structure of the computation |
| 158 | graph from input ids/embeddings to the final hidden states, |
| 159 | excluding anything before input ids/embeddings and after |
| 160 | the final hidden states. |
| 161 | """ |
| 162 | from vllm.config.utils import get_hash_factors, hash_factors |
| 163 | |
| 164 | factors = get_hash_factors(self, set()) |
| 165 | return hash_factors(factors) |
| 166 | |
| 167 | @field_validator("backend", mode="before") |
| 168 | @classmethod |
| 169 | def validate_backend_before(cls, value: Any) -> Any: |
| 170 | """Enable parsing of the `backend` enum type from string. |
| 171 | |
| 172 | The special value "auto" is treated as None, which triggers |
| 173 | automatic backend selection. |
| 174 | """ |
| 175 | if isinstance(value, str): |
| 176 | if value.lower() == "auto": |
| 177 | return None |
| 178 | return AttentionBackendEnum[value.upper()] |
| 179 | return value |
| 180 | |
| 181 | @field_validator("mla_prefill_backend", mode="before") |
| 182 | @classmethod |
| 183 | def validate_mla_prefill_backend_before(cls, value: Any) -> Any: |
| 184 | """Enable parsing of the `mla_prefill_backend` enum type from string.""" |
| 185 | if isinstance(value, str): |
| 186 | return MLAPrefillBackendEnum[value.upper()] |
| 187 | return value |
| 188 | |
| 189 | @field_validator("backend_per_kind", mode="before") |
| 190 | @classmethod |
| 191 | def validate_backend_per_kind_before(cls, value: Any) -> Any: |
| 192 | """Parse the `backend_per_kind` map from strings. |
| 193 | |
| 194 | Keys must be valid `KVCacheSpecKind` values; values are parsed like |
| 195 | `backend` (enum name, case-insensitive). |
| 196 | """ |
| 197 | from vllm.v1.kv_cache_interface import KVCacheSpecKind |
| 198 | |
| 199 | if not isinstance(value, dict): |
| 200 | return value |
| 201 | valid_kinds = {kind.value for kind in KVCacheSpecKind} |
| 202 | parsed: dict[str, AttentionBackendEnum] = {} |
| 203 | for kind, backend in value.items(): |
| 204 | if kind not in valid_kinds: |
| 205 | raise ValueError( |
| 206 | f"Unknown KV cache group kind '{kind}' in " |
| 207 | f"backend_per_kind. Valid kinds are: " |
| 208 | f"{', '.join(sorted(valid_kinds))}." |
| 209 | ) |
| 210 | if isinstance(backend, str): |
| 211 | backend = AttentionBackendEnum[backend.upper()] |
| 212 | parsed[kind] = backend |
| 213 | return parsed |