46a95815d3 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 +import functools

5 from dataclasses import dataclass

6 from typing import ClassVar, Final

7

@@ -9,6 +10,7 @@ import torch

10 from vllm._aiter_ops import rocm_aiter_ops

11 from vllm.config import VllmConfig

12 from vllm.config.cache import CacheDType

13 +from vllm.logger import init_logger

14 from vllm.model_executor.layers.attention.mla_attention import (

15 MLACommonBackend,

16 MLACommonDecodeMetadata,

@@ -26,6 +28,28 @@ from vllm.v1.attention.backend import (

28 )

29 from vllm.v1.kv_cache_interface import AttentionSpec

30

31 +logger = init_logger(__name__)

32 +

33 +

34 +@functools.lru_cache(maxsize=1)

35 +def _fp8_mla_prefill_supported() -> bool:

36 + """Auto-detect FP8 MLA prefill via mla_prefill_ps_asm_fwd + mla_reduce_v1.

37 +

38 + Requires gfx950 plus an AITER build that exports both kernels. When

39 + either is missing we silently fall back to ``flash_attn_varlen_func``.

40 + """

41 + try:

42 + from vllm.platforms.rocm import on_gfx950

43 + except Exception: # noqa: BLE001

44 + return False

45 + if not on_gfx950():

46 + return False

47 + try:

48 + from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1 # noqa: F401

49 + except Exception: # noqa: BLE001

50 + return False

51 + return True

52 +

53

54 class AiterMLABackend(MLACommonBackend):

55 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

@@ -91,6 +115,25 @@ class AiterMLAMetadata(MLACommonMetadata[AiterMLADecodeMetadata]):

115 reduce_final_map: torch.Tensor | None = None

116 reduce_partial_map: torch.Tensor | None = None

117

118 + # FP8 ASM prefill persistent-scheduling (PS) metadata. Populated by

119 + # AiterMLAMetadataBuilder._build_fp8_prefill_ps_metadata when prefill

120 + # tokens are present and FP8 MLA prefill is supported on the device.

121 + # Left as None on hosts/configs that fall back to flash_attn_varlen_func.

122 + fp8_prefill_qo_indptr: torch.Tensor | None = None

123 + fp8_prefill_kv_indptr: torch.Tensor | None = None

124 + fp8_prefill_kv_indices: torch.Tensor | None = None

125 + fp8_prefill_work_indptr: torch.Tensor | None = None

126 + fp8_prefill_work_info_set: torch.Tensor | None = None

127 + fp8_prefill_reduce_indptr: torch.Tensor | None = None

128 + fp8_prefill_reduce_final_map: torch.Tensor | None = None

129 + fp8_prefill_reduce_partial_map: torch.Tensor | None = None

130 + fp8_prefill_max_q_len: int | None = None

131 + fp8_prefill_num_partial_tiles: int | None = None

132 +

133 +

134 +# Tile size used by the mla_prefill_ps_asm_fwd assembly kernel.

135 +_FP8_PREFILL_TILE_Q = 256

136 +

137

138 class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

139 # TODO(luka, lucas): audit this as part of:

@@ -193,6 +236,14 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

236 device=device,

237 )

238

239 + self._fp8_prefill_enabled = _fp8_mla_prefill_supported()

240 + if self._fp8_prefill_enabled:

241 + max_prefill_qlen = min(

242 + vllm_config.model_config.max_model_len,

243 + vllm_config.scheduler_config.max_num_batched_tokens,

244 + )

245 + self._init_fp8_prefill_ps_buffers(max_num_reqs, max_prefill_qlen, device)

246 +

247 if self.compilation_config.cudagraph_mode.has_full_cudagraphs():

248 self.paged_kv_indptr = torch.zeros(

249 max_num_reqs + 1, dtype=torch.int32, device=device

@@ -202,6 +253,162 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

253 max_num_reqs + 1, dtype=torch.int32, device=device

254 )

255

