| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """MiniMax CUTLASS sparse decode using per-query-token page indices.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | from dataclasses import dataclass, field |
| 8 | from typing import Any |
| 9 | |
| 10 | import torch |
| 11 | |
| 12 | from vllm.config.attention import MiniMaxM3MSADecodeBackend |
| 13 | from vllm.platforms import current_platform |
| 14 | from vllm.triton_utils import tl, triton |
| 15 | from vllm.v1.kv_cache_interface import get_kv_quant_mode |
| 16 | |
| 17 | _MAX_NUM_Q_HEADS = 64 |
| 18 | _MAX_NUM_KV_HEADS = 4 |
| 19 | _HEAD_DIM = 128 |
| 20 | _PAGE_SIZE = 128 |
| 21 | _TOPK = 16 |
| 22 | # fmha_sm100 plans one row per query head. Keep every cached plan within the |
| 23 | # fixed planner allocation used by the MSA decode kernel. |
| 24 | _MAX_QUERY_HEAD_ROWS = 65536 |
| 25 | _MAX_DECODE_QUERY_LEN = 32 |
| 26 | # Kernel benchmarks put the CUTLASS crossover at 16 requests for TP1 and TP4. |
| 27 | # NVFP4 has no Triton fallback, so it takes CUTLASS at every batch size. |
| 28 | _MIN_CUTLASS_BATCH_SIZE = 16 |
| 29 | |
| 30 | |
| 31 | def is_nvfp4_kv_cache(kv_cache_dtype: str) -> bool: |
| 32 | return get_kv_quant_mode(kv_cache_dtype).is_nvfp4 |
| 33 | |
| 34 | |
| 35 | def nvfp4_kv_cache_views( |
| 36 | kv_cache: torch.Tensor, |
| 37 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 38 | """Return the K data, K scale, V data and V scale views of an NVFP4 cache. |
| 39 | |
| 40 | ``kv_cache`` is uint8 ``[pages, 2 * Hkv, page_size, 72]``; slot |
| 41 | ``2 * head + side`` (K = 0, V = 1) is that head's E2M1 data |
| 42 | ``[page_size, 64]`` followed by its E4M3 block scales ``[page_size, 8]``, |
| 43 | the one page layout ``fmha_sm100`` reads. |
| 44 | """ |
| 45 | from vllm.third_party.fmha_sm100.nvfp4_kv import nvfp4_head_slot_views |
| 46 | |
| 47 | return nvfp4_head_slot_views(*nvfp4_kv_cache_slots(kv_cache)) |
| 48 | |
| 49 | |
| 50 | def nvfp4_kv_cache_slots( |
| 51 | kv_cache: torch.Tensor, |
| 52 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 53 | """Return the per-head K and V slots of an NVFP4 cache.""" |
| 54 | return kv_cache[:, 0::2], kv_cache[:, 1::2] |
| 55 | |
| 56 | |
| 57 | @dataclass |
| 58 | class MSACutlassDecodeMetadata: |
| 59 | plan: Any |
| 60 | page_table: torch.Tensor |
| 61 | |
| 62 | |
| 63 | @triton.jit |
| 64 | def _update_runtime_metadata_kernel( |
| 65 | seq_lens_ptr, |
| 66 | kv_segment_lens_ptr, |
| 67 | qo_offset_ptr, |
| 68 | num_rows: tl.constexpr, |
| 69 | decode_query_len: tl.constexpr, |
| 70 | BLOCK_SIZE: tl.constexpr, |
| 71 | ): |
| 72 | offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 73 | mask = offsets < num_rows |
| 74 | request = offsets // decode_query_len |
| 75 | local_query = offsets % decode_query_len |
| 76 | seq_len = tl.load(seq_lens_ptr + request, mask=mask) |
| 77 | tl.store(kv_segment_lens_ptr + offsets, seq_len, mask=mask) |
| 78 | tl.store( |
| 79 | qo_offset_ptr + offsets, |
| 80 | seq_len - decode_query_len + local_query, |
| 81 | mask=mask, |
| 82 | ) |
| 83 | |
| 84 | |
| 85 | @dataclass |
| 86 | class MSACutlassDecodePlanCache: |
| 87 | """Reusable plans whose mutable tensors retain cudagraph-stable addresses.""" |
| 88 | |
| 89 | plans: dict[tuple[int, ...], Any] = field(init=False, default_factory=dict) |
| 90 | |
| 91 | def _build_plan( |
| 92 | self, |
| 93 | *, |
| 94 | batch: int, |
| 95 | decode_query_len: int, |
| 96 | page_table_stride: int, |
| 97 | initial_seq_lens_cpu: torch.Tensor, |
| 98 | device: torch.device, |
| 99 | num_q_heads: int, |
| 100 | num_kv_heads: int, |
| 101 | page_size: int, |
| 102 | topk_blocks: int, |
| 103 | ) -> Any: |
| 104 | from vllm.third_party.fmha_sm100.api import fmha_sm100_plan |
| 105 | |
| 106 | qo_lens_cpu = torch.full((batch,), decode_query_len, dtype=torch.int32) |
| 107 | kv_lens_cpu = initial_seq_lens_cpu |
| 108 | plan = fmha_sm100_plan( |
| 109 | qo_lens_cpu, |
| 110 | kv_lens_cpu, |
| 111 | num_q_heads, |
| 112 | num_kv_heads=num_kv_heads, |
| 113 | qo_offset=kv_lens_cpu - qo_lens_cpu, |
| 114 | page_size=page_size, |
| 115 | output_maxscore=False, |
| 116 | kv_block_num=topk_blocks, |
| 117 | causal=True, |
| 118 | sparse_kernel_mode="decode", |
| 119 | use_fp8_kvcache=True, |
| 120 | split_prefill_decode=False, |
| 121 | device=device, |
| 122 | # The Q8KV4 route keeps plan-time lengths and its own page offsets, |
| 123 | # which this cached plan does not refresh. |
| 124 | decode_backend="kv_mode3", |
| 125 | ) |
| 126 | |
| 127 | plan_info = plan[3] |
| 128 | row_starts = ( |
| 129 | torch.arange(batch, dtype=torch.int32, device=device) |
| 130 | .mul_(page_table_stride) |
| 131 | .repeat_interleave(decode_query_len) |
| 132 | ) |
| 133 | page_indptr = torch.cat( |
| 134 | ( |
| 135 | row_starts, |
| 136 | torch.tensor( |
| 137 | [batch * page_table_stride], |
| 138 | dtype=torch.int32, |
| 139 | device=device, |
| 140 | ), |
| 141 | ) |
| 142 | ) |
| 143 | plan_info["kv_page_indptr"].copy_(page_indptr) |
| 144 | return plan |
| 145 | |
| 146 | def prepare( |
| 147 | self, |
| 148 | block_table: torch.Tensor, |
| 149 | seq_lens: torch.Tensor, |
| 150 | seq_lens_cpu: torch.Tensor, |
| 151 | decode_query_len: int, |
| 152 | *, |
| 153 | num_q_heads: int, |
| 154 | num_kv_heads: int, |
| 155 | page_size: int, |
| 156 | topk_blocks: int, |
| 157 | ) -> MSACutlassDecodeMetadata: |
| 158 | batch = int(seq_lens.shape[0]) |
| 159 | if ( |
| 160 | block_table.device.type != "cuda" |
| 161 | or block_table.dtype != torch.int32 |
| 162 | or not block_table.is_contiguous() |
| 163 | or block_table.shape[0] != batch |
| 164 | ): |
| 165 | raise ValueError( |
| 166 | "MSA sparse decode requires a contiguous CUDA int32 block " |
| 167 | "table with one row per request" |
| 168 | ) |
| 169 | if ( |
| 170 | seq_lens.dtype != torch.int32 |
| 171 | or not seq_lens.is_contiguous() |
| 172 | or seq_lens.device != block_table.device |
| 173 | ): |
| 174 | raise ValueError( |
| 175 | "MSA sparse decode requires contiguous CUDA int32 sequence " |
| 176 | "lengths on the block table device" |
| 177 | ) |
| 178 | if ( |
| 179 | seq_lens_cpu.device.type != "cpu" |
| 180 | or seq_lens_cpu.dtype != torch.int32 |
| 181 | or not seq_lens_cpu.is_contiguous() |
| 182 | or seq_lens_cpu.shape != seq_lens.shape |
| 183 | ): |
| 184 | raise ValueError( |
| 185 | "MSA sparse decode requires contiguous CPU int32 sequence " |
| 186 | "lengths matching the device sequence lengths" |
| 187 | ) |
| 188 | |
| 189 | page_table_stride = int(block_table.stride(0)) |
| 190 | key = ( |
| 191 | batch, |
| 192 | decode_query_len, |
| 193 | page_table_stride, |
| 194 | num_q_heads, |
| 195 | num_kv_heads, |
| 196 | page_size, |
| 197 | topk_blocks, |
| 198 | ) |
| 199 | plan = self.plans.get(key) |
| 200 | if plan is None: |
| 201 | plan = self._build_plan( |
| 202 | batch=batch, |
| 203 | decode_query_len=decode_query_len, |
| 204 | page_table_stride=page_table_stride, |
| 205 | initial_seq_lens_cpu=seq_lens_cpu, |
| 206 | device=seq_lens.device, |
| 207 | num_q_heads=num_q_heads, |
| 208 | num_kv_heads=num_kv_heads, |
| 209 | page_size=page_size, |
| 210 | topk_blocks=topk_blocks, |
| 211 | ) |
| 212 | self.plans[key] = plan |
| 213 | |
| 214 | plan_info = plan[3] |
| 215 | num_rows = batch * decode_query_len |
| 216 | _update_runtime_metadata_kernel[(triton.cdiv(num_rows, 128),)]( |
| 217 | seq_lens, |
| 218 | plan_info["kv_segment_lens"], |
| 219 | plan_info["qo_offset"], |
| 220 | num_rows=num_rows, |
| 221 | decode_query_len=decode_query_len, |
| 222 | BLOCK_SIZE=128, |
| 223 | ) |
| 224 | return MSACutlassDecodeMetadata( |
| 225 | plan=plan, |
| 226 | page_table=block_table.view(-1), |
| 227 | ) |
| 228 | |
| 229 | |
| 230 | def _supported_head_geometry(num_q_heads: int, num_kv_heads: int) -> bool: |
| 231 | return ( |
| 232 | 0 < num_q_heads <= _MAX_NUM_Q_HEADS |
| 233 | and 0 < num_kv_heads <= _MAX_NUM_KV_HEADS |
| 234 | and num_q_heads % num_kv_heads == 0 |
| 235 | ) |
| 236 | |
| 237 | |
| 238 | def supports_cutlass_sparse_decode( |
| 239 | *, |
| 240 | decode_backend: MiniMaxM3MSADecodeBackend, |
| 241 | num_q_heads: int, |
| 242 | num_kv_heads: int, |
| 243 | kv_cache_dtype: str, |
| 244 | page_size: int, |
| 245 | topk_blocks: int, |
| 246 | ) -> bool: |
| 247 | """Return whether static model geometry supports CUTLASS sparse decode.""" |
| 248 | nvfp4 = is_nvfp4_kv_cache(kv_cache_dtype) |
| 249 | return ( |
| 250 | (decode_backend == "cutlass" or nvfp4) |
| 251 | and current_platform.is_cuda() |
| 252 | and current_platform.is_device_capability_family(100) |
| 253 | and (kv_cache_dtype in ("fp8", "fp8_e4m3") or nvfp4) |
| 254 | and _supported_head_geometry(num_q_heads, num_kv_heads) |
| 255 | and page_size == _PAGE_SIZE |
| 256 | and topk_blocks == _TOPK |
| 257 | ) |
| 258 | |
| 259 | |
| 260 | def should_prepare_decode_metadata( |
| 261 | batch_size: int, |
| 262 | decode_query_len: int, |
| 263 | *, |
| 264 | decode_backend: MiniMaxM3MSADecodeBackend, |
| 265 | num_q_heads: int, |
| 266 | num_kv_heads: int, |
| 267 | kv_cache_dtype: str, |
| 268 | page_size: int, |
| 269 | topk_blocks: int, |
| 270 | ) -> bool: |
| 271 | """Return whether a graph shape can use the CUTLASS decode path.""" |
| 272 | total_q = batch_size * decode_query_len |
| 273 | return ( |
| 274 | supports_cutlass_sparse_decode( |
| 275 | decode_backend=decode_backend, |
| 276 | num_q_heads=num_q_heads, |
| 277 | num_kv_heads=num_kv_heads, |
| 278 | kv_cache_dtype=kv_cache_dtype, |
| 279 | page_size=page_size, |
| 280 | topk_blocks=topk_blocks, |
| 281 | ) |
| 282 | and 1 <= decode_query_len <= _MAX_DECODE_QUERY_LEN |
| 283 | and (batch_size >= _MIN_CUTLASS_BATCH_SIZE or is_nvfp4_kv_cache(kv_cache_dtype)) |
| 284 | and total_q * num_q_heads <= _MAX_QUERY_HEAD_ROWS |
| 285 | ) |
| 286 | |
| 287 | |
| 288 | @torch.no_grad() |
| 289 | def prepare_decode_metadata( |
| 290 | block_table: torch.Tensor, |
| 291 | seq_lens: torch.Tensor, |
| 292 | seq_lens_cpu: torch.Tensor, |
| 293 | decode_query_len: int, |
| 294 | *, |
| 295 | num_q_heads: int, |
| 296 | num_kv_heads: int, |
| 297 | page_size: int, |
| 298 | topk_blocks: int, |
| 299 | plan_cache: MSACutlassDecodePlanCache | None = None, |
| 300 | ) -> MSACutlassDecodeMetadata: |
| 301 | """Prepare graph-stable runtime metadata for one sparse decode step.""" |
| 302 | cache = plan_cache or MSACutlassDecodePlanCache() |
| 303 | return cache.prepare( |
| 304 | block_table, |
| 305 | seq_lens, |
| 306 | seq_lens_cpu, |
| 307 | decode_query_len, |
| 308 | num_q_heads=num_q_heads, |
| 309 | num_kv_heads=num_kv_heads, |
| 310 | page_size=page_size, |
| 311 | topk_blocks=topk_blocks, |
| 312 | ) |
| 313 | |
| 314 | |
| 315 | @torch.no_grad() |
| 316 | def msa_cutlass_sparse_decode( |
| 317 | query_fp8: torch.Tensor, |
| 318 | kv_cache: torch.Tensor, |
| 319 | topk: torch.Tensor, |
| 320 | output: torch.Tensor, |
| 321 | metadata: MSACutlassDecodeMetadata, |
| 322 | *, |
| 323 | scale: float, |
| 324 | q_scale_float: float, |
| 325 | k_scale_float: float, |
| 326 | v_scale_float: float, |
| 327 | k_scale: torch.Tensor | None = None, |
| 328 | v_scale: torch.Tensor | None = None, |
| 329 | ) -> None: |
| 330 | """Run CUTLASS sparse decode with metadata prepared by the MSA builder. |
| 331 | |
| 332 | NVFP4 caches are uint8 head-slot pages (see ``nvfp4_kv_cache_views``) and |
| 333 | take the device global scales ``k_scale``/``v_scale``. |
| 334 | """ |
| 335 | if kv_cache.dtype == torch.uint8: |
| 336 | assert k_scale is not None and v_scale is not None |
| 337 | key, value = nvfp4_kv_cache_slots(kv_cache) |
| 338 | k_scale_arg: float | torch.Tensor = k_scale |
| 339 | v_scale_arg: float | torch.Tensor = v_scale |
| 340 | else: |
| 341 | key, value = kv_cache.split(_HEAD_DIM, dim=-1) |
| 342 | k_scale_arg, v_scale_arg = k_scale_float, v_scale_float |
| 343 | |
| 344 | from vllm.third_party.fmha_sm100.api import fmha_sm100 |
| 345 | |
| 346 | fmha_sm100( |
| 347 | query_fp8, |
| 348 | key, |
| 349 | value, |
| 350 | metadata.plan, |
| 351 | kv_indices=metadata.page_table, |
| 352 | kv_block_indexes=topk, |
| 353 | out=output, |
| 354 | output_maxscore=False, |
| 355 | output_o=True, |
| 356 | sm_scale=scale, |
| 357 | q_scale=q_scale_float, |
| 358 | k_scale=k_scale_arg, |
| 359 | v_scale=v_scale_arg, |
| 360 | o_scale=1.0, |
| 361 | ) |