| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """FlashInfer sparse MLA backend for SM90 (Hopper) NoPE models. |
| 4 | |
| 5 | Wraps FlashInfer's ``BatchMLAPagedAttentionWrapper`` (FA2/FA3 paths), which |
| 6 | as of FlashInfer 0.6.18 supports ``head_dim_kpe=0`` (GLM-5.3-Flash NoPE MLA) |
| 7 | and FP8 E4M3 KV caches on SM90 with in-kernel dequantization: the FP8 cache |
| 8 | is read directly (half the bf16 HBM traffic) and converted to BF16 in shared |
| 9 | memory, while queries stay BF16 (no query quantization). |
| 10 | |
| 11 | Sparsity rides the same trick the FA-based sparse backend uses: with |
| 12 | ``page_size=1`` the per-token top-k slot indices ARE the page table, so each |
| 13 | query token becomes one varlen batch row whose ``kv_indices`` slice is its |
| 14 | top-k row and whose ``kv_len`` is its valid count. Causality is already |
| 15 | encoded by the indexer's selection, so ``causal=False``. |
| 16 | |
| 17 | CUDA-graph handling: ``plan()`` copies its inputs to host unconditionally, |
| 18 | so it must stay outside graph capture. Each metadata builder owns a wrapper, |
| 19 | reserved capture-stable device buffers, and the plan parameters. The wrapper |
| 20 | bakes the per-row ``kv_len`` into its int schedule at plan() time — ``run()`` |
| 21 | never reads the device-side buffer. The builder plans (outside capture) from |
| 22 | sync-free host upper bounds and clamps each scheduled work item's ``kv_end`` |
| 23 | on device to the exact valid count, so the kernel never reads past a row's |
| 24 | valid prefix (the -1 tail of the converted index buffer) and async scheduling |
| 25 | needs no D2H sync. Per-step content (top-k slots) is written into the |
| 26 | reserved buffers by kernels inside the captured forward, and captured runs |
| 27 | read the refreshed plan buffers on replay. |
| 28 | |
| 29 | KV cache format: plain contiguous E4M3 ``[num_blocks, block_size, 512]`` |
| 30 | (uint8 storage) with a per-tensor ``k_scale``; BF16 caches also work. The |
| 31 | per-token x 128-channel-group ``ckv_scale_arr`` layout is supported by the |
| 32 | kernel but not wired yet (it needs a group-quantizing cache-write op). |
| 33 | """ |
| 34 | |
| 35 | from dataclasses import dataclass |
| 36 | from typing import Any, ClassVar |
| 37 | |
| 38 | import torch |
| 39 | |
| 40 | from vllm.config import VllmConfig |
| 41 | from vllm.config.cache import CacheDType |
| 42 | from vllm.model_executor.layers.attention.sparse_mla_attention import ( |
| 43 | SparseMLACommonImpl, |
| 44 | ) |
| 45 | from vllm.platforms.interface import DeviceCapability |
| 46 | from vllm.triton_utils import tl, triton |
| 47 | from vllm.utils.flashinfer import has_flashinfer_sm90_nope_mla |
| 48 | from vllm.v1.attention.backend import ( |
| 49 | AttentionBackend, |
| 50 | AttentionLayer, |
| 51 | CommonAttentionMetadata, |
| 52 | MLAAttentionImpl, |
| 53 | MultipleOf, |
| 54 | ) |
| 55 | from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( |
| 56 | FlashInferMLASparseMetadata, |
| 57 | FlashInferMLASparseMetadataBuilder, |
| 58 | ) |
| 59 | from vllm.v1.attention.backends.mla.sparse_utils import ( |
| 60 | triton_convert_req_index_to_global_index, |
| 61 | ) |
| 62 | from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout |
| 63 | |
| 64 | _FP8_KV_DTYPES = ("fp8", "fp8_e4m3") |
| 65 | _WORKSPACE_BYTES = 128 * 1024 * 1024 |
| 66 | # TODO: FlashInfer bakes per-row kv_len into the plan on the host, which |
| 67 | # forces a D2H sync under async scheduling. The plan-info layout constants |
| 68 | # below and the int-workspace clamp in _SM90State work around this by |
| 69 | # planning from sync-free upper bounds; delete them once FlashInfer accepts |
| 70 | # an upper-bound plan() that clamps kv_end internally, or exposes a |
| 71 | # GPU-side plan API (upstream ask: |
| 72 | # https://github.com/flashinfer-ai/flashinfer/issues/5559, see |
| 73 | # vllm-project/vllm#58684). |
| 74 | # FlashInfer MLAPlanInfo layout (scheduler.cuh) field indices. |
| 75 | _PLAN_INFO_LEN = 18 |
| 76 | _PI_NUM_BLKS_Y = 1 |
| 77 | _PI_Q_INDPTR = 2 |
| 78 | _PI_KV_INDPTR = 3 |
| 79 | _PI_KV_START = 13 |
| 80 | _PI_KV_END = 14 |
| 81 | _PI_WORK_INDPTR = 15 |
| 82 | # Planned lengths exceed the host upper bound by up to this many tokens so a |
| 83 | # plan stays reusable across draft steps and consecutive decode steps; the |
| 84 | # device-side clamp restores the exact lengths. |
| 85 | _PLAN_SLACK = 32 |
| 86 | |
| 87 | |
| 88 | @triton.jit |
| 89 | def _pack_topk_indices( |
| 90 | slots, |
| 91 | indptr, |
| 92 | indices, |
| 93 | WIDTH: tl.constexpr, |
| 94 | BLOCK_SIZE: tl.constexpr, |
| 95 | ): |
| 96 | row = tl.program_id(0).to(tl.int64) |
| 97 | cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 98 | start = tl.load(indptr + row) |
| 99 | end = tl.load(indptr + row + 1) |
| 100 | mask = (cols < end - start) & (cols < WIDTH) |
| 101 | values = tl.load(slots + row * WIDTH + cols, mask=mask, other=0) |
| 102 | tl.store(indices + start + cols, tl.maximum(values, 0), mask=mask) |
| 103 | |
| 104 | |
| 105 | @triton.jit |
| 106 | def _clamp_work_kv_end( |
| 107 | int_ws, |
| 108 | saved_kv_end, |
| 109 | seq_lens, |
| 110 | query_start_loc, |
| 111 | req_id_per_token, |
| 112 | num_works_idx, |
| 113 | q_indptr_off, |
| 114 | kv_start_off, |
| 115 | kv_end_off, |
| 116 | TOPK: tl.constexpr, |
| 117 | KPOOL: tl.constexpr, |
| 118 | BLOCK_SIZE: tl.constexpr, |
| 119 | ): |
| 120 | # Clamp each planned work item to its row's exact valid count, derived |
| 121 | # from the device batch layout exactly as the convert kernel counts it. |
| 122 | works = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 123 | mask = works < tl.load(int_ws + num_works_idx) |
| 124 | row = tl.load(int_ws + q_indptr_off + works, mask=mask, other=0) |
| 125 | req = tl.load(req_id_per_token + row, mask=mask, other=0) |
| 126 | seq_len = tl.load(seq_lens + req, mask=mask, other=0) |
| 127 | query_end = tl.load(query_start_loc + req + 1, mask=mask, other=0) |
| 128 | ctx = seq_len - query_end + row + 1 |
| 129 | valid = tl.where(ctx <= TOPK, ctx, TOPK + ctx % KPOOL) |
| 130 | start = tl.load(int_ws + kv_start_off + works, mask=mask, other=0) |
| 131 | end = tl.load(saved_kv_end + works, mask=mask, other=0) |
| 132 | tl.store( |
| 133 | int_ws + kv_end_off + works, |
| 134 | tl.maximum(tl.minimum(end, valid), start), |
| 135 | mask=mask, |
| 136 | ) |
| 137 | |
| 138 | |
| 139 | class FlashInferMLASparseSM90Backend(AttentionBackend): |
| 140 | supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] |
| 141 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 142 | "auto", |
| 143 | "bfloat16", |
| 144 | "fp8", |
| 145 | "fp8_e4m3", |
| 146 | ] |
| 147 | |
| 148 | @staticmethod |
| 149 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 150 | return [MultipleOf(64)] |
| 151 | |
| 152 | @staticmethod |
| 153 | def get_name() -> str: |
| 154 | return "FLASHINFER_MLA_SPARSE_SM90" |
| 155 | |
| 156 | @staticmethod |
| 157 | def get_builder_cls() -> type["FlashInferMLASparseSM90Builder"]: |
| 158 | return FlashInferMLASparseSM90Builder |
| 159 | |
| 160 | @staticmethod |
| 161 | def get_impl_cls() -> type[MLAAttentionImpl]: |
| 162 | return FlashInferMLASparseSM90Impl |
| 163 | |
| 164 | @classmethod |
| 165 | def get_supported_head_sizes(cls) -> list[int]: |
| 166 | # 512 = ckv 512 + kpe 0 (NoPE); 576 = ckv 512 + kpe 64. |
| 167 | return [512, 576] |
| 168 | |
| 169 | @classmethod |
| 170 | def is_mla(cls) -> bool: |
| 171 | return True |
| 172 | |
| 173 | @classmethod |
| 174 | def is_sparse(cls) -> bool: |
| 175 | return True |
| 176 | |
| 177 | @classmethod |
| 178 | def supports_compute_capability(cls, capability: DeviceCapability) -> bool: |
| 179 | return capability.major == 9 |
| 180 | |
| 181 | @classmethod |
| 182 | def supports_combination( |
| 183 | cls, |
| 184 | head_size: int, |
| 185 | dtype: torch.dtype, |
| 186 | kv_cache_dtype: CacheDType | None, |
| 187 | block_size: int | None, |
| 188 | use_mla: bool, |
| 189 | has_sink: bool, |
| 190 | use_sparse: bool, |
| 191 | use_mm_prefix: bool, |
| 192 | device_capability: DeviceCapability, |
| 193 | ) -> str | None: |
| 194 | if not has_flashinfer_sm90_nope_mla(): |
| 195 | return ( |
| 196 | "FLASHINFER_MLA_SPARSE_SM90 requires FlashInfer with SM90 " |
| 197 | "MLA support (ckv_scale_arr in " |
| 198 | "BatchMLAPagedAttentionWrapper.run, FlashInfer >= 0.6.18)" |
| 199 | ) |
| 200 | if not use_sparse: |
| 201 | return "FLASHINFER_MLA_SPARSE_SM90 requires sparse MLA" |
| 202 | from vllm.config import get_current_vllm_config |
| 203 | |
| 204 | vllm_config = get_current_vllm_config() |
| 205 | if vllm_config.model_config is not None: |
| 206 | hf = vllm_config.model_config.hf_text_config |
| 207 | # The SM90 FA2/FA3 kernel covers ckv=512 with kpe in {0, 64} |
| 208 | # (NoPE models and DeepSeek-style rope MLA alike). |
| 209 | if hf.kv_lora_rank != 512: |
| 210 | return "FLASHINFER_MLA_SPARSE_SM90 requires kv_lora_rank=512" |
| 211 | if hf.qk_rope_head_dim not in (0, 64): |
| 212 | return "FLASHINFER_MLA_SPARSE_SM90 requires qk_rope_head_dim in (0, 64)" |
| 213 | return None |
| 214 | |
| 215 | @staticmethod |
| 216 | def get_kv_cache_shape( |
| 217 | num_blocks: int, |
| 218 | block_size: int, |
| 219 | num_kv_heads: int, |
| 220 | head_size: int, |
| 221 | cache_dtype_str: str = "auto", |
| 222 | ) -> tuple[int, ...]: |
| 223 | return (num_blocks, block_size, head_size) |
| 224 | |
| 225 | @classmethod |
| 226 | def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: |
| 227 | return (KVCacheLayout.LBHNC,) |
| 228 | |
| 229 | |
| 230 | class _SM90State: |
| 231 | """Builder-owned wrapper, capture-stable buffers, and plan parameters. |
| 232 | |
| 233 | One instance serves every MLA layer in an attention group because the plan |
| 234 | depends only on the batch shape, not the layer. |
| 235 | """ |
| 236 | |
| 237 | def __init__( |
| 238 | self, |
| 239 | device: torch.device, |
| 240 | num_heads: int, |
| 241 | kv_dtype: torch.dtype, |
| 242 | max_tokens: int, |
| 243 | topk_width: int, |
| 244 | kv_lora_rank: int, |
| 245 | qk_rope_head_dim: int, |
| 246 | sm_scale: float, |
| 247 | index_topk: int, |
| 248 | index_kpool: int, |
| 249 | ) -> None: |
| 250 | from flashinfer.mla import BatchMLAPagedAttentionWrapper |
| 251 | |
| 252 | self.workspace = torch.empty(_WORKSPACE_BYTES, dtype=torch.uint8, device=device) |
| 253 | self.device = device |
| 254 | self.num_heads = num_heads |
| 255 | self.kv_dtype = kv_dtype |
| 256 | self.max_tokens = max_tokens |
| 257 | self.topk_width = topk_width |
| 258 | self.kv_lora_rank = kv_lora_rank |
| 259 | self.qk_rope_head_dim = qk_rope_head_dim |
| 260 | self.sm_scale = sm_scale |
| 261 | self.index_topk = index_topk |
| 262 | self.index_kpool = max(index_kpool, 1) |
| 263 | # Largest valid count the convert kernel can produce for any context. |
| 264 | self.max_valid = min(self.index_topk + self.index_kpool - 1, topk_width) |
| 265 | # User-reserved buffers: with use_cuda_graph=True plan() refreshes |
| 266 | # these in place, so run()'s captured kernels always read them. |
| 267 | self.kv_indices = torch.zeros( |
| 268 | max_tokens * topk_width, dtype=torch.int32, device=device |
| 269 | ) |
| 270 | self.kv_len_arr = torch.full( |
| 271 | (max_tokens,), topk_width, dtype=torch.int32, device=device |
| 272 | ) |
| 273 | self.kv_indptr = torch.zeros(max_tokens + 1, dtype=torch.int32, device=device) |
| 274 | self.wrapper = BatchMLAPagedAttentionWrapper( |
| 275 | self.workspace, |
| 276 | qo_indptr=torch.zeros(max_tokens + 1, dtype=torch.int32, device=device), |
| 277 | kv_indptr=self.kv_indptr, |
| 278 | kv_indices=self.kv_indices, |
| 279 | kv_len_arr=self.kv_len_arr, |
| 280 | use_cuda_graph=True, |
| 281 | backend="fa3", |
| 282 | ) |
| 283 | self._arange_cpu = torch.arange(self.max_tokens + 1, dtype=torch.int32) |
| 284 | self._qo_cpu = torch.empty(self.max_tokens + 1, dtype=torch.int32) |
| 285 | self._kv_cpu = torch.empty(self.max_tokens + 1, dtype=torch.int32) |
| 286 | self._lens_cpu = torch.full( |
| 287 | (self.max_tokens,), self.topk_width, dtype=torch.int32 |
| 288 | ) |
| 289 | # Planned per-work kv_end of the current plan, restored by each clamp; |
| 290 | # sized on first plan to the planner's work capacity. |
| 291 | self._saved_kv_end = torch.empty(0, dtype=torch.int32, device=device) |
| 292 | self._planned_num_tokens = -1 |
| 293 | self._int_ws: torch.Tensor | None = None |
| 294 | self._plan_offsets: tuple[int, int, int, int] | None = None |
| 295 | |
| 296 | def plan( |
| 297 | self, |
| 298 | num_tokens: int, |
| 299 | kv_lens: torch.Tensor, |
| 300 | cam: CommonAttentionMetadata | None, |
| 301 | req_id_per_token: torch.Tensor | None, |
| 302 | ) -> None: |
| 303 | """Plan per-row KV lengths (CPU int32, ``[num_tokens]``). |
| 304 | |
| 305 | The wrapper bakes kv_len into its int schedule from host values, so |
| 306 | rows past their valid count would send the kernel into the -1 tail of |
| 307 | the converted index buffer. With ``cam`` and ``req_id_per_token`` |
| 308 | (device batch layout), ``kv_lens`` only needs to upper-bound the |
| 309 | valid counts: the plan is padded by ``_PLAN_SLACK`` and reused while |
| 310 | the bound still fits, and every call clamps each work item's kv_end |
| 311 | on device to the exact count, so no D2H sync is needed. Without them |
| 312 | the lengths must be exact. Must run outside CUDA graph capture: the |
| 313 | in-place refreshed plan_info/indptr buffers are what captured runs |
| 314 | read. |
| 315 | """ |
| 316 | if torch.cuda.is_current_stream_capturing(): |
| 317 | raise RuntimeError( |
| 318 | "FlashInferMLASparseSM90 plan() called inside CUDA graph " |
| 319 | "capture; lengths must be planned host-side before capture." |
| 320 | ) |
| 321 | kv_lens = kv_lens.to(torch.int32) |
| 322 | if cam is None or req_id_per_token is None: |
| 323 | self._plan(num_tokens, kv_lens) |
| 324 | self._planned_num_tokens = -1 |
| 325 | return |
| 326 | planned = self._lens_cpu[:num_tokens] |
| 327 | if num_tokens != self._planned_num_tokens or not bool( |
| 328 | ((planned >= kv_lens) & (planned <= kv_lens + _PLAN_SLACK)).all() |
| 329 | ): |
| 330 | self._plan( |
| 331 | num_tokens, torch.clamp(kv_lens + _PLAN_SLACK, max=self.max_valid) |
| 332 | ) |
| 333 | self._planned_num_tokens = num_tokens |
| 334 | self._clamp_kv_end(cam, req_id_per_token) |
| 335 | |
| 336 | def _plan(self, num_tokens: int, kv_lens: torch.Tensor) -> None: |
| 337 | # CPU staging buffers are filled in place: plan() runs per step |
| 338 | # (once per draft/verify metadata build), so per-call allocations and |
| 339 | # device round trips are on the hot path. Passing CPU tensors lets |
| 340 | # the wrapper's internal .to("cpu") no-op; its reserved-buffer |
| 341 | # copy_ then performs the single H2D transfer per tensor. |
| 342 | # use_cuda_graph=True makes the wrapper copy qo/kv indptr into its |
| 343 | # fixed (max_tokens+1)-sized buffers with exact-size copy_, so the |
| 344 | # indptr must always be full-size. Rows past num_tokens are padded |
| 345 | # empty (qo_indptr flat at num_tokens) — zero-query rows read no q |
| 346 | # and schedule no work. FlashInfer 0.7 requires packed CSR offsets |
| 347 | # matching the exact lengths, including zero-length padded rows. |
| 348 | torch.clamp(self._arange_cpu, max=num_tokens, out=self._qo_cpu) |
| 349 | self._lens_cpu.zero_() |
| 350 | self._lens_cpu[:num_tokens] = kv_lens |
| 351 | self._kv_cpu[0] = 0 |
| 352 | torch.cumsum(self._lens_cpu, dim=0, out=self._kv_cpu[1:]) |
| 353 | # Only the live prefix of kv_indices: the wrapper snapshots and |
| 354 | # copies whatever is passed. |
| 355 | num_indices = max(int(self._kv_cpu[-1]), 1) |
| 356 | self.wrapper.plan( |
| 357 | self._qo_cpu, |
| 358 | self._kv_cpu, |
| 359 | self.kv_indices[:num_indices], |
| 360 | self._lens_cpu, |
| 361 | self.num_heads, |
| 362 | self.kv_lora_rank, # head_dim_ckv |
| 363 | self.qk_rope_head_dim, # 0 (NoPE) or 64 (rope MLA) |
| 364 | 1, # page_size: top-k slots are the page table |
| 365 | False, # causal: encoded by the indexer's selection |
| 366 | self.sm_scale, |
| 367 | q_data_type=torch.bfloat16, |
| 368 | kv_data_type=self.kv_dtype, |
| 369 | ) |
| 370 | plan_info = [int(x) for x in self.wrapper._plan_info] |
| 371 | assert len(plan_info) == _PLAN_INFO_LEN, plan_info |
| 372 | # Per-work arrays are laid out back to back at the planner's capacity. |
| 373 | max_works = (plan_info[_PI_KV_INDPTR] - plan_info[_PI_Q_INDPTR]) // 4 |
| 374 | if self._saved_kv_end.numel() != max_works: |
| 375 | self._saved_kv_end = torch.empty( |
| 376 | max_works, dtype=torch.int32, device=self.device |
| 377 | ) |
| 378 | self._int_ws = self.wrapper._planned_backend._int_workspace_buffer.view( |
| 379 | torch.int32 |
| 380 | ) |
| 381 | self._plan_offsets = ( |
| 382 | plan_info[_PI_WORK_INDPTR] // 4 + plan_info[_PI_NUM_BLKS_Y], |
| 383 | plan_info[_PI_Q_INDPTR] // 4, |
| 384 | plan_info[_PI_KV_START] // 4, |
| 385 | plan_info[_PI_KV_END] // 4, |
| 386 | ) |
| 387 | kv_end_off = self._plan_offsets[3] |
| 388 | self._saved_kv_end.copy_(self._int_ws[kv_end_off : kv_end_off + max_works]) |
| 389 | |
| 390 | def _clamp_kv_end( |
| 391 | self, cam: CommonAttentionMetadata, req_id_per_token: torch.Tensor |
| 392 | ) -> None: |
| 393 | assert self._int_ws is not None and self._plan_offsets is not None |
| 394 | block = 1024 |
| 395 | _clamp_work_kv_end[(triton.cdiv(self._saved_kv_end.numel(), block),)]( |
| 396 | self._int_ws, |
| 397 | self._saved_kv_end, |
| 398 | cam.seq_lens, |
| 399 | cam.query_start_loc, |
| 400 | req_id_per_token, |
| 401 | *self._plan_offsets, |
| 402 | TOPK=self.index_topk, |
| 403 | KPOOL=self.index_kpool, |
| 404 | BLOCK_SIZE=block, |
| 405 | ) |
| 406 | |
| 407 | def pack_indices(self, topk_slots: torch.Tensor) -> None: |
| 408 | num_tokens, width = topk_slots.shape |
| 409 | _pack_topk_indices[(num_tokens, triton.cdiv(width, 256))]( |
| 410 | topk_slots, self.kv_indptr, self.kv_indices, width, 256 |
| 411 | ) |
| 412 | |
| 413 | |
| 414 | @dataclass |
| 415 | class FlashInferMLASparseSM90Metadata(FlashInferMLASparseMetadata): |
| 416 | state: _SM90State | None = None |
| 417 | |
| 418 | |
| 419 | class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder): |
| 420 | """Reuse the common sparse metadata (req ids, topk buffer access).""" |
| 421 | |
| 422 | metadata_cls = FlashInferMLASparseSM90Metadata |
| 423 | |
| 424 | def __init__( |
| 425 | self, |
| 426 | kv_cache_spec: "AttentionSpec", |
| 427 | layer_names: list[str], |
| 428 | vllm_config: "VllmConfig", |
| 429 | device: torch.device, |
| 430 | ) -> None: |
| 431 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 432 | attention_layer = vllm_config.compilation_config.static_forward_context[ |
| 433 | layer_names[0] |
| 434 | ] |
| 435 | self._attention_layer = attention_layer |
| 436 | impl = attention_layer.impl |
| 437 | if not isinstance(impl, FlashInferMLASparseSM90Impl): |
| 438 | raise TypeError( |
| 439 | "FlashInferMLASparseSM90Builder requires an SM90 FlashInfer " |
| 440 | f"implementation, got {type(impl).__name__}." |
| 441 | ) |
| 442 | topk_indices_buffer = impl.topk_indices_buffer |
| 443 | assert topk_indices_buffer is not None |
| 444 | hf_config = vllm_config.model_config.hf_text_config |
| 445 | assert hf_config.index_topk is not None |
| 446 | # fp8 KV cache is stored as uint8; plan() needs the logical dtype. |
| 447 | kv_plan_dtype = ( |
| 448 | torch.float8_e4m3fn |
| 449 | if kv_cache_spec.dtype == torch.uint8 |
| 450 | else kv_cache_spec.dtype |
| 451 | ) |
| 452 | self.state = _SM90State( |
| 453 | device, |
| 454 | impl.num_heads, |
| 455 | kv_plan_dtype, |
| 456 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 457 | topk_indices_buffer.shape[1], |
| 458 | kv_lora_rank=impl.kv_lora_rank, |
| 459 | qk_rope_head_dim=impl.qk_rope_head_dim, |
| 460 | sm_scale=impl.scale, |
| 461 | index_topk=int(hf_config.index_topk), |
| 462 | index_kpool=int(getattr(hf_config, "index_kpool", 1) or 1), |
| 463 | ) |
| 464 | spec_config = vllm_config.speculative_config |
| 465 | self._adaptive_verification = bool( |
| 466 | spec_config is not None and spec_config.enable_adaptive_verification |
| 467 | ) |
| 468 | |
| 469 | def _kv_lens_host( |
| 470 | self, cam: CommonAttentionMetadata |
| 471 | ) -> tuple[int, torch.Tensor, bool]: |
| 472 | """Host per-row KV lengths for the plan and whether they are exact. |
| 473 | |
| 474 | A row for the j-th query token of request i attends |
| 475 | ``seq_lens[i] - q_len[i] + j + 1`` tokens. The indexer's selection |
| 476 | then bounds the valid count: contexts up to ``index_topk`` select |
| 477 | everything (valid == context); longer contexts keep the top |
| 478 | ``index_topk`` pool-expanded tokens plus the trailing incomplete |
| 479 | pool (valid == ``index_topk + context % index_kpool``). Both match |
| 480 | the count of non -1 entries the convert kernel produces. |
| 481 | |
| 482 | With ``seq_lens_cpu_upper_bound`` (optimistic under async spec |
| 483 | decode) returns sync-free upper bounds, which the plan clamps on |
| 484 | device. Otherwise returns exact lengths via a D2H copy. |
| 485 | """ |
| 486 | num_reqs = cam.num_reqs |
| 487 | qsl = cam.query_start_loc_cpu[: num_reqs + 1] |
| 488 | num_rows = int(qsl[-1]) |
| 489 | if num_rows == 0: |
| 490 | return 0, torch.zeros(0, dtype=torch.int32), True |
| 491 | state = self.state |
| 492 | sl_host = cam.seq_lens_cpu_upper_bound |
| 493 | if sl_host is None: |
| 494 | ctx = self._row_contexts(qsl, cam.seq_lens[:num_reqs].cpu()) |
| 495 | topk, kpool = state.index_topk, state.index_kpool |
| 496 | lens = torch.where(ctx <= topk, ctx, topk + ctx % kpool) |
| 497 | return num_rows, lens.to(torch.int32), True |
| 498 | if self._adaptive_verification: |
| 499 | # The host query split may not match the device one, so bound |
| 500 | # every row by the longest request. |
| 501 | ctx = torch.full((num_rows,), int(sl_host[:num_reqs].max())) |
| 502 | else: |
| 503 | ctx = self._row_contexts(qsl, sl_host[:num_reqs]) |
| 504 | return num_rows, torch.clamp(ctx, max=state.max_valid), False |
| 505 | |
| 506 | @staticmethod |
| 507 | def _row_contexts(qsl: torch.Tensor, seq_lens: torch.Tensor) -> torch.Tensor: |
| 508 | num_reqs = qsl.shape[0] - 1 |
| 509 | q_lens = (qsl[1:] - qsl[:-1]).to(torch.int64) |
| 510 | first_ctx = seq_lens.to(torch.int64) - q_lens + 1 |
| 511 | req_of_row = torch.repeat_interleave( |
| 512 | torch.arange(num_reqs, dtype=torch.int64), q_lens |
| 513 | ) |
| 514 | rows = torch.arange(int(qsl[-1]), dtype=torch.int64) |
| 515 | return first_ctx[req_of_row] + rows - qsl.to(torch.int64)[req_of_row] |
| 516 | |
| 517 | def build( |
| 518 | self, |
| 519 | common_prefix_len: int, |
| 520 | common_attn_metadata: CommonAttentionMetadata, |
| 521 | fast_build: bool = False, |
| 522 | ) -> FlashInferMLASparseSM90Metadata: |
| 523 | metadata = super().build(common_prefix_len, common_attn_metadata, fast_build) |
| 524 | assert isinstance(metadata, FlashInferMLASparseSM90Metadata) |
| 525 | # Plan outside any CUDA graph capture; captured runs read the |
| 526 | # refreshed buffers. |
| 527 | num_rows, kv_lens, exact = self._kv_lens_host(common_attn_metadata) |
| 528 | # MHA prefills are removed from q before forward_mqa. Plan only the |
| 529 | # decode prefix in that case, using the same routing decision as MLA. |
| 530 | if metadata.num_prefills > 0 and self._attention_layer._use_sparse_mha( |
| 531 | metadata |
| 532 | ): |
| 533 | num_rows = metadata.num_decode_tokens |
| 534 | kv_lens = kv_lens[:num_rows] |
| 535 | if exact: |
| 536 | self.state.plan(num_rows, kv_lens, None, None) |
| 537 | else: |
| 538 | self.state.plan( |
| 539 | num_rows, kv_lens, common_attn_metadata, metadata.req_id_per_token |
| 540 | ) |
| 541 | metadata.state = self.state |
| 542 | return metadata |
| 543 | |
| 544 | |
| 545 | class FlashInferMLASparseSM90Impl(SparseMLACommonImpl[FlashInferMLASparseSM90Metadata]): |
| 546 | def __init__( |
| 547 | self, |
| 548 | num_heads: int, |
| 549 | head_size: int, |
| 550 | scale: float, |
| 551 | num_kv_heads: int, |
| 552 | alibi_slopes: list[float] | None, |
| 553 | sliding_window: int | None, |
| 554 | kv_cache_dtype: str, |
| 555 | logits_soft_cap: float | None, |
| 556 | attn_type: str, |
| 557 | kv_sharing_target_layer_name: str | None, |
| 558 | topk_indices_buffer: torch.Tensor | None = None, |
| 559 | indexer: Any | None = None, |
| 560 | **mla_args: Any, |
| 561 | ) -> None: |
| 562 | if any([alibi_slopes, sliding_window, logits_soft_cap]): |
| 563 | raise NotImplementedError( |
| 564 | "FlashInferMLASparseSM90Impl does not support alibi, sliding " |
| 565 | "window, or logits soft cap." |
| 566 | ) |
| 567 | super().__init__( |
| 568 | num_heads, |
| 569 | head_size, |
| 570 | scale, |
| 571 | num_kv_heads, |
| 572 | alibi_slopes, |
| 573 | sliding_window, |
| 574 | kv_cache_dtype, |
| 575 | logits_soft_cap, |
| 576 | attn_type, |
| 577 | kv_sharing_target_layer_name, |
| 578 | indexer=indexer, |
| 579 | topk_indices_buffer=topk_indices_buffer, |
| 580 | **mla_args, |
| 581 | ) |
| 582 | assert self.topk_indices_buffer is not None |
| 583 | self.supports_quant_query_input = False |
| 584 | self.use_fp8_kv_cache = self.kv_cache_dtype in _FP8_KV_DTYPES |
| 585 | |
| 586 | def forward_mqa( |
| 587 | self, |
| 588 | q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], |
| 589 | kv_c_and_k_pe_cache: torch.Tensor, |
| 590 | attn_metadata: FlashInferMLASparseSM90Metadata, |
| 591 | layer: AttentionLayer, |
| 592 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 593 | if not isinstance(q, tuple): |
| 594 | raise NotImplementedError( |
| 595 | "FlashInferMLASparseSM90Impl expects split (q_nope, q_rope)." |
| 596 | ) |
| 597 | q_nope, q_rope = q |
| 598 | num_tokens = q_rope.shape[0] |
| 599 | # NoPE models hand a zero-width rope tensor through; rope MLA hands |
| 600 | # the real 64-dim part. The kernel takes both as-is. |
| 601 | q_pe = q_rope.reshape(num_tokens, self.num_heads, self.qk_rope_head_dim) |
| 602 | |
| 603 | assert self.topk_indices_buffer is not None |
| 604 | topk_indices = self.topk_indices_buffer[:num_tokens] |
| 605 | # return_valid_counts=True keeps the compacted-prefix layout: valid |
| 606 | # entries at [0, valid_count), -1 past it — exactly the prefix the |
| 607 | # planned per-row lengths address. |
| 608 | topk_slots, _ = triton_convert_req_index_to_global_index( |
| 609 | attn_metadata.req_id_per_token[:num_tokens], |
| 610 | attn_metadata.block_table, |
| 611 | topk_indices, |
| 612 | BLOCK_SIZE=attn_metadata.block_size, |
| 613 | NUM_TOPK_TOKENS=topk_indices.shape[1], |
| 614 | return_valid_counts=True, |
| 615 | ) |
| 616 | state = attn_metadata.state |
| 617 | assert state is not None |
| 618 | # Pack valid prefixes at the offsets refreshed by plan() before replay. |
| 619 | state.pack_indices(topk_slots) |
| 620 | |
| 621 | flat = ( |
| 622 | kv_c_and_k_pe_cache.view(torch.float8_e4m3fn) |
| 623 | if self.use_fp8_kv_cache |
| 624 | else kv_c_and_k_pe_cache |
| 625 | ).reshape(-1, 1, self.head_size) |
| 626 | ckv = flat[..., : self.kv_lora_rank] |
| 627 | kpe = flat[..., self.kv_lora_rank :] |
| 628 | |
| 629 | scale_kwargs = ( |
| 630 | {"ckv_scale": float(layer._k_scale_float or 1.0), "kpe_scale": 1.0} |
| 631 | if self.use_fp8_kv_cache |
| 632 | else {} |
| 633 | ) |
| 634 | out = state.wrapper.run(q_nope, q_pe, ckv, kpe, **scale_kwargs) |
| 635 | return out, None |