256 + def _init_fp8_prefill_ps_buffers(

257 + self,

258 + max_num_reqs: int,

259 + max_prefill_qlen: int,

260 + device: torch.device,

261 + ) -> None:

262 + """Pre-allocate persistent buffers for FP8 MLA prefill PS metadata.

263 +

264 + Uses ``get_ps_metadata_info_v1`` with max values so the buffers are

265 + large enough for any batch. ``get_ps_metadata_v1`` fills them

266 + per-batch in ``build()``.

267 +

268 + Args:

269 + max_num_reqs: Maximum number of concurrent requests.

270 + max_prefill_qlen: Maximum Q-length for a single request in one

271 + prefill batch. Should be ``min(max_model_len,

272 + max_num_batched_tokens)`` — the chunked-prefill scheduler

273 + never emits more than ``max_num_batched_tokens`` new tokens

274 + per batch.

275 + device: Target device for the buffers.

276 + """

277 + from aiter import get_ps_metadata_info_v1

278 +

279 + # After kv_b_proj decompression, K has num_heads heads (same as Q).

280 + # So gqa_ratio=1 and num_head_k=num_heads for the PS kernel.

281 + num_head_k = self.num_heads

282 + # gqa_ratio = 1

283 + # qlen_granularity = _FP8_PREFILL_TILE_Q // max(gqa_ratio, 1)

284 + qlen_granularity = _FP8_PREFILL_TILE_Q

285 +

286 + (

287 + (work_metadata_size, work_metadata_dtype),

288 + (work_indptr_size, work_indptr_dtype),

289 + (work_info_size, work_info_dtype),

290 + (reduce_indptr_size, reduce_indptr_dtype),

291 + (reduce_final_map_size, reduce_final_map_dtype),

292 + (reduce_partial_map_size, reduce_partial_map_dtype),

293 + ) = get_ps_metadata_info_v1(

294 + batch_size=max_num_reqs,

295 + num_head_k=num_head_k,

296 + max_qlen=max_prefill_qlen,

297 + qlen_granularity=qlen_granularity,

298 + )

299 +

300 + self.fp8_ps_work_metadata = torch.empty(

301 + work_metadata_size, dtype=work_metadata_dtype, device=device

302 + )

303 + self.fp8_ps_work_indptr = torch.empty(

304 + work_indptr_size, dtype=work_indptr_dtype, device=device

305 + )

306 + self.fp8_ps_work_info = torch.empty(

307 + *work_info_size, dtype=work_info_dtype, device=device

308 + )

309 + self.fp8_ps_reduce_indptr = torch.empty(

310 + reduce_indptr_size, dtype=reduce_indptr_dtype, device=device

311 + )

312 + self.fp8_ps_reduce_final_map = torch.empty(

313 + *reduce_final_map_size, dtype=reduce_final_map_dtype, device=device

314 + )

315 + self.fp8_ps_reduce_partial_map = torch.empty(

316 + reduce_partial_map_size,

317 + dtype=reduce_partial_map_dtype,

318 + device=device,

319 + )

320 +

321 + logger.info(

322 + "FP8 MLA prefill PS buffers allocated "

323 + "(max_batch=%d, max_qlen=%d, num_head_k=%d)",

324 + max_num_reqs,

325 + max_prefill_qlen,

326 + num_head_k,

327 + )

328 +

329 + def _build_fp8_prefill_ps_metadata(

330 + self,

331 + metadata: AiterMLAMetadata,

332 + common_attn_metadata: CommonAttentionMetadata,

333 + ) -> None:

334 + """Build per-batch FP8 MLA prefill PS metadata and attach to *metadata*.

335 +

336 + Called from ``build()`` when prefill tokens are present and

337 + FP8 MLA prefill is enabled (auto-detected via

338 + ``_fp8_mla_prefill_supported()``).

339 + """

340 + from aiter import get_ps_metadata_v1

341 +

342 + prefill = metadata.prefill

343 + # Caller (build()) only invokes this when prefill tokens exist, so

344 + # metadata.prefill is guaranteed non-None. Assert to narrow for mypy.

345 + assert prefill is not None

346 + qo_indptr = prefill.query_start_loc

347 + kv_indptr = qo_indptr # new tokens: KV length == Q length

348 +

349 + # Reuse the existing CPU view of query_start_loc instead of forcing a

350 + # device->host copy. Prefill batches sit at the tail of the request

351 + # list, so we slice from num_decodes onwards and rebase to zero, the

352 + # same transform the parent build applies on device tensors.

353 + num_decodes = metadata.num_decodes

354 + qsl_cpu = common_attn_metadata.query_start_loc_cpu

355 + qo_indptr_cpu = (qsl_cpu[num_decodes:] - qsl_cpu[num_decodes]).to(torch.int32)

356 + kv_indptr_cpu = qo_indptr_cpu.clone()

357 + seq_lens_cpu = (qo_indptr_cpu[1:] - qo_indptr_cpu[:-1]).to(torch.int32)

358 +

359 + num_head_k = self.num_heads

360 + # gqa_ratio = 1

361 + # qhead_granularity = max(gqa_ratio, 1)

362 + # qlen_granularity = _FP8_PREFILL_TILE_Q // qhead_granularity

363 + gqa_ratio = 1

364 + qhead_granularity = 1

365 + qlen_granularity = _FP8_PREFILL_TILE_Q

366 + kvlen_granularity = 128

367 + block_size = 1 # non-paged: each "page" is one token

368 +

369 + get_ps_metadata_v1(

370 + qo_indptr_cpu,

371 + kv_indptr_cpu,

372 + seq_lens_cpu,

373 + gqa_ratio,

374 + num_head_k,

375 + self.fp8_ps_work_metadata,

376 + self.fp8_ps_work_indptr,

377 + self.fp8_ps_work_info,

378 + self.fp8_ps_reduce_indptr,

379 + self.fp8_ps_reduce_final_map,

380 + self.fp8_ps_reduce_partial_map,

381 + qhead_granularity=qhead_granularity,

382 + qlen_granularity=qlen_granularity,

383 + kvlen_granularity=kvlen_granularity,

384 + block_size=block_size,

385 + is_causal=True,

386 + )

387 +

388 + total_prefill_tokens = int(qo_indptr_cpu[-1].item())

389 + kv_indices = torch.arange(

390 + total_prefill_tokens, device=qo_indptr.device, dtype=torch.int32

391 + )

392 +

393 + # The actual number of active partial tiles for this batch is the

394 + # final value of reduce_indptr. Resolving it here (during metadata

395 + # build) keeps it off the per-layer forward path where a sync would

396 + # break CUDA Graph capture. Using the device-side reduce_indptr is

397 + # acceptable since build is allowed to incur an occasional sync.

398 + num_partial_tiles = int(self.fp8_ps_reduce_indptr[-1].item())

399 +

400 + # Attach PS metadata to the metadata object so forward_mha can read it.

401 + metadata.fp8_prefill_qo_indptr = qo_indptr

402 + metadata.fp8_prefill_kv_indptr = kv_indptr

403 + metadata.fp8_prefill_kv_indices = kv_indices

404 + metadata.fp8_prefill_work_indptr = self.fp8_ps_work_indptr

405 + metadata.fp8_prefill_work_info_set = self.fp8_ps_work_info

406 + metadata.fp8_prefill_reduce_indptr = self.fp8_ps_reduce_indptr

407 + metadata.fp8_prefill_reduce_final_map = self.fp8_ps_reduce_final_map

408 + metadata.fp8_prefill_reduce_partial_map = self.fp8_ps_reduce_partial_map

409 + metadata.fp8_prefill_max_q_len = prefill.max_query_len

410 + metadata.fp8_prefill_num_partial_tiles = num_partial_tiles

411 +

412 def _build_decode(

413 self,

414 block_table_tensor: torch.Tensor,

@@ -335,6 +542,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

542 attn_metadata.reduce_indptr = self._mla_reduce_indptr

543 attn_metadata.reduce_final_map = self._mla_reduce_final_map

544 attn_metadata.reduce_partial_map = self._mla_reduce_partial_map

545 + if self._fp8_prefill_enabled and attn_metadata.prefill is not None:

546 + self._build_fp8_prefill_ps_metadata(attn_metadata, common_attn_metadata)

547 return attn_metadata

548

549

@@ -479,6 +688,15 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

688

689 self.flash_attn_varlen_func = flash_attn_varlen_func

690

691 + # FP8 MLA prefill kernel imports (lazy, only when enabled).

692 + # Auto-enabled on gfx950 when AITER ships the kernels.

693 + self._fp8_prefill_enabled = _fp8_mla_prefill_supported()

694 + if self._fp8_prefill_enabled:

695 + from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1

696 +

697 + self._mla_prefill_ps_asm_fwd = mla_prefill_ps_asm_fwd

698 + self._mla_reduce_v1 = mla_reduce_v1

699 +

700 def _flash_attn_varlen_diff_headdims(

701 self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs

702 ):

@@ -493,6 +711,157 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

711

712 return output

713

714 + def _mla_fp8_prefill_attn(

715 + self,

716 + q: torch.Tensor,

717 + k: torch.Tensor,

718 + v: torch.Tensor,

719 + attn_metadata: AiterMLAMetadata,

720 + out: torch.Tensor,

721 + ) -> None:

722 + """Run FP8 MLA prefill via mla_prefill_ps_asm_fwd + mla_reduce_v1.

723 +

724 + Q, K, V are already decompressed (post-kv_b_proj), so K and V have

725 + ``num_heads`` heads (same as Q) and gqa_ratio=1. Writes the

726 + result in-place to ``out``, which is the [total_q, nhead * v_head_dim]

727 + output buffer supplied by ``forward_mha``; no extra allocation or

728 + copy is required.

729 + """

730 + from vllm.platforms import current_platform

731 + from vllm.v1.worker.workspace import current_workspace_manager

732 +

733 + fp8_dtype = current_platform.fp8_dtype()

734 + total_q = q.shape[0]

735 + nhead = self.num_heads

736 + v_head_dim = self.v_head_dim

737 + tile_q = _FP8_PREFILL_TILE_Q

738 +

739 + # The FP8 ASM kernel expects FP8 inputs; the q_scale/k_scale/v_scale

740 + # parameters select per-tensor dequant scales. Q/K/V arrive as

741 + # bf16 from kv_b_proj, so cast here (one_scale=1.0 disables scaling).

742 + if q.dtype != fp8_dtype:

743 + q = q.to(fp8_dtype)

744 + if k.dtype != fp8_dtype:

745 + k = k.to(fp8_dtype)

746 + if v.dtype != fp8_dtype:

747 + v = v.to(fp8_dtype)

748 +

749 + one_scale = torch.ones((), dtype=torch.float32, device=q.device)

750 +

751 + # num_partial_tiles is resolved during metadata build to avoid an

752 + # in-forward .item() sync that would prevent CUDA Graph capture.

753 + # forward_mha gates the FP8 path on fp8_prefill_qo_indptr being set,

754 + # and the builder always sets every fp8_prefill_* field together, so

755 + # num_partial_tiles is non-None here.

756 + num_partial_tiles = attn_metadata.fp8_prefill_num_partial_tiles

757 + assert num_partial_tiles is not None

758 +

759 + # Reuse the caller's output buffer to skip the per-call alloc + copy.

760 + # The ASM and reduce kernels both write to a [total_q, nhead, v_head_dim]

761 + # view, which aliases the [total_q, nhead * v_head_dim] storage of out.

762 + out_3d = out.view(total_q, nhead, v_head_dim)

763 +

764 + # Per-call scratch (logits, attn_lse, final_lse) is served from the

765 + # workspace manager so allocator churn in the prefill hot path is

766 + # bounded after warmup, matching the pattern in PR #41002.

767 + logits, attn_lse, final_lse = current_workspace_manager().get_simultaneous(

768 + ((num_partial_tiles * tile_q, nhead, v_head_dim), torch.float32),

769 + ((num_partial_tiles * tile_q, nhead), torch.float32),

770 + ((total_q, nhead), torch.float32),

771 + )

772 +

773 + # Phase 1: persistent-scheduling assembly prefill kernel.

774 + self._mla_prefill_ps_asm_fwd(

775 + q,

776 + k,

777 + v,

778 + attn_metadata.fp8_prefill_qo_indptr,

779 + attn_metadata.fp8_prefill_kv_indptr,

780 + attn_metadata.fp8_prefill_kv_indices,

781 + attn_metadata.fp8_prefill_work_indptr,

782 + attn_metadata.fp8_prefill_work_info_set,

783 + attn_metadata.fp8_prefill_max_q_len,

784 + self.scale,

785 + True, # is_causal

786 + logits,

787 + attn_lse,

788 + out_3d,

789 + one_scale,

790 + one_scale,

791 + one_scale,

792 + )

793 +

794 + # Phase 2: reduction across KV splits.

795 + self._mla_reduce_v1(

796 + logits,

797 + attn_lse,

798 + attn_metadata.fp8_prefill_reduce_indptr,

799 + attn_metadata.fp8_prefill_reduce_final_map,

800 + attn_metadata.fp8_prefill_reduce_partial_map,

801 + tile_q,

802 + out_3d,

803 + final_lse,

804 + )

805 +

806 + def forward_mha(

807 + self,

808 + q: torch.Tensor,

809 + kv_c_normed: torch.Tensor,

810 + k_pe: torch.Tensor,

811 + kv_c_and_k_pe_cache: torch.Tensor,

812 + attn_metadata: MLACommonMetadata,

813 + k_scale: torch.Tensor,

814 + output: torch.Tensor,

815 + ) -> None:

816 + """Dispatch prefill to the FP8 ASM kernel when available.

817 +

818 + Falls back to the parent (``flash_attn_varlen_func``) when FP8

819 + MLA prefill is disabled, PS metadata is missing, or chunked

820 + context requires two-pass merge.

821 +

822 + The annotation uses the base ``MLACommonMetadata`` to honour LSP

823 + with ``MLACommonImpl.forward_mha``; the AITER builder always

824 + produces ``AiterMLAMetadata`` instances at runtime, so we narrow

825 + with ``isinstance`` before reading the AITER-specific FP8 fields.

826 + """

827 + if (

828 + not self._fp8_prefill_enabled

829 + or not isinstance(attn_metadata, AiterMLAMetadata)

830 + or attn_metadata.fp8_prefill_qo_indptr is None

831 + ):

832 + return super().forward_mha(

833 + q,

834 + kv_c_normed,

835 + k_pe,

836 + kv_c_and_k_pe_cache,

837 + attn_metadata,

838 + k_scale,

839 + output,

840 + )

841 +

842 + assert attn_metadata.prefill is not None

843 + prefill_metadata = attn_metadata.prefill

844 + has_context = prefill_metadata.chunked_context is not None

845 +

846 + if has_context:

847 + return super().forward_mha(

848 + q,

849 + kv_c_normed,

850 + k_pe,

851 + kv_c_and_k_pe_cache,

852 + attn_metadata,

853 + k_scale,

854 + output,

855 + )

856 +

857 + kv_nope = self.kv_b_proj(kv_c_normed)[0].view(

858 + -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim

859 + )

860 + k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)

861 + k = self._concat_k_nope_k_pe(k_nope, k_pe)

862 +

863 + self._mla_fp8_prefill_attn(q, k, v, attn_metadata, output)

864 +

865 def forward_mqa(

866 self,

867 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

[ROCm][MLA] FP8 ASM prefill for AITER dense MLA backend on gfx950 (#42509)