@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 288cc6c234
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 22:15 UTC

[Attention] MLA with chunked prefill (#12639)

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Patrick Horn <patrick.horn@gmail.com> Co-authored-by: simon-mo <xmo@berkeley.edu> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>

Lucas Wilkinson committed Feb 21, 2025 at 18:30 UTC 288cc6c234d0509e2caed574729b2829ed9921a1
18 files changed +1910 -1275
csrc/cache.h
+7
@@ -39,3 +39,10 @@ void concat_and_cache_mla(torch::Tensor& kv_c, torch::Tensor& k_pe,
39 // Just for unittest
40 void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,
41 const double scale, const std::string& kv_cache_dtype);
42 +
43 +void gather_cache(
44 + torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
45 + torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
46 + torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
47 + torch::Tensor const& cu_seq_lens, // [BATCH+1]
48 + int64_t batch_size, std::optional<torch::Tensor> seq_starts = std::nullopt);
\ No newline at end of file
csrc/cache_kernels.cu
+159
@@ -2,6 +2,7 @@
2 #include <ATen/cuda/CUDAContext.h>
3 #include <c10/cuda/CUDAGuard.h>
4
5 +#include "cuda_utils.h"
6 #include "cuda_compat.h"
7 #include "dispatch_utils.h"
8
@@ -570,3 +571,161 @@ void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,
571 TORCH_CHECK(false, "Unsupported data type: ", kv_cache_dtype);
572 }
573 }
574 +
575 +namespace vllm {
576 +
577 +// grid is launched with dimensions (batch, num_splits)
578 +template <typename scalar_t>
579 +__global__ void gather_cache(
580 + const scalar_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
581 + // ENTRIES...]
582 + scalar_t* __restrict__ dst, // [TOT_TOKENS, ENTRIES...]
583 + const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES]
584 + const int32_t* __restrict__ cu_seq_lens, // [BATCH+1]
585 + const int32_t block_size, const int32_t entry_size,
586 + const int64_t block_table_stride, const int64_t cache_block_stride,
587 + const int64_t cache_entry_stride, const int64_t dst_entry_stride,
588 + const int32_t* __restrict__ seq_starts) { // Optional: starting offsets per
589 + // batch
590 +
591 + const int64_t bid = blockIdx.x; // Batch ID
592 + const int32_t num_splits = gridDim.y;
593 + const int32_t split = blockIdx.y;
594 + const int32_t seq_start = cu_seq_lens[bid];
595 + const int32_t seq_end = cu_seq_lens[bid + 1];
596 + const int32_t seq_len = seq_end - seq_start;
597 + const int32_t tot_blocks = cuda_utils::ceil_div(seq_len, block_size);
598 + const int32_t split_blocks = cuda_utils::ceil_div(tot_blocks, num_splits);
599 +
600 + const int32_t split_start = split * split_blocks;
601 + const int32_t split_end = min((split + 1) * split_blocks, tot_blocks);
602 +
603 + const bool is_active_split = (split_start < tot_blocks);
604 + const bool is_last_split = (split_end == tot_blocks);
605 +
606 + if (!is_active_split) return;
607 +
608 + int32_t full_blocks_end = split_end;
609 + int32_t partial_block_size = 0;
610 +
611 + // Adjust the pointer for the block_table for this batch.
612 + // If seq_starts is provided, compute an offset based on (seq_starts[bid] /
613 + // page_size)
614 + const int32_t batch_offset = bid * block_table_stride;
615 + int32_t offset = 0;
616 + if (seq_starts != nullptr) {
617 + offset = seq_starts[bid] / block_size;
618 + }
619 + const int32_t* batch_block_table = block_table + batch_offset + offset;
620 +
621 + // Adjust dst pointer based on the cumulative sequence lengths.
622 + dst += seq_start * dst_entry_stride;
623 +
624 + if (is_last_split) {
625 + partial_block_size = seq_len % block_size;
626 + if (partial_block_size) full_blocks_end -= 1;
627 + }
628 +
629 + auto copy_entry = [&](const scalar_t* __restrict__ _src,
630 + scalar_t* __restrict__ _dst) {
631 + for (int i = threadIdx.x; i < entry_size; i += blockDim.x)
632 + _dst[i] = _src[i];
633 + };
634 +
635 + for (int pid = split_start; pid < full_blocks_end; ++pid) {
636 + auto block_id = batch_block_table[pid];
637 + auto block_start_ptr = src_cache + block_id * cache_block_stride;
638 + auto block_dst_ptr = dst + pid * block_size * dst_entry_stride;
639 + for (int eid = 0; eid < block_size; ++eid) {
640 + copy_entry(block_start_ptr + eid * cache_entry_stride,
641 + block_dst_ptr + eid * dst_entry_stride);
642 + }
643 + }
644 +
645 + if (partial_block_size) {
646 + auto block_id = batch_block_table[full_blocks_end];
647 + auto block_start_ptr = src_cache + block_id * cache_block_stride;
648 + auto block_dst_ptr = dst + full_blocks_end * block_size * dst_entry_stride;
649 + for (int eid = 0; eid < partial_block_size; ++eid) {
650 + copy_entry(block_start_ptr + eid * cache_entry_stride,
651 + block_dst_ptr + eid * dst_entry_stride);
652 + }
653 + }
654 +}
655 +
656 +} // namespace vllm
657 +
658 +// Macro to dispatch the kernel based on the data type.
659 +#define CALL_GATHER_CACHE(CPY_DTYPE) \
660 + vllm::gather_cache<CPY_DTYPE><<<grid, block, 0, stream>>>( \
661 + reinterpret_cast<CPY_DTYPE*>(src_cache.data_ptr()), \
662 + reinterpret_cast<CPY_DTYPE*>(dst.data_ptr()), \
663 + block_table.data_ptr<int32_t>(), cu_seq_lens.data_ptr<int32_t>(), \
664 + block_size, entry_size, block_table_stride, cache_block_stride, \
665 + cache_entry_stride, dst_entry_stride, seq_starts_ptr);
666 +
667 +// Gather sequences from the cache into the destination tensor.
668 +// - cu_seq_lens contains the cumulative sequence lengths for each batch
669 +// - block_table contains the cache block indices for each sequence
670 +// - Optionally, seq_starts (if provided) offsets the starting block index by
671 +// (seq_starts[bid] / page_size)
672 +void gather_cache(
673 + torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
674 + torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
675 + torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
676 + torch::Tensor const& cu_seq_lens, // [BATCH+1]
677 + int64_t batch_size,
678 + std::optional<torch::Tensor> seq_starts = std::nullopt) {
679 + at::cuda::OptionalCUDAGuard device_guard(src_cache.device());
680 + const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
681 +
682 + int32_t block_size = src_cache.size(1);
683 + int32_t entry_size = src_cache.flatten(2, -1).size(2);
684 +
685 + TORCH_CHECK(block_table.dtype() == torch::kInt32,
686 + "block_table must be int32");
687 + TORCH_CHECK(cu_seq_lens.dtype() == torch::kInt32,
688 + "cu_seq_lens must be int32");
689 + if (seq_starts.has_value()) {
690 + TORCH_CHECK(seq_starts.value().dtype() == torch::kInt32,
691 + "seq_starts must be int32");
692 + }
693 +
694 + TORCH_CHECK(src_cache.device() == dst.device(),
695 + "src_cache and dst must be on the same device");
696 + TORCH_CHECK(src_cache.device() == block_table.device(),
697 + "src_cache and block_table must be on the same device");
698 + TORCH_CHECK(src_cache.device() == cu_seq_lens.device(),
699 + "src_cache and cu_seq_lens must be on the same device");
700 + if (seq_starts.has_value()) {
701 + TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
702 + "src_cache and seq_starts must be on the same device");
703 + }
704 +
705 + int64_t block_table_stride = block_table.stride(0);
706 + int64_t cache_block_stride = src_cache.stride(0);
707 + int64_t cache_entry_stride = src_cache.stride(1);
708 + int64_t dst_entry_stride = dst.stride(0);
709 +
710 + // Decide on the number of splits based on the batch size.
711 + int num_splits = batch_size > 128 ? 2 : batch_size > 64 ? 4 : 16;
712 + dim3 grid(batch_size, num_splits);
713 + dim3 block(1024);
714 +
715 + TORCH_CHECK(src_cache.dtype() == dst.dtype(),
716 + "src_cache and dst must have the same dtype");
717 +
718 + const int dtype_bits = src_cache.element_size() * 8;
719 + const int32_t* seq_starts_ptr =
720 + seq_starts.has_value() ? seq_starts.value().data_ptr<int32_t>() : nullptr;
721 +
722 + if (dtype_bits == 32) {
723 + CALL_GATHER_CACHE(uint32_t);
724 + } else if (dtype_bits == 16) {
725 + CALL_GATHER_CACHE(uint16_t);
726 + } else if (dtype_bits == 8) {
727 + CALL_GATHER_CACHE(uint8_t);
728 + } else {
729 + TORCH_CHECK(false, "Unsupported data type width: ", dtype_bits);
730 + }
731 +}
csrc/core/math.hpp
-5
@@ -7,8 +7,3 @@ inline constexpr uint32_t next_pow_2(uint32_t const num) {
7 if (num <= 1) return num;
8 return 1 << (CHAR_BIT * sizeof(num) - __builtin_clz(num - 1));
9 }
10 -
11 -template <typename T>
12 -inline constexpr std::enable_if_t<std::is_integral_v<T>, T> ceil_div(T a, T b) {
13 - return (a + b - 1) / b;
14 -}
\ No newline at end of file
csrc/cuda_utils.h
+18 -4
@@ -2,10 +2,14 @@
2
3 #include <stdio.h>
4
5 -#if defined(__CUDACC__) || defined(_NVHPC_CUDA)
6 - #define HOST_DEVICE_INLINE __forceinline__ __host__ __device__
7 - #define DEVICE_INLINE __forceinline__ __device__
8 - #define HOST_INLINE __forceinline__ __host__
5 +#if defined(__HIPCC__)
6 + #define HOST_DEVICE_INLINE __host__ __device__
7 + #define DEVICE_INLINE __device__
8 + #define HOST_INLINE __host__
9 +#elif defined(__CUDACC__) || defined(_NVHPC_CUDA)
10 + #define HOST_DEVICE_INLINE __host__ __device__ __forceinline__
11 + #define DEVICE_INLINE __device__ __forceinline__
12 + #define HOST_INLINE __host__ __forceinline__
13 #else
14 #define HOST_DEVICE_INLINE inline
15 #define DEVICE_INLINE inline
@@ -25,3 +29,13 @@
29 int64_t get_device_attribute(int64_t attribute, int64_t device_id);
30
31 int64_t get_max_shared_memory_per_block_device_attribute(int64_t device_id);
32 +
33 +namespace cuda_utils {
34 +
35 +template <typename T>
36 +HOST_DEVICE_INLINE constexpr std::enable_if_t<std::is_integral_v<T>, T>
37 +ceil_div(T a, T b) {
38 + return (a + b - 1) / b;
39 +}
40 +
41 +}; // namespace cuda_utils
\ No newline at end of file
csrc/quantization/cutlass_w8a8/scaled_mm_c3x.cu
+3 -2
@@ -1,7 +1,7 @@
1 #include <cudaTypedefs.h>
2 #include "c3x/scaled_mm_kernels.hpp"
3
4 -#include "core/math.hpp"
4 +#include "cuda_utils.h"
5
6 /*
7 This file defines quantized GEMM operations using the CUTLASS 3.x API, for
@@ -33,7 +33,8 @@ void cutlass_scaled_mm_sm90(torch::Tensor& c, torch::Tensor const& a,
33 auto make_group_shape = [](torch::Tensor const& x,
34 torch::Tensor const& s) -> GroupShape {
35 TORCH_CHECK(s.dim() == 2, "cutlass_scaled_mm group scales must be 2D");
36 - return {ceil_div(x.size(0), s.size(0)), ceil_div(x.size(1), s.size(1))};
36 + return {cuda_utils::ceil_div(x.size(0), s.size(0)),
37 + cuda_utils::ceil_div(x.size(1), s.size(1))};
38 };
39
40 GroupShape a_scale_group_shape = make_group_shape(a, a_scales);
csrc/torch_bindings.cpp
+6
@@ -493,6 +493,12 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
493 "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, "
494 "str kv_cache_dtype) -> ()");
495 cache_ops.impl("convert_fp8", torch::kCUDA, &convert_fp8);
496 +
497 + // Gather cache blocks from src_cache to dst.
498 + cache_ops.def(
499 + "gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, "
500 + "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()");
501 + cache_ops.impl("gather_cache", torch::kCUDA, &gather_cache);
502 }
503
504 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
tests/kernels/test_cache.py
+73 -2
@@ -682,8 +682,6 @@ def test_swap_blocks_mla(
682 torch.ops._C_cache_ops.swap_blocks,
683 (src_cache, dst_cache, block_mapping_tensor),
684 test_utils=DEFAULT_OPCHECK_TEST_UTILS,
685 - cond=(kv_lora_rank == KV_LORA_RANKS[0]
686 - and qk_rope_head_dim == QK_ROPE_HEAD_DIMS[0]),
685 )
686
687 ops.swap_blocks(src_cache, dst_cache, block_mapping_tensor)
@@ -694,3 +692,76 @@ def test_swap_blocks_mla(
692 dst_cache[dst].cpu(),
693 msg=f"Block {src} from src should have been swapped to block "
694 f"{dst} in dst_cache.")
695 +
696 +
697 +@pytest.mark.parametrize("kv_lora_rank", [512])
698 +@pytest.mark.parametrize("qk_rope_head_dim", [64])
699 +@pytest.mark.parametrize("block_size", [16])
700 +@pytest.mark.parametrize("num_blocks", [1024])
701 +@pytest.mark.parametrize("max_seq_len", [512])
702 +@pytest.mark.parametrize("batch_size", [8])
703 +@pytest.mark.parametrize("dtype", [torch.float32])
704 +@pytest.mark.parametrize("kv_cache_dtype",
705 + ["auto"]) # You can also test "fp8" if needed.
706 +@pytest.mark.parametrize("align_cache", [True, False])
707 +@pytest.mark.parametrize("device", CUDA_DEVICES)
708 +@torch.inference_mode()
709 +def test_gather_cache_mla(kv_lora_rank, qk_rope_head_dim, block_size,
710 + num_blocks, max_seq_len, batch_size, dtype,
711 + kv_cache_dtype, align_cache, device):
712 + entry_size = kv_lora_rank + qk_rope_head_dim
713 + src_cache = _create_mla_cache(num_blocks, block_size, entry_size, dtype,
714 + kv_cache_dtype, device, align_cache)
715 + _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype)
716 +
717 + seq_len_tensor = torch.randint(0,
718 + max_seq_len + 1, (batch_size, ),
719 + device=device)
720 +
721 + total_tokens = seq_len_tensor.sum()
722 + cu_seq_lens = torch.empty((batch_size + 1),
723 + dtype=torch.int32,
724 + device=device)
725 + cu_seq_lens[0] = 0
726 + cu_seq_lens[1:] = seq_len_tensor.cumsum(dim=0).to(dtype=torch.int32)
727 + print("seq_len_tensor", seq_len_tensor)
728 +
729 + tot_blocks_tensor = (seq_len_tensor + block_size - 1) // block_size
730 + block_table = torch.empty((batch_size, num_blocks),
731 + dtype=torch.int32,
732 + device=device)
733 +
734 + for b in range(batch_size):
735 + perm = torch.randperm(num_blocks, device=device)
736 + block_table[b, :] = perm
737 +
738 + dst = torch.zeros((total_tokens, entry_size),
739 + dtype=src_cache.dtype,
740 + device=device)
741 +
742 + expected_batches = []
743 + for b in range(batch_size):
744 + s = seq_len_tensor[b]
745 + if s == 0:
746 + continue
747 + tot = tot_blocks_tensor[b]
748 + blocks = block_table[b, :tot].tolist()
749 +
750 + gathered_rows = []
751 + for i in range(tot - 1):
752 + gathered_rows.append(src_cache[blocks[i]])
753 + remaining = s - (tot - 1) * block_size
754 + gathered_rows.append(src_cache[blocks[-1], :remaining, :])
755 +
756 + batch_expected = torch.cat(gathered_rows, dim=0)
757 + expected_batches.append(batch_expected)
758 + expected = torch.cat(expected_batches, dim=0)
759 +
760 + opcheck(
761 + torch.ops._C_cache_ops.gather_cache,
762 + (src_cache, dst, block_table, cu_seq_lens, batch_size, None),
763 + test_utils=DEFAULT_OPCHECK_TEST_UTILS,
764 + )
765 +
766 + ops.gather_cache(src_cache, dst, block_table, cu_seq_lens, batch_size)
767 + torch.testing.assert_close(dst, expected)
vllm/_custom_ops.py
+10
@@ -1099,6 +1099,16 @@ def convert_fp8(output: torch.Tensor,
1099 torch.ops._C_cache_ops.convert_fp8(output, input, scale, kv_dtype)
1100
1101
1102 +def gather_cache(src_cache: torch.Tensor,
1103 + dst: torch.Tensor,
1104 + block_table: torch.Tensor,
1105 + cu_seq_lens: torch.Tensor,
1106 + batch_size: int,
1107 + seq_starts: Optional[torch.Tensor] = None) -> None:
1108 + torch.ops._C_cache_ops.gather_cache(src_cache, dst, block_table,
1109 + cu_seq_lens, batch_size, seq_starts)
1110 +
1111 +
1112 def get_device_attribute(attribute: int, device: int) -> int:
1113 return torch.ops._C_cuda_utils.get_device_attribute(attribute, device)
1114
vllm/attention/__init__.py
+4 -8
@@ -4,16 +4,12 @@ from vllm.attention.backends.abstract import (AttentionBackend,
4 AttentionMetadata,
5 AttentionMetadataBuilder,
6 AttentionState, AttentionType)
7 +from vllm.attention.backends.utils import get_flash_attn_version
8 from vllm.attention.layer import Attention
9 from vllm.attention.selector import get_attn_backend
10
11 __all__ = [
11 - "Attention",
12 - "AttentionBackend",
13 - "AttentionMetadata",
14 - "AttentionType",
15 - "AttentionMetadataBuilder",
16 - "Attention",
17 - "AttentionState",
18 - "get_attn_backend",
12 + "Attention", "AttentionBackend", "AttentionMetadata", "AttentionType",
13 + "AttentionMetadataBuilder", "Attention", "AttentionState",
14 + "get_attn_backend", "get_flash_attn_version"
15 ]
vllm/attention/backends/mla/common.py new
+1503
@@ -0,0 +1,1503 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""
3 +This file implements common components for MLA implementations.
4 +
5 +First we define:
6 +
7 +Sq as Q sequence length
8 +Skv as KV sequence length
9 +
10 +MLA has two possible ways of computing, a data-movement friendly approach and a
11 +compute friendly approach, we generally want to use the compute friendly
12 +approach for "prefill" (i.e. the ratio Sq / Skv is "small", is near 1)
13 +and the data-movement friendly approach for "decode" (i.e. the ratio
14 +Sq / Skv is "large").
15 +
16 +NOTE what we deem small and large is currently determined by if its labelled
17 +prefill or decode by the scheduler, but this is something we should probably
18 +tune.
19 +
20 +Main reference: DeepseekV2 paper, and FlashInfer Implementation
21 +(https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551).
22 +
23 +Deepseek's MLA attention works the following way:
24 +* Use a single latent vector to represent the per-token entry of the KV cache.
25 +* For decode (i.e. the memory friendly approach) the attention "simulates" a
26 +multi-head attention, while the compute is similar to multi-query attention.
27 +
28 +Below is example of both paths assuming batchsize = 1
29 +
30 +## More Extent Definitions:
31 +
32 +C Context length, `Skv - Sq`
33 +H hidden size
34 +N number of attention heads
35 +Lq latent dimension for Q 1536 in DSV3
36 +Lkv latent dimension for K/V 512 in DSV3
37 +P nope dimension, no rope. 128 in DSV3
38 +R rope dimension, goes through rope. 64 in DSV3
39 +V V head dim. 128 in DSV3
40 +
41 +## Vector/Matrix Definitions
42 +
43 +h_t hidden states (input to attention) shape [Sq, H]
44 +q_c latent/compressed Q shape [Sq, Lq]
45 +q_nope uncompressed Q (no-rope) shape [Sq, N, P]
46 +q_pe uncompressed Q (rope) shape [Sq, N, R]
47 +kv_c latent/compressed KV shape [Skv, Lkv]
48 +k_pe decoupled k position embeddings shape [Skv, R]
49 +new_kv_c new kv_c from current iter shape [Sq, Lkv]
50 +new_k_pe new k_pe from current iter shape [Sq, R]
51 +cache_kv_c cached k_c from previous iters shape [C, Lkv]
52 +cache_k_pe cached k_pe from previous iters shape [C, R]
53 +W_DQ project h_t to q_c shape [H, Lq]
54 +W_UQ project q_c to q_nope shape [Lq, N * P]
55 +W_QR project q_c to q_pe shape [Lq, N * R]
56 +W_DKV project h_t to kv_c shape [H, Lkv]
57 +W_UK project kv_c to k_nope shape [Lkv, N * P]
58 +W_KR project h_t to k_pe shape [H, N * R]
59 +W_UV project kv_c to v shape [Lkv, N * V]
60 +W_O project v to h_t shape [N * V, H]
61 +
62 +
63 +## Compute Friendly Approach (i.e. "_forward_prefill"):
64 +
65 +q_c = h_t @ W_DQ
66 +q_nope = (q_c @ W_UQ).view(Sq, N, P)
67 +q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)
68 +new_kv_c = h_t @ W_DKV
69 +new_k_pe = RoPE(h_t @ W_KR)
70 +kv_c = torch.cat([new_kv_c, cache_kv_c], dim=0)
71 +k_pe = torch.cat([new_k_pe, cache_k_pe], dim=0)
72 +k_nope = (kv_c @ W_UK).view(Skv, N, P)
73 +v = (kv_c @ W_UV).view(Skv, N, V)
74 +
75 +// MHA with QK headdim = P + R
76 +// V headdim = V
77 +// spda_o shape [Sq, N, V]
78 +spda_o = scaled_dot_product_attention(
79 + torch.cat([q_nope, q_pe], dim=-1),
80 + torch.cat([k_nope, k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1),
81 + v
82 +)
83 +return spda_o @ W_O
84 +
85 +NOTE: in the actual code,
86 + `kv_b_proj` is [W_UK; W_UV] concatnated per head
87 + `q_b_proj` is [W_UQ; W_QR] concatnated per head
88 + `out_proj` is W_O
89 +
90 +
91 +## Data-Movement Friendly Approach (i.e. "_forward_decode"):
92 +
93 +Ahead of time, compute:
94 +
95 +% this projects from q_c to [Sq, N * Lkv]
96 +W_UQ_UK = einsum("qnp,knp -> qnk"
97 + W_UQ.view(Lq, N, P), W_UK.view(Lkv, N, P)
98 + ).view(Lkv, N * Lkv)
99 +% this projects from attn output [Sq, N * Lkv] to [Sq, H]
100 +W_UV_O = einsum("knv,nvh -> nkh"
101 + W_UV.view(Lkv, N, V), W_O.view(N, V, H)
102 + ).view(N * Lkv, H)
103 +
104 +Runtime
105 +q_c = h_t @ W_DQ
106 +q_latent = q_c @ W_UQ_UK.view(Sq, N, Lkv)
107 +q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)
108 +new_kv_c = h_t @ W_DKV
109 +new_k_pe = RoPE(h_t @ W_KR)
110 +kv_c = torch.cat([new_kv_c, cache_kv_c], dim=0)
111 +k_pe = torch.cat([new_k_pe, cache_k_pe], dim=0)
112 +
113 +// MQA with QK headdim = Lkv + R
114 +// V headdim = Lkv
115 +// spda_o shape [Sq, N, Lkv]
116 +// NOTE: this is less compute-friendly since Lkv > P
117 +// but is more data-movement friendly since its MQA vs MHA
118 +spda_o = scaled_dot_product_attention(
119 + torch.cat([q_latent, q_pe], dim=-1),
120 + torch.cat([kv_c, k_pe], dim=-1),
121 + kv_c
122 +)
123 +return spda_o.reshape(-1, N * Lkv) @ W_UV_O
124 +
125 +
126 +## Chunked Prefill
127 +
128 +For chunked prefill we want to use the compute friendly algorithm. We are
129 +assuming sufficiently large Sq / Skv ratio, in the future may want to switch to
130 +the data-movement friendly approach if the chunk (i.e. `Sq`) is small.
131 +
132 +However, the compute-friendly approach can potentially run out of memory if Skv
133 +is large due to: `k_nope = (kv_c @ W_UK).view(Skv, N, P)`
134 +
135 +To mitigate this, we chunk the computation of attention with respect to the
136 +current context (i.e. `cache_kv_c` and `cache_k_pe`) so that we can used a
137 +fixed workspace size.
138 +
139 +The chunked prefill approach is as follows:
140 +
141 +MCC Max chunk of context to process per iter, computed dynamically,
142 + used to bound the memory usage
143 +
144 +q_c = h_t @ W_DQ
145 +q_nope = (q_c @ W_UQ).view(Sq, N, P)
146 +q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)
147 +new_kv_c = h_t @ W_DKV
148 +new_k_pe = RoPE(h_t @ W_KR)
149 +new_k_nope = (new_kv_c @ W_UK).view(Sq, N, P)
150 +new_v = (new_kv_c @ W_UV).view(Sq, N, V)
151 +
152 +// MHA between queries and new KV
153 +// with QK headdim = P + R
154 +// V headdim = V
155 +// curr_o shape [Sq, N, V]
156 +// curr_lse shape [N, Sq], this is just order FA returns
157 +curr_o, curr_lse = scaled_dot_product_attention(
158 + torch.cat([q_nope, q_pe], dim=-1),
159 + torch.cat([new_k_nope, new_k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1),
160 + new_v,
161 + casual=True,
162 + return_softmax_lse=True
163 +)
164 +
165 +// Compute attention with the already existing context
166 +for chunk_idx in range(cdiv(C, MCC)):
167 + chunk_start = chunk_idx * MCC
168 + chunk_end = min(chunk_start + MCC, C)
169 + Sc = chunk_end - chunk_start
170 + cache_kv_c_chunk = cache_kv_c[chunk_start:chunk_end]
171 + cache_k_pe_chunk = cache_k_pe[chunk_start:chunk_end]
172 + cache_k_nope_chunk = (cache_kv_c_chunk @ W_UK).view(-1, N, P)
173 + cache_v_chunk = (cache_kv_c_chunk @ W_UV).view(-1, N, V)
174 +
175 + chunk_o, chunk_lse = scaled_dot_product_attention(
176 + torch.cat([q_nope, q_pe], dim=-1),
177 + torch.cat([cache_k_nope_chunk,
178 + cache_k_pe_chunk.unsqueeze(1).expand(-1, N, -1)],
179 + dim=-1),
180 + cache_v_chunk,
181 + casual=False,
182 + return_softmax_lse=True
183 + )
184 +
185 + curr_o, curr_lse = merge_attn_states(
186 + suffix_output=curr_o,
187 + suffix_lse=curr_lse,
188 + prefix_output=chunk_o,
189 + prefix_lse=chunk_lse,
190 + )
191 +
192 +return curr_o @ W_O
193 +"""
194 +
195 +import functools
196 +from abc import abstractmethod
197 +from collections import defaultdict
198 +from contextlib import contextmanager
199 +from dataclasses import dataclass
200 +from itertools import accumulate
201 +from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Optional, Tuple,
202 + Type, TypeVar)
203 +
204 +import torch
205 +from compressed_tensors.quantization import QuantizationStrategy
206 +
207 +from vllm import _custom_ops as ops
208 +from vllm import envs
209 +from vllm.attention.backends.abstract import (AttentionBackend, AttentionLayer,
210 + AttentionMetadata,
211 + AttentionMetadataBuilder,
212 + AttentionState, MLAAttentionImpl)
213 +from vllm.attention.backends.utils import (PAD_SLOT_ID, compute_slot_mapping,
214 + compute_slot_mapping_start_idx,
215 + get_flash_attn_version,
216 + is_block_tables_empty)
217 +from vllm.attention.ops.triton_merge_attn_states import merge_attn_states
218 +from vllm.distributed import (get_tensor_model_parallel_world_size,
219 + tensor_model_parallel_all_reduce)
220 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,
221 + LinearBase, RowParallelLinear,
222 + UnquantizedLinearMethod)
223 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
224 + CompressedTensorsLinearMethod)
225 +from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
226 + CompressedTensorsW8A8Fp8)
227 +from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
228 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
229 + apply_fp8_linear_generic, current_platform_fp8_dtype, is_fp8)
230 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
231 + scaled_quantize)
232 +from vllm.model_executor.layers.rotary_embedding import (
233 + DeepseekScalingRotaryEmbedding, RotaryEmbedding)
234 +from vllm.multimodal import MultiModalPlaceholderMap
235 +from vllm.utils import async_tensor_h2d, cdiv, make_tensor_with_pad, round_down
236 +
237 +try:
238 + from vllm.vllm_flash_attn import flash_attn_varlen_func
239 +except ImportError:
240 + # For rocm use upstream flash attention
241 + from flash_attn import flash_attn_varlen_func
242 +
243 +if TYPE_CHECKING:
244 + from vllm.worker.model_runner import (ModelInputForGPUBuilder,
245 + ModelInputForGPUWithSamplingMetadata)
246 +
247 +
248 +class MLACommonBackend(AttentionBackend):
249 +
250 + @staticmethod
251 + def get_name() -> str:
252 + return "TRITON_MLA"
253 +
254 + @staticmethod
255 + def get_metadata_cls() -> Type["AttentionMetadata"]:
256 + return MLACommonMetadata
257 +
258 + @staticmethod
259 + def get_builder_cls() -> Type["MLACommonMetadataBuilder"]:
260 + return MLACommonMetadataBuilder
261 +
262 + @staticmethod
263 + def get_state_cls() -> Type["MLACommonState"]:
264 + return MLACommonState
265 +
266 + @staticmethod
267 + def get_kv_cache_shape(
268 + num_blocks: int,
269 + block_size: int,
270 + num_kv_heads: int, # assumed to be 1 for MLA
271 + head_size: int,
272 + ) -> Tuple[int, ...]:
273 + return (num_blocks, block_size, head_size)
274 +
275 + @staticmethod
276 + def swap_blocks(
277 + src_kv_cache: torch.Tensor,
278 + dst_kv_cache: torch.Tensor,
279 + src_to_dst: torch.Tensor,
280 + ) -> None:
281 + ops.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)
282 +
283 + @staticmethod
284 + def copy_blocks(
285 + kv_caches: List[torch.Tensor],
286 + src_to_dists: torch.Tensor,
287 + ) -> None:
288 + ops.copy_blocks_mla(kv_caches, src_to_dists)
289 +
290 + @staticmethod
291 + def get_supported_head_sizes() -> List[int]:
292 + return [576]
293 +
294 +
295 +class MLACommonState(AttentionState):
296 +
297 + def __init__(self, runner):
298 + self.runner = runner
299 + self._is_graph_capturing = False
300 +
301 + scheduler_config = runner.scheduler_config
302 + self.model_config = runner.model_config
303 + cache_config = runner.cache_config
304 +
305 + self.chunked_prefill_enabled = scheduler_config.chunked_prefill_enabled
306 +
307 + if self.chunked_prefill_enabled:
308 + self.chunked_prefill_workspace_size = min(
309 + # Max sure there is enough for 8 full length request or at least
310 + # 4 pages of cache per request
311 + max(
312 + 8 * self.model_config.max_model_len, 4 *
313 + scheduler_config.max_num_seqs * cache_config.block_size),
314 + # For long-context models try not to over-allocate limiting
315 + # kv-cache space, limiting it to 64k tokens,
316 + # which would result in the workspace being:
317 + # 2*(576)*(64*1024) = 144mb
318 + # (assuming 576 MLA head dim, and fp16)
319 + # which would result in up-projected context being
320 + # 2*(192*128)*(64*1024) = 3gb
321 + # (assuming 192 QK head dim, 128 heads, and fp16)
322 + 128 * 1024)
323 + assert self.chunked_prefill_workspace_size >= \
324 + scheduler_config.max_num_seqs * cache_config.block_size
325 +
326 + @contextmanager
327 + def graph_capture(self, max_batch_size: int):
328 + self._is_graph_capturing = True
329 +
330 + self._graph_slot_mapping = torch.full((max_batch_size, ),
331 + PAD_SLOT_ID,
332 + dtype=torch.long,
333 + device=self.runner.device)
334 + self._graph_seq_lens = torch.ones(max_batch_size,
335 + dtype=torch.int32,
336 + device=self.runner.device)
337 + self._graph_block_tables = torch.from_numpy(
338 + self.runner.graph_block_tables).to(device=self.runner.device)
339 +
340 + self._positions = torch.zeros((max_batch_size, ),
341 + dtype=torch.long,
342 + device=self.runner.device)
343 +
344 + yield
345 +
346 + self._is_graph_capturing = False
347 + del self._graph_slot_mapping
348 + del self._graph_seq_lens
349 + del self._graph_block_tables
350 + del self._positions
351 +
352 + def graph_clone(self, batch_size: int):
353 + assert self._is_graph_capturing
354 + return self.__class__(self.runner)
355 +
356 + def graph_capture_get_metadata_for_batch(
357 + self, batch_size: int, is_encoder_decoder_model: bool = False):
358 + assert self._is_graph_capturing
359 +
360 + attn_metadata = self.runner.attn_backend.make_metadata(
361 + multi_modal_placeholder_index_maps=None,
362 + enable_kv_scales_calculation=False,
363 + use_cuda_graph=True,
364 + num_prefills=0,
365 + num_prefill_tokens=0,
366 + num_decode_tokens=batch_size,
367 + slot_mapping=self._graph_slot_mapping[:batch_size],
368 + seq_lens=None,
369 + seq_lens_tensor=self._graph_seq_lens[:batch_size],
370 + max_query_len=1,
371 + max_decode_query_len=1,
372 + max_prefill_seq_len=0,
373 + max_decode_seq_len=self.runner.max_seq_len_to_capture,
374 + query_start_loc=None,
375 + seq_start_loc=None,
376 + context_lens_tensor=None,
377 + block_tables=self._graph_block_tables[:batch_size],
378 + input_positions=self._positions[:batch_size],
379 + head_dim=self.runner.model_config.get_head_size())
380 +
381 + if is_encoder_decoder_model:
382 + raise NotImplementedError(
383 + "MLACommonState does not support encoder/decoder yet")
384 +
385 + return attn_metadata
386 +
387 + def get_graph_input_buffers(self,
388 + attn_metadata,
389 + is_encoder_decoder_model: bool = False):
390 + input_buffers = {
391 + "slot_mapping": attn_metadata.slot_mapping,
392 + "seq_lens_tensor": attn_metadata.decode_metadata.seq_lens_tensor,
393 + "block_tables": attn_metadata.decode_metadata.block_tables,
394 + "input_positions": attn_metadata.decode_metadata.input_positions,
395 + }
396 + if is_encoder_decoder_model:
397 + raise NotImplementedError(
398 + "MLACommonState does not support encoder/decoder yet")
399 +
400 + return input_buffers
401 +
402 + def prepare_graph_input_buffers(self,
403 + input_buffers,
404 + attn_metadata,
405 + is_encoder_decoder_model: bool = False):
406 + input_positions = attn_metadata.input_positions
407 + num_positions = input_positions.shape[0]
408 + input_buffers["seq_lens_tensor"].copy_(
409 + attn_metadata.decode_metadata.seq_lens_tensor, non_blocking=True)
410 + input_buffers["block_tables"].copy_(
411 + attn_metadata.decode_metadata.block_tables, non_blocking=True)
412 + # CUDA graph buffer is padded so only perform a partial copy based on
413 + # num_positions
414 + input_buffers["input_positions"][:num_positions].copy_(
415 + input_positions, non_blocking=True)
416 + if is_encoder_decoder_model:
417 + raise NotImplementedError(
418 + "TritonMLAState does not support encoder/decoder yet")
419 +
420 + def begin_forward(self, model_input):
421 + if self.chunked_prefill_enabled:
422 + if not hasattr(self, "chunked_prefill_workspace"):
423 + # not self.runner.device does not return the correct device
424 + # for this process, (init_device sets the correct device but
425 + # only on the Worker). The only way Ive figured out to get the
426 + # correct device is to allocate the workspace on the first call
427 + # to begin_forward and use the device of the input tokens
428 + assert model_input.input_tokens is not None
429 + self.chunked_prefill_workspace = torch.empty(
430 + (self.chunked_prefill_workspace_size,
431 + self.model_config.get_head_size()),
432 + dtype=self.model_config.dtype,
433 + device=model_input.input_tokens.device,
434 + )
435 +
436 + model_input.attn_metadata.chunked_prefill_workspace = \
437 + self.chunked_prefill_workspace
438 +
439 +
440 +@dataclass
441 +class MLACommonMetadata(AttentionMetadata):
442 + """Metadata for MLACommon.
443 +
444 + NOTE: Please read the comment at the top of the file before trying to
445 + understand this class
446 +
447 + NOTE: Any python object stored here is not updated when it is
448 + cuda-graph replayed. If you have values that need to be changed
449 + dynamically, it should be stored in tensor. The tensor has to be
450 + updated from `CUDAGraphRunner.forward` API.
451 + """
452 + # Whether or not if cuda graph is enabled.
453 + # Cuda-graph is currently enabled for decoding only.
454 + # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.
455 + use_cuda_graph: bool
456 +
457 + # New for MLA (compared to FlashAttention)
458 + # Input positions for rotrary embeddings since for MLA the rotary
459 + # position embeddings are applied inside the attention backend
460 + input_positions: torch.Tensor
461 +
462 + # NOTE(sang): Definition of context_len, query_len, and seq_len.
463 + # |---------- N-1 iteration --------|
464 + # |---------------- N iteration ---------------------|
465 + # |- tokenA -|......................|-- newTokens ---|
466 + # |---------- context_len ----------|
467 + # |-------------------- seq_len ---------------------|
468 + # |-- query_len ---|
469 +
470 + # (batch_size,). The sequence length per sequence. Sequence length means
471 + # the computed tokens + new tokens None if it is a decoding.
472 + seq_lens: Optional[List[int]]
473 + # seq_lens stored as a tensor.
474 + seq_lens_tensor: Optional[torch.Tensor]
475 +
476 + # Maximum sequence length among prefill batch. 0 if there are decoding
477 + # requests only.
478 + max_prefill_seq_len: int
479 + # Maximum sequence length among decode batch. 0 if there are prefill
480 + # requests only.
481 + max_decode_seq_len: int
482 + # (batch_size,) A tensor of context lengths (tokens that are computed
483 + # so far).
484 + context_lens_tensor: Optional[torch.Tensor]
485 +
486 + # (batch_size, max_blocks_per_seq).
487 + # Block addresses per sequence. (Seq id -> list of physical block)
488 + # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks
489 + # in the kv cache. Each block can contain up to block_size tokens.
490 + # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph
491 + # captured.
492 + block_tables: Optional[torch.Tensor]
493 +
494 + # Maximum query length in the batch.
495 + max_query_len: Optional[int] = None
496 +
497 + # Max number of query tokens among request in the batch.
498 + max_decode_query_len: Optional[int] = None
499 +
500 + # (batch_size + 1,). The cumulative subquery lengths of the sequences in
501 + # the batch, used to index into subquery. E.g., if the subquery length
502 + # is [4, 6], it is [0, 4, 10].
503 + query_start_loc: Optional[torch.Tensor] = None
504 + # (batch_size + 1,). The cumulative sequence lengths of the sequences in
505 + # the batch, used to index into sequence. E.g., if the sequence length is
506 + # [4, 6], it is [0, 4, 10].
507 + seq_start_loc: Optional[torch.Tensor] = None
508 +
509 + _cached_prefill_metadata: Optional["MLACommonMetadata"] = None
510 + _cached_decode_metadata: Optional["MLACommonMetadata"] = None
511 +
512 + num_prefill_tokens: int
513 +
514 + # The dimension of the attention heads
515 + head_dim: Optional[int] = None
516 +
517 + # Used when chunked prefill is enabled to simulate worst case workspace
518 + # allocations, hopefully to avoid going OOM
519 + is_profile_run: bool = False
520 +
521 + # New for MLA (compared to FlashAttention)
522 + # For chunked prefill
523 + context_chunk_cu_seq_lens: Optional[torch.Tensor] = None
524 + context_chunk_starts: Optional[torch.Tensor] = None
525 + context_chunk_seq_tot: Optional[List[int]] = None
526 + context_chunk_max_seq_lens: Optional[List[int]] = None
527 + # Set by MLAAttentionState in `begin_forward` so it doesn't get broadcasted
528 + chunked_prefill_workspace: Optional[torch.Tensor] = None
529 +
530 + def __post_init__(self):
531 + supported_head_sizes = MLACommonBackend.get_supported_head_sizes()
532 + if self.head_dim is not None and self.head_dim \
533 + not in supported_head_sizes:
534 + raise ValueError(
535 + f"Only {supported_head_sizes} are supported for head_dim,",
536 + f"received {self.head_dim}.")
537 +
538 + @property
539 + def prefill_metadata(self) -> Optional["MLACommonMetadata"]:
540 + if self.num_prefills == 0:
541 + return None
542 +
543 + if self._cached_prefill_metadata is not None:
544 + return self._cached_prefill_metadata
545 +
546 + assert self.seq_lens is not None
547 + assert self.seq_lens_tensor is not None
548 +
549 + # Compute some attn_metadata fields which default to None
550 + query_start_loc = (None if self.query_start_loc is None else
551 + self.query_start_loc[:self.num_prefills + 1])
552 + slot_mapping = (None if self.slot_mapping is None else
553 + self.slot_mapping[:self.num_prefill_tokens])
554 + seq_lens = (None if self.seq_lens is None else
555 + self.seq_lens[:self.num_prefills])
556 + seq_lens_tensor = (None if self.seq_lens_tensor is None else
557 + self.seq_lens_tensor[:self.num_prefills])
558 + seq_start_loc = (None if self.seq_start_loc is None else
559 + self.seq_start_loc[:self.num_prefills + 1])
560 + context_lens_tensor = (None if self.context_lens_tensor is None else
561 + self.context_lens_tensor[:self.num_prefills])
562 + block_tables = (None if self.block_tables is None else
563 + self.block_tables[:self.num_prefills])
564 + input_positions = (None if self.input_positions is None else
565 + self.input_positions[:self.num_prefill_tokens])
566 +
567 + self._cached_prefill_metadata = MLACommonMetadata(
568 + # Required by ModelRunner
569 + use_cuda_graph=False, # Not Attention Related
570 + # Required by Attention Metadata
571 + num_prefills=self.num_prefills,
572 + num_prefill_tokens=self.num_prefill_tokens,
573 + num_decode_tokens=0,
574 + slot_mapping=slot_mapping,
575 + # Required by Attention Metadata (not used)
576 + multi_modal_placeholder_index_maps=None,
577 + enable_kv_scales_calculation=False,
578 + # MLACommonMetadata
579 + input_positions=input_positions,
580 + seq_lens=seq_lens,
581 + seq_lens_tensor=seq_lens_tensor,
582 + max_query_len=self.max_query_len,
583 + max_prefill_seq_len=self.max_prefill_seq_len,
584 + max_decode_query_len=0,
585 + max_decode_seq_len=0,
586 + query_start_loc=query_start_loc,
587 + seq_start_loc=seq_start_loc,
588 + context_lens_tensor=context_lens_tensor,
589 + block_tables=block_tables,
590 + head_dim=self.head_dim,
591 + is_profile_run=self.is_profile_run,
592 + # MLACommonMetadata Chunk prefill specific
593 + context_chunk_cu_seq_lens=self.context_chunk_cu_seq_lens,
594 + context_chunk_starts=self.context_chunk_starts,
595 + context_chunk_seq_tot=self.context_chunk_seq_tot,
596 + context_chunk_max_seq_lens=self.context_chunk_max_seq_lens,
597 + )
598 + return self._cached_prefill_metadata
599 +
600 + @property
601 + def decode_metadata(self) -> Optional["MLACommonMetadata"]:
602 + if self.num_decode_tokens == 0:
603 + return None
604 +
605 + if self._cached_decode_metadata is not None:
606 + return self._cached_decode_metadata
607 + assert self.seq_lens_tensor is not None
608 +
609 + # Compute some attn_metadata fields which default to None
610 + slot_mapping = (None if self.slot_mapping is None else
611 + self.slot_mapping[self.num_prefill_tokens:])
612 + seq_lens_tensor = (None if self.seq_lens_tensor is None else
613 + self.seq_lens_tensor[self.num_prefills:])
614 + block_tables = (None if self.block_tables is None else
615 + self.block_tables[self.num_prefills:])
616 + input_positions = (None if self.input_positions is None else
617 + self.input_positions[self.num_prefill_tokens:])
618 +
619 + self._cached_decode_metadata = MLACommonMetadata(
620 + # Required by ModelRunner
621 + use_cuda_graph=self.use_cuda_graph, # Not Attention Related
622 + # Required by Attention Metadata
623 + num_prefills=0,
624 + num_prefill_tokens=0,
625 + num_decode_tokens=self.num_decode_tokens,
626 + slot_mapping=slot_mapping,
627 + # Required by Attention Metadata (not used)
628 + multi_modal_placeholder_index_maps=None,
629 + enable_kv_scales_calculation=False,
630 + # MLACommonMetadata
631 + seq_lens=None,
632 + seq_lens_tensor=seq_lens_tensor,
633 + max_decode_query_len=self.max_decode_query_len,
634 + max_query_len=self.max_query_len,
635 + max_prefill_seq_len=0,
636 + max_decode_seq_len=self.max_decode_seq_len,
637 + # Batch may be composed of prefill|decodes, adjust query start
638 + # indices to refer to the start of decodes. E.g.
639 + # in tokens:[3 prefills|6 decodes], query_start_loc=[3,9] => [0,6].
640 + query_start_loc=(self.query_start_loc[self.num_prefills:] -
641 + self.query_start_loc[self.num_prefills])
642 + if self.query_start_loc is not None else None,
643 + seq_start_loc=self.seq_start_loc[self.num_prefills:]
644 + if self.seq_start_loc is not None else None,
645 + context_lens_tensor=None,
646 + block_tables=block_tables,
647 + input_positions=input_positions,
648 + head_dim=self.head_dim,
649 + is_profile_run=self.is_profile_run)
650 + return self._cached_decode_metadata
651 +
652 + def advance_step(self,
653 + model_input: "ModelInputForGPUWithSamplingMetadata",
654 + sampled_token_ids: Optional[torch.Tensor],
655 + block_size: int,
656 + num_seqs: int,
657 + num_queries: int,
658 + turn_prefills_into_decodes: bool = False):
659 + """
660 + Update metadata in-place to advance one decode step.
661 + """
662 + # When using cudagraph, the num_seqs is padded to the next captured
663 + # batch sized, but num_queries tracks the actual number of requests in
664 + # the batch. For --enforce-eager mode, num_seqs == num_queries
665 + if num_seqs != num_queries:
666 + assert num_seqs > num_queries
667 +
668 + if turn_prefills_into_decodes:
669 + # When Mutli-Step is enabled with Chunked-Prefill, prefills and
670 + # decodes are scheduled together. In the first step, all the
671 + # prefills turn into decodes. This update reflects that
672 + # conversion.
673 + assert self.num_decode_tokens + self.num_prefills == num_seqs
674 + self.num_decode_tokens += self.num_prefills
675 + self.num_prefills = 0
676 + self.num_prefill_tokens = 0
677 + self.max_prefill_seq_len = 0
678 + self.max_query_len = 1
679 +
680 + self.slot_mapping = self.slot_mapping[:num_seqs]
681 + else:
682 + assert self.seq_lens is not None
683 + assert self.max_decode_seq_len == max(self.seq_lens)
684 +
685 + assert self.num_prefills == 0
686 + assert self.num_prefill_tokens == 0
687 + assert self.num_decode_tokens == num_seqs
688 + assert self.slot_mapping.shape == (num_seqs, )
689 +
690 + assert self.seq_lens is not None
691 + assert len(self.seq_lens) == num_seqs
692 + assert self.seq_lens_tensor is not None
693 + assert self.seq_lens_tensor.shape == (num_seqs, )
694 + assert self.max_query_len == 1
695 + assert self.max_prefill_seq_len == 0
696 +
697 + assert self.query_start_loc is not None
698 + assert self.query_start_loc.shape == (num_queries + 1, )
699 + assert self.seq_start_loc is not None
700 + assert self.seq_start_loc.shape == (num_seqs + 1, )
701 +
702 + assert self.context_lens_tensor is not None
703 + assert self.context_lens_tensor.shape == (num_queries, )
704 +
705 + assert self.block_tables is not None
706 + assert self.block_tables.shape[0] == num_seqs
707 +
708 + # Update query lengths. Note that we update only queries and not seqs,
709 + # since tensors may be padded due to captured cuda graph batch size
710 + for i in range(num_queries):
711 + self.seq_lens[i] += 1
712 + self.max_decode_seq_len = max(self.seq_lens)
713 +
714 + ops.advance_step_flashattn(num_seqs=num_seqs,
715 + num_queries=num_queries,
716 + block_size=block_size,
717 + input_tokens=model_input.input_tokens,
718 + sampled_token_ids=sampled_token_ids,
719 + input_positions=model_input.input_positions,
720 + seq_lens=self.seq_lens_tensor,
721 + slot_mapping=self.slot_mapping,
722 + block_tables=self.block_tables)
723 +
724 +
725 +T = TypeVar("T", bound=MLACommonMetadata)
726 +
727 +
728 +class MLACommonMetadataBuilder(AttentionMetadataBuilder[MLACommonMetadata]):
729 + """
730 + NOTE: Please read the comment at the top of the file before trying to
731 + understand this class
732 + """
733 +
734 + def __init__(self, input_builder: "ModelInputForGPUBuilder"):
735 + self.input_builder = input_builder
736 + self.runner = input_builder.runner
737 + self.sliding_window = input_builder.sliding_window
738 + self.block_size = input_builder.block_size
739 + self.chunked_prefill_enabled = \
740 + self.runner.scheduler_config.chunked_prefill_enabled
741 +
742 + if self.chunked_prefill_enabled:
743 + attn_state = self.input_builder.runner.attn_state
744 + self.chunked_prefill_workspace_size = \
745 + attn_state.chunked_prefill_workspace_size
746 + self.page_size = self.runner.block_size
747 +
748 + def prepare(self):
749 + self.slot_mapping: List[int] = []
750 + self.prefill_seq_lens: List[int] = []
751 + self.context_lens: List[int] = []
752 + self.block_tables: List[List[int]] = []
753 + self.curr_seq_lens: List[int] = []
754 + self.input_positions: List[int] = []
755 + self.multimodal_placeholder_maps: Dict[
756 + str,
757 + MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap)
758 + self.num_prefills = 0
759 + self.num_prefill_tokens = 0
760 + self.num_decode_tokens = 0
761 + self.has_prefix_cache_hit = False
762 +
763 + def _add_seq_group(
764 + self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup",
765 + chunked_prefill_enabled: bool, prefix_cache_hit: bool):
766 + """Add a sequence group to the metadata. Specifically update/append
767 + 1. context length.
768 + 2. block table.
769 + 3. slot mapping.
770 + """
771 + is_prompt = inter_data.is_prompt
772 + block_tables = inter_data.block_tables
773 +
774 + for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,
775 + curr_sliding_window_block, input_positions) in zip(
776 + inter_data.seq_ids, [len(t) for t in inter_data.input_tokens],
777 + inter_data.orig_seq_lens, inter_data.seq_lens,
778 + inter_data.query_lens, inter_data.context_lens,
779 + inter_data.curr_sliding_window_blocks,
780 + inter_data.input_positions):
781 + self.input_positions.extend(input_positions)
782 + self.context_lens.append(context_len)
783 + if is_prompt:
784 + self.num_prefills += 1
785 + self.num_prefill_tokens += token_len
786 + self.prefill_seq_lens.append(seq_len)
787 + else:
788 + self.num_decode_tokens += query_len
789 + self.curr_seq_lens.append(curr_seq_len)
790 +
791 + # Compute block table.
792 + # TODO(sang): Combine chunked prefill and prefix caching by
793 + # only allowing multiple of block_size chunk size.
794 + # NOTE: This only works for oooooooxxx style attention.
795 + block_table = []
796 + if prefix_cache_hit:
797 + # NOTE(woosuk): For flash-attn, the block table should
798 + # include the entries for the incoming prefill tokens.
799 + block_table = block_tables[seq_id]
800 + elif ((chunked_prefill_enabled or not is_prompt)
801 + and block_tables is not None):
802 + if curr_sliding_window_block == 0:
803 + block_table = block_tables[seq_id]
804 + else:
805 + block_table = block_tables[seq_id][
806 + -curr_sliding_window_block:]
807 + self.block_tables.append(block_table)
808 +
809 + # Compute slot mapping.
810 + is_profile_run = is_block_tables_empty(block_tables)
811 + start_idx = compute_slot_mapping_start_idx(is_prompt, query_len,
812 + context_len,
813 + self.sliding_window)
814 + compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,
815 + seq_len, context_len, start_idx,
816 + self.block_size, inter_data.block_tables)
817 +
818 + def _get_graph_runner_block_tables(
819 + self, num_seqs: int,
820 + block_tables: List[List[int]]) -> torch.Tensor:
821 + # The shape of graph_block_tables is
822 + # [max batch size, max context len // block size].
823 + max_batch_size, max_blocks = self.runner.graph_block_tables.shape
824 + assert max_batch_size >= num_seqs
825 +
826 + graph_block_tables = self.runner.graph_block_tables[:num_seqs]
827 + for i, block_table in enumerate(block_tables):
828 + if block_table:
829 + num_blocks = len(block_table)
830 + if num_blocks <= max_blocks:
831 + graph_block_tables[i, :num_blocks] = block_table
832 + else:
833 + # It may be possible to have more blocks allocated due
834 + # to lookahead slots of multi-step, however, they are
835 + # not used anyway, so can be safely ignored.
836 + graph_block_tables[
837 + i, :max_blocks] = block_table[:max_blocks]
838 +
839 + return torch.from_numpy(graph_block_tables).to(
840 + device=self.runner.device, non_blocking=True)
841 +
842 + def build(self, seq_lens: List[int], query_lens: List[int],
843 + cuda_graph_pad_size: int, batch_size: int):
844 + """Build attention metadata with on-device tensors.
845 +
846 + Args:
847 + seq_lens: The maybe padded sequence lengths of the input sequences.
848 + query_lens: The query lengths of the input sequences.
849 + cuda_graph_pad_size: The padding size for cuda graph.
850 + -1 if cuda graph is not used.
851 + batch_size: The maybe padded batch size.
852 + """
853 + prefix_cache_hit = any([
854 + inter_data.prefix_cache_hit
855 + for inter_data in self.input_builder.inter_data_list
856 + ])
857 +
858 + for inter_data in self.input_builder.inter_data_list:
859 + self._add_seq_group(inter_data,
860 + self.input_builder.chunked_prefill_enabled,
861 + prefix_cache_hit)
862 +
863 + device = self.runner.device
864 + use_captured_graph = cuda_graph_pad_size != -1
865 +
866 + max_query_len = max(query_lens)
867 + decode_query_lens = query_lens[self.num_prefills:]
868 + if len(decode_query_lens) > 0:
869 + max_decode_query_len = max(decode_query_lens)
870 + else:
871 + max_decode_query_len = 1
872 + max_prefill_seq_len = max(self.prefill_seq_lens, default=0)
873 + max_decode_seq_len = max(self.curr_seq_lens, default=0)
874 + num_decode_tokens = self.num_decode_tokens
875 + query_start_loc = list(accumulate(query_lens, initial=0))
876 + seq_start_loc = list(accumulate(seq_lens, initial=0))
877 +
878 + num_seqs = len(seq_lens)
879 + if use_captured_graph:
880 + self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size)
881 + self.block_tables.extend([] * cuda_graph_pad_size)
882 + num_decode_tokens = batch_size - self.num_prefill_tokens
883 + block_tables = self._get_graph_runner_block_tables(
884 + num_seqs, self.block_tables)
885 + else:
886 + block_tables = make_tensor_with_pad(
887 + self.block_tables,
888 + pad=0,
889 + dtype=torch.int,
890 + device=device,
891 + )
892 + assert max_query_len > 0, ("query_lens: {}".format(query_lens))
893 +
894 + assert device is not None
895 + context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int,
896 + device, self.runner.pin_memory)
897 + seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device,
898 + self.runner.pin_memory)
899 + input_positions = async_tensor_h2d(self.input_positions, torch.long,
900 + device, self.runner.pin_memory)
901 + slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.long,
902 + device, self.runner.pin_memory)
903 + query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32,
904 + device,
905 + self.runner.pin_memory)
906 + seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32,
907 + device, self.runner.pin_memory)
908 +
909 + context_chunk_cu_seq_lens = None
910 + context_chunk_starts = None
911 + context_chunk_seq_tot = None
912 + context_chunk_max_seq_lens = None
913 +
914 + if self.chunked_prefill_enabled and self.num_prefills > 0 \
915 + and context_lens_tensor is not None \
916 + and context_lens_tensor[:self.num_prefills].max() > 0:
917 +
918 + # NOTE: it is recommend you read the `Chunked Prefill` section in
919 + # the comment at the top of the file before trying to understand
920 + # the following code
921 +
922 + num_prefills_with_context = \
923 + (context_lens_tensor[:self.num_prefills] > 0).sum().item()
924 +
925 + # currently we allocate an equal amount of workspace for each
926 + # prefill in the batch, we could probably use a more advanced
927 + # algorithm here and allocate more workspace to prefills with
928 + # longer context lengths
929 + max_context_chunk = \
930 + self.chunked_prefill_workspace_size // num_prefills_with_context
931 +
932 + # align max_context_chunk to page_size by rounding down,
933 + # currently the `gather_cache` kernel cannot handle
934 + # `context_chunk_starts` that are not aligned to page_size
935 + max_context_chunk = round_down(max_context_chunk, self.page_size)
936 + assert max_context_chunk > 0
937 + num_chunks = cdiv(context_lens_tensor.max(), max_context_chunk)
938 +
939 + # if `max_context_chunk = 256`, `num_chunks = 3`, and
940 + # `num_prefills_with_context = 4`, create a tensor that looks like
941 + # [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512]]
942 + context_chunk_starts = \
943 + torch.arange(num_chunks, device=device, dtype=torch.int32)\
944 + .unsqueeze(1).expand(-1, self.num_prefills)\
945 + * max_context_chunk
946 + chunk_ends = torch.min(context_lens_tensor[:self.num_prefills]\
947 + .unsqueeze(0), context_chunk_starts + max_context_chunk)
948 + chunk_seq_lens = (chunk_ends - context_chunk_starts).clamp(min=0)
949 + _context_chunk_cu_seq_lens = chunk_seq_lens.cumsum(dim=1).to(
950 + torch.int32)
951 + zero = torch.zeros(num_chunks, dtype=torch.int32, device=device)\
952 + .unsqueeze(-1)
953 + context_chunk_cu_seq_lens = \
954 + torch.cat([zero, _context_chunk_cu_seq_lens], dim=1)
955 + context_chunk_max_seq_lens = \
956 + chunk_seq_lens.max(dim=1).values.tolist()
957 + context_chunk_seq_tot = chunk_seq_lens.sum(dim=1).tolist()
958 + assert max(context_chunk_seq_tot) <= \
959 + self.chunked_prefill_workspace_size
960 +
961 + return MLACommonMetadata(
962 + # Required by ModelRunner
963 + use_cuda_graph=use_captured_graph, # Not Attention Related
964 + # Required by Attention Metadata
965 + num_prefills=self.num_prefills,
966 + slot_mapping=slot_mapping_tensor,
967 + num_prefill_tokens=self.num_prefill_tokens,
968 + num_decode_tokens=num_decode_tokens,
969 + # Required by Attention Metadata (not used)
970 + multi_modal_placeholder_index_maps=None, # Not Attention Related
971 + enable_kv_scales_calculation=False,
972 + # MLACommonMetadata
973 + input_positions=input_positions,
974 + seq_lens=seq_lens,
975 + seq_lens_tensor=seq_lens_tensor,
976 + max_query_len=max_query_len,
977 + max_decode_query_len=max_decode_query_len,
978 + max_prefill_seq_len=max_prefill_seq_len,
979 + max_decode_seq_len=max_decode_seq_len,
980 + query_start_loc=query_start_loc_tensor,
981 + seq_start_loc=seq_start_loc_tensor,
982 + context_lens_tensor=context_lens_tensor,
983 + block_tables=block_tables,
984 + head_dim=self.runner.model_config.get_head_size(),
985 + is_profile_run=self.runner.in_profile_run,
986 + # MLACommonMetadata Chunk prefill specific
987 + context_chunk_cu_seq_lens=context_chunk_cu_seq_lens,
988 + context_chunk_starts=context_chunk_starts,
989 + context_chunk_seq_tot=context_chunk_seq_tot,
990 + context_chunk_max_seq_lens=context_chunk_max_seq_lens,
991 + )
992 +
993 +
994 +class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
995 + """
996 + NOTE: Please read the comment at the top of the file before trying to
997 + understand this class
998 + """
999 +
1000 + def __init__(
1001 + self,
1002 + num_heads: int,
1003 + head_size: int,
1004 + scale: float,
1005 + num_kv_heads: int,
1006 + alibi_slopes: Optional[List[float]],
1007 + sliding_window: Optional[int],
1008 + kv_cache_dtype: str,
1009 + blocksparse_params: Optional[Dict[str, Any]],
1010 + logits_soft_cap: Optional[float],
1011 + attn_type: str,
1012 + # MLA Specific Arguments
1013 + q_lora_rank: Optional[int],
1014 + kv_lora_rank: int,
1015 + qk_nope_head_dim: int,
1016 + qk_rope_head_dim: int,
1017 + qk_head_dim: int,
1018 + v_head_dim: int,
1019 + rotary_emb: RotaryEmbedding,
1020 + # q_proj should be q_b_proj if q_lora_rank is not None, but from an
1021 + # attention backend perspective we rely on the layer to pass in the
1022 + # correct matrix
1023 + q_proj: ColumnParallelLinear,
1024 + kv_b_proj: ColumnParallelLinear,
1025 + o_proj: RowParallelLinear,
1026 + ) -> None:
1027 + self.num_heads = num_heads
1028 + self.head_size = head_size
1029 + self.scale = float(scale)
1030 + self.num_kv_heads = num_kv_heads
1031 + self.kv_cache_dtype = kv_cache_dtype
1032 +
1033 + self.q_lora_rank = q_lora_rank
1034 + self.kv_lora_rank = kv_lora_rank
1035 + self.qk_nope_head_dim = qk_nope_head_dim
1036 + self.qk_rope_head_dim = qk_rope_head_dim
1037 + self.qk_head_dim = qk_head_dim
1038 + self.v_head_dim = v_head_dim
1039 +
1040 + self.rotary_emb = rotary_emb
1041 + self.use_yarn_rope = isinstance(rotary_emb,
1042 + DeepseekScalingRotaryEmbedding)
1043 + self.q_proj = q_proj
1044 + self.kv_b_proj = kv_b_proj
1045 + self.o_proj = o_proj
1046 + self.vllm_flash_attn_version = get_flash_attn_version()
1047 +
1048 + # Handle the differences between the flash_attn_varlen from flash_attn
1049 + # and the one from vllm_flash_attn. The former is used on RoCM and the
1050 + # latter has an additional parameter to control FA2 vs FA3
1051 + self.flash_attn_varlen_func = flash_attn_varlen_func
1052 + if self.vllm_flash_attn_version is not None:
1053 + self.flash_attn_varlen_func = \
1054 + functools.partial(flash_attn_varlen_func,
1055 + fa_version=self.vllm_flash_attn_version)
1056 +
1057 + def _v_up_proj_and_o_proj(self, x):
1058 + if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
1059 + if is_fp8(self.W_UV_O):
1060 + output_parallel = apply_fp8_linear_generic(
1061 + x.flatten(start_dim=1), self.W_UV_O, self.W_UV_O_scales,
1062 + self.reqaunt_input_group_shape,
1063 + self.reqaunt_weight_group_shape)
1064 + else:
1065 + output_parallel = torch.matmul(x.flatten(start_dim=1),
1066 + self.W_UV_O)
1067 + if self.tp_size > 1:
1068 + output = tensor_model_parallel_all_reduce(output_parallel)
1069 + else:
1070 + output = output_parallel
1071 + return output
1072 + else:
1073 + x = torch.einsum("bnl,lnv->bnv", x, self.W_UV)
1074 + return self.o_proj(x.reshape(-1,
1075 + self.num_heads * self.v_head_dim))[0]
1076 +
1077 + def _q_proj_and_k_up_proj(self, x):
1078 + if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
1079 + if is_fp8(self.W_Q_UK):
1080 + return apply_fp8_linear_generic(
1081 + x, self.W_Q_UK, self.W_Q_UK_scales,
1082 + self.reqaunt_input_group_shape,
1083 + self.reqaunt_weight_group_shape).view(
1084 + -1, self.num_heads, self.kv_lora_rank)
1085 + return torch.matmul(x, self.W_Q_UK)\
1086 + .view(-1, self.num_heads, self.kv_lora_rank)
1087 + else:
1088 + x = torch.matmul(x, self.W_Q)\
1089 + .view(-1, self.num_heads, self.qk_nope_head_dim)
1090 + return torch.einsum("bnp,lnp->bnl", x, self.W_UK)\
1091 + .view(-1, self.num_heads, self.kv_lora_rank)
1092 +
1093 + def process_weights_after_loading(self, act_dtype: torch.dtype):
1094 +
1095 + # TODO(lucas) This is very gross, we need a more wide scale refactor of
1096 + # all the FP8 code with a more standard way of
1097 + # defining schemes/group-shapes, we should also potentially force
1098 + # quant_methods to support a decompress function
1099 + #
1100 + # returns input_group_shape, weight_group_shape
1101 + def get_scale_group_shapes_for_fp8(layer: LinearBase) -> \
1102 + Tuple[Tuple[int, int], Tuple[int, int]]:
1103 + if isinstance(layer.quant_method, Fp8LinearMethod):
1104 + if layer.quant_method.block_quant:
1105 + weight_block_size = \
1106 + layer.quant_method.quant_config.weight_block_size
1107 + # per-token-group (1, X), block-quantized (X, Y)
1108 + return (1, weight_block_size[-1]), weight_block_size
1109 + else:
1110 + return (-1, -1), (-1, -1) # per-tensor, per-tensor
1111 + elif isinstance(layer.quant_method, CompressedTensorsLinearMethod)\
1112 + and isinstance(layer.scheme, CompressedTensorsW8A8Fp8):
1113 + # this is hacky but we always assume the for
1114 + # CompressedTensorsW8A8Fp8 the input is dynamic per-token
1115 + # we ignore if it is static-per-tensor since we are going to
1116 + # requantize after later anyways
1117 + strategy = layer.scheme.strategy
1118 + if strategy == QuantizationStrategy.TENSOR:
1119 + return (1, -1), (-1, -1) # per-token, per-tensor
1120 + elif strategy == QuantizationStrategy.CHANNEL:
1121 + return (1, -1), (-1, 1) # per-token, per-channel
1122 + else:
1123 + raise NotImplementedError(
1124 + f"QuantizationStrategy.{strategy} is not supported for "
1125 + "fp8 MLA, please run with VLLM_MLA_DISABLE=1")
1126 + else:
1127 + raise NotImplementedError(
1128 + "Can't determine scale group shapes for "
1129 + f"{layer.quant_method}, please run with VLLM_MLA_DISABLE=1"
1130 + )
1131 +
1132 + def get_layer_weight(layer):
1133 + if hasattr(layer, "weight"):
1134 + return layer.weight
1135 + elif hasattr(layer, "qweight"):
1136 + return layer.qweight
1137 + else:
1138 + raise AttributeError(
1139 + f"Layer '{layer}' has neither weight nor qweight")
1140 +
1141 + def get_and_maybe_dequant_weights(layer: LinearBase):
1142 + if not isinstance(layer.quant_method, UnquantizedLinearMethod):
1143 + # NOTE: This should only be used offline, since it's O(N^3)
1144 + eye = torch.eye(layer.input_size_per_partition,
1145 + dtype=act_dtype,
1146 + device=get_layer_weight(layer).device)
1147 + dequant_weights = layer.quant_method.apply(layer,
1148 + eye,
1149 + bias=None)
1150 + del eye
1151 + # standardize to (output, input)
1152 + return dequant_weights.T
1153 + return layer.weight
1154 +
1155 + weight_dtype = get_layer_weight(self.kv_b_proj).dtype
1156 + assert get_layer_weight(self.o_proj).dtype == weight_dtype
1157 + assert get_layer_weight(self.q_proj).dtype == weight_dtype
1158 +
1159 + kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T
1160 + assert kv_b_proj_weight.shape == (
1161 + self.kv_lora_rank,
1162 + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim)), (
1163 + f"{kv_b_proj_weight.shape=}, "
1164 + f"{self.kv_lora_rank=}, "
1165 + f"{self.num_heads=}, "
1166 + f"{self.qk_nope_head_dim=}, "
1167 + f"{self.v_head_dim=}")
1168 + kv_b_proj_weight = kv_b_proj_weight.view(
1169 + self.kv_lora_rank,
1170 + self.num_heads,
1171 + self.qk_nope_head_dim + self.v_head_dim,
1172 + )
1173 +
1174 + W_UK, W_UV = kv_b_proj_weight.split(
1175 + [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
1176 +
1177 + q_proj_weight = get_and_maybe_dequant_weights(self.q_proj).T\
1178 + .view(-1, self.num_heads, self.qk_head_dim)
1179 +
1180 + # can be W_Q or W_UQ depending q_lora_rank, the former if
1181 + # q_lora_rank is None, the latter otherwise. From the Attention backend
1182 + # perspective though we call these both W_Q and rely on the layer
1183 + # to pass in the correct matrix
1184 + W_Q = q_proj_weight[..., :self.qk_nope_head_dim]
1185 + self.W_QR = q_proj_weight[..., self.qk_nope_head_dim:]\
1186 + .flatten(start_dim=1).contiguous()
1187 +
1188 + # W_QR is small so for simplicity we dont bother requantizing it
1189 + self.W_QR = self.W_QR.to(act_dtype)
1190 +
1191 + if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
1192 + requantization_enabled = not envs.VLLM_MLA_DISABLE_REQUANTIZATION
1193 + if is_fp8(weight_dtype) and requantization_enabled:
1194 + # This assumes it wise to requantize using the same group shapes
1195 + # (i.e. strategy, per-tensor, per-channel, block etc.) that the
1196 + # weights were originally quantized
1197 + requant_input_group_shape, requant_weight_group_shape = \
1198 + get_scale_group_shapes_for_fp8(self.q_proj)
1199 + assert (requant_input_group_shape, requant_weight_group_shape)\
1200 + == get_scale_group_shapes_for_fp8(self.kv_b_proj)
1201 + assert (requant_input_group_shape, requant_weight_group_shape)\
1202 + == get_scale_group_shapes_for_fp8(self.o_proj)
1203 + self.reqaunt_input_group_shape = requant_input_group_shape
1204 + self.reqaunt_weight_group_shape = requant_weight_group_shape
1205 +
1206 + #
1207 + # Perform matrix-absorption following
1208 + # https://github.com/flashinfer-ai/flashinfer/pull/551
1209 + # for decode, as a result we end up with absorbed weights for decode
1210 + # and another copy of raw weights for prefill.
1211 + #
1212 + self.W_UK, self.W_UV = kv_b_proj_weight.split(
1213 + [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
1214 + # We absorb `W_UK` into `W_Q` resulting in either W_Q_UK or W_UQ_UK
1215 + # depending q_lora_rank, the former if q_lora_rank is None, the
1216 + # latter otherwise
1217 + # basically if q_lora_rank is none we are absorbing into q_proj
1218 + # instead of UQ
1219 + W_Q_UK = torch.einsum("qnd,lnd -> qnl", W_Q, W_UK)\
1220 + .flatten(start_dim=1).contiguous()
1221 +
1222 + if is_fp8(weight_dtype) and requantization_enabled:
1223 + W_Q_UK, W_Q_UK_scales = scaled_quantize(
1224 + W_Q_UK,
1225 + self.reqaunt_weight_group_shape,
1226 + quant_dtype=current_platform_fp8_dtype)
1227 + # For FP8 save the transpose so we can use
1228 + # `apply_w8a8_block_fp8_linear` directly
1229 + self.W_Q_UK = W_Q_UK.T.contiguous()
1230 + self.W_Q_UK_scales = W_Q_UK_scales.T.contiguous()
1231 + else:
1232 + self.W_Q_UK = W_Q_UK.to(act_dtype)
1233 +
1234 + W_O = get_and_maybe_dequant_weights(self.o_proj)\
1235 + .view(-1, self.num_heads, self.v_head_dim)
1236 + W_UV_O = torch.einsum("lnd,hnd -> nlh", W_UV, W_O)\
1237 + .flatten(start_dim=0, end_dim=1).contiguous()
1238 +
1239 + if is_fp8(weight_dtype) and requantization_enabled:
1240 + W_UV_O, W_UV_O_scales = scaled_quantize(
1241 + W_UV_O,
1242 + self.reqaunt_weight_group_shape,
1243 + quant_dtype=current_platform_fp8_dtype)
1244 + # For FP8 save the transpose so we can use
1245 + # `apply_w8a8_block_fp8_linear` directly
1246 + self.W_UV_O = W_UV_O.T.contiguous()
1247 + self.W_UV_O_scales = W_UV_O_scales.T.contiguous()
1248 + else:
1249 + self.W_UV_O = W_UV_O.to(act_dtype)
1250 +
1251 + self.tp_size = get_tensor_model_parallel_world_size()
1252 + else:
1253 + if is_fp8(weight_dtype):
1254 + raise NotImplementedError(
1255 + "Currently fp8 requires matrix absorption")
1256 +
1257 + self.W_UV = W_UV
1258 + self.W_UK = W_UK
1259 + self.W_Q = W_Q.flatten(start_dim=1)
1260 +
1261 + def _compute_prefill_context(
1262 + self,
1263 + q: torch.Tensor,
1264 + kv_c_and_k_pe_cache: torch.Tensor,
1265 + attn_metadata: MLACommonMetadata,
1266 + ):
1267 + prefill_metadata = attn_metadata.prefill_metadata
1268 + assert prefill_metadata is not None
1269 + assert prefill_metadata.context_chunk_seq_tot is not None
1270 + assert prefill_metadata.context_chunk_cu_seq_lens is not None
1271 + assert prefill_metadata.context_chunk_starts is not None
1272 + assert prefill_metadata.context_chunk_max_seq_lens is not None
1273 + assert prefill_metadata.context_lens_tensor is not None
1274 +
1275 + output = None
1276 + iters = len(prefill_metadata.context_chunk_seq_tot)
1277 +
1278 + # Fetch from attn_metadata directly, since it late bound by
1279 + # MLAAttentionState, grabbing it directly `attn_metadata` can avoid
1280 + # any weirdness around prefill_metadata caching
1281 + assert attn_metadata.chunked_prefill_workspace is not None
1282 + workspace = attn_metadata.chunked_prefill_workspace
1283 +
1284 + for i in range(iters):
1285 + toks = prefill_metadata.context_chunk_seq_tot[i]
1286 +
1287 + ops.gather_cache(
1288 + src_cache=kv_c_and_k_pe_cache,
1289 + dst=workspace,
1290 + block_table=prefill_metadata.block_tables,
1291 + cu_seq_lens=prefill_metadata.context_chunk_cu_seq_lens[i],
1292 + batch_size=prefill_metadata.num_prefills,
1293 + seq_starts=prefill_metadata.context_chunk_starts[i],
1294 + )
1295 +
1296 + kv_c_normed = workspace[:toks]\
1297 + [..., :self.kv_lora_rank].unsqueeze(1)
1298 + k_pe = workspace[:toks]\
1299 + [..., self.kv_lora_rank:].unsqueeze(1)
1300 +
1301 + kv_nope = self.kv_b_proj(kv_c_normed)[0].view( \
1302 + -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
1303 + k_nope, v = kv_nope\
1304 + .split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
1305 +
1306 + k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))),
1307 + dim=-1)
1308 +
1309 + # For MLA the v head dim is smaller than qk head dim so we pad
1310 + # out v with 0s to match the qk head dim
1311 + v_padded = torch.nn.functional.pad(v,
1312 + [0, q.shape[-1] - v.shape[-1]],
1313 + value=0)
1314 +
1315 + attn_output, attn_softmax_lse = self.flash_attn_varlen_func(
1316 + q=q,
1317 + k=k,
1318 + v=v_padded,
1319 + cu_seqlens_q=prefill_metadata.query_start_loc,
1320 + cu_seqlens_k=prefill_metadata.context_chunk_cu_seq_lens[i],
1321 + max_seqlen_q=prefill_metadata.max_query_len,
1322 + max_seqlen_k=prefill_metadata.context_chunk_max_seq_lens[i],
1323 + softmax_scale=self.scale,
1324 + causal=False, # Context is unmasked
1325 + return_softmax_lse=True,
1326 + )
1327 +
1328 + if output is None:
1329 + output = attn_output
1330 + output_lse = attn_softmax_lse
1331 + else:
1332 + output_tmp = torch.empty_like(output)
1333 + output_lse_tmp = torch.empty_like(output_lse)
1334 + merge_attn_states(
1335 + output=output_tmp,
1336 + output_lse=output_lse_tmp,
1337 + prefix_output=output,
1338 + prefix_lse=output_lse,
1339 + suffix_output=attn_output,
1340 + suffix_lse=attn_softmax_lse,
1341 + )
1342 + output = output_tmp
1343 + output_lse = output_lse_tmp
1344 +
1345 + return output, output_lse
1346 +
1347 + def _forward_prefill(
1348 + self,
1349 + q: torch.Tensor,
1350 + kv_c_normed: torch.Tensor,
1351 + k_pe: torch.Tensor,
1352 + kv_c_and_k_pe_cache: torch.Tensor,
1353 + attn_metadata: MLACommonMetadata,
1354 + ) -> torch.Tensor:
1355 +
1356 + prefill_metadata = attn_metadata.prefill_metadata
1357 + assert prefill_metadata is not None
1358 +
1359 + has_context = prefill_metadata.context_lens_tensor is not None \
1360 + and prefill_metadata.context_lens_tensor.max() > 0
1361 +
1362 + kv_nope = self.kv_b_proj(kv_c_normed)[0].view(\
1363 + -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
1364 + k_nope, v = kv_nope\
1365 + .split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
1366 +
1367 + k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
1368 +
1369 + # For MLA the v head dim is smaller than qk head dim so we pad out
1370 + # v with 0s to match the qk head dim
1371 + v_padded = torch.nn.functional.pad(v, [0, q.shape[-1] - v.shape[-1]],
1372 + value=0)
1373 +
1374 + output = self.flash_attn_varlen_func(
1375 + q=q,
1376 + k=k,
1377 + v=v_padded,
1378 + cu_seqlens_q=prefill_metadata.query_start_loc,
1379 + cu_seqlens_k=prefill_metadata.query_start_loc,
1380 + max_seqlen_q=prefill_metadata.max_prefill_seq_len,
1381 + max_seqlen_k=prefill_metadata.max_prefill_seq_len,
1382 + softmax_scale=self.scale,
1383 + causal=True,
1384 + return_softmax_lse=has_context,
1385 + )
1386 +
1387 + if has_context:
1388 + suffix_output, suffix_lse = output
1389 + context_output, context_lse = self._compute_prefill_context( \
1390 + q, kv_c_and_k_pe_cache, attn_metadata)
1391 +
1392 + output = torch.empty_like(suffix_output)
1393 + merge_attn_states(
1394 + output=output,
1395 + prefix_output=context_output,
1396 + prefix_lse=context_lse,
1397 + suffix_output=suffix_output,
1398 + suffix_lse=suffix_lse,
1399 + )
1400 +
1401 + # slice by `:v.shape[-1]` in order to remove v headdim padding
1402 + output = output\
1403 + .view(-1, self.num_heads, q.shape[-1])[..., :v.shape[-1]]\
1404 + .reshape(-1, self.num_heads * v.shape[-1])
1405 +
1406 + return self.o_proj(output)[0]
1407 +
1408 + @abstractmethod
1409 + def _forward_decode(
1410 + self,
1411 + q_nope: torch.Tensor,
1412 + q_pe: torch.Tensor,
1413 + kv_c_and_k_pe_cache: torch.Tensor,
1414 + attn_metadata: T,
1415 + ) -> torch.Tensor:
1416 + raise NotImplementedError
1417 +
1418 + def forward(
1419 + self,
1420 + layer: AttentionLayer,
1421 + hidden_states_or_q_c: torch.Tensor, # query in unified attn
1422 + k_c_normed: torch.Tensor, # key in unified attn
1423 + k_pe: torch.Tensor, # value in unified attn
1424 + kv_cache: torch.Tensor,
1425 + attn_metadata: T,
1426 + output: Optional[torch.Tensor] = None,
1427 + ) -> torch.Tensor:
1428 + if output is not None:
1429 + raise NotImplementedError(
1430 + "output is not yet supported for MLAImplBase")
1431 +
1432 + if attn_metadata.is_profile_run and \
1433 + attn_metadata.chunked_prefill_workspace is not None:
1434 + # During the profile run try to simulate to worse case output size
1435 + # for `self.kv_b_proj(kv_c_normed)` in `_compute_prefill_context`
1436 + # since this can be large
1437 + _ = torch.empty(
1438 + (attn_metadata.chunked_prefill_workspace.shape[0],
1439 + self.num_heads, self.qk_nope_head_dim + self.v_head_dim),
1440 + device=k_c_normed.device,
1441 + dtype=k_c_normed.dtype,
1442 + )
1443 +
1444 + has_decode = attn_metadata.decode_metadata is not None
1445 + has_prefill = attn_metadata.prefill_metadata is not None
1446 +
1447 + # Restore head dim (for rotary embedding)
1448 + k_pe = k_pe.unsqueeze(1)
1449 + assert hasattr(attn_metadata, "input_positions")
1450 +
1451 + num_prefill_tokens: int = attn_metadata.num_prefill_tokens
1452 +
1453 + decode_hs_or_q_c = hidden_states_or_q_c[num_prefill_tokens:]
1454 + decode_k_pe = k_pe[num_prefill_tokens:]
1455 + decode_input_positions = \
1456 + attn_metadata.input_positions[num_prefill_tokens:]
1457 +
1458 + prefill_hs_or_q_c = hidden_states_or_q_c[:num_prefill_tokens]
1459 + prefill_k_pe = k_pe[:num_prefill_tokens]
1460 + prefill_input_positions = \
1461 + attn_metadata.input_positions[:num_prefill_tokens]
1462 + prefill_k_c_normed = k_c_normed[:num_prefill_tokens]
1463 +
1464 + if has_decode:
1465 + decode_q_nope = self._q_proj_and_k_up_proj(decode_hs_or_q_c)
1466 + decode_q_pe = torch.matmul(decode_hs_or_q_c, self.W_QR)\
1467 + .view(-1, self.num_heads, self.qk_rope_head_dim)
1468 + decode_q_pe[...], decode_k_pe[...] = self.rotary_emb(
1469 + decode_input_positions, decode_q_pe, decode_k_pe)
1470 +
1471 + if has_prefill:
1472 + prefill_q = self.q_proj(prefill_hs_or_q_c)[0]\
1473 + .view(-1, self.num_heads, self.qk_head_dim)
1474 + prefill_q_pe = prefill_q[..., self.qk_nope_head_dim:]
1475 + prefill_q_pe[...], prefill_k_pe[...] = self.rotary_emb(
1476 + prefill_input_positions, prefill_q_pe, prefill_k_pe)
1477 +
1478 + # write the latent and rope to kv cache
1479 + if kv_cache.numel() > 0:
1480 + ops.concat_and_cache_mla(
1481 + k_c_normed,
1482 + k_pe.squeeze(1),
1483 + kv_cache,
1484 + attn_metadata.slot_mapping.flatten(),
1485 + kv_cache_dtype=self.kv_cache_dtype,
1486 + scale=layer._k_scale,
1487 + )
1488 +
1489 + output = torch.empty(attn_metadata.num_prefill_tokens +
1490 + attn_metadata.num_decode_tokens,
1491 + self.o_proj.output_size,
1492 + device=hidden_states_or_q_c.device,
1493 + dtype=hidden_states_or_q_c.dtype)
1494 + if has_prefill:
1495 + output[:num_prefill_tokens] = self._forward_prefill(
1496 + prefill_q, prefill_k_c_normed, prefill_k_pe, kv_cache,
1497 + attn_metadata)
1498 +
1499 + if has_decode:
1500 + output[num_prefill_tokens:] = self._forward_decode(
1501 + decode_q_nope, decode_q_pe, kv_cache, attn_metadata)
1502 +
1503 + return output
vllm/attention/backends/mla/utils.py deleted
-515
@@ -1,515 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -
3 -import functools
4 -from abc import abstractmethod
5 -from dataclasses import dataclass
6 -from typing import Any, Dict, Generic, List, Optional, Tuple
7 -
8 -import torch
9 -from compressed_tensors.quantization import QuantizationStrategy
10 -
11 -from vllm import _custom_ops as ops
12 -from vllm import envs
13 -from vllm.attention.backends.abstract import (AttentionLayer,
14 - AttentionMetadata,
15 - MLAAttentionImpl, T)
16 -from vllm.attention.backends.utils import get_flash_attn_version
17 -from vllm.distributed import (get_tensor_model_parallel_world_size,
18 - tensor_model_parallel_all_reduce)
19 -from vllm.model_executor.layers.linear import (ColumnParallelLinear,
20 - LinearBase, RowParallelLinear,
21 - UnquantizedLinearMethod)
22 -from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
23 - CompressedTensorsLinearMethod)
24 -from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
25 - CompressedTensorsW8A8Fp8)
26 -from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
27 -from vllm.model_executor.layers.quantization.utils.fp8_utils import (
28 - apply_fp8_linear_generic, current_platform_fp8_dtype, is_fp8)
29 -from vllm.model_executor.layers.quantization.utils.quant_utils import (
30 - scaled_quantize)
31 -from vllm.model_executor.layers.rotary_embedding import (
32 - DeepseekScalingRotaryEmbedding, RotaryEmbedding)
33 -
34 -try:
35 - from vllm.vllm_flash_attn import flash_attn_varlen_func
36 -except ImportError:
37 - from flash_attn import flash_attn_varlen_func
38 -
39 -
40 -@dataclass
41 -class MLACommonMetadata(AttentionMetadata):
42 - # Input positions for rotrary embeddings since for MLA the rotary
43 - # position embeddings are applied inside the attention backend
44 - input_positions: torch.Tensor
45 -
46 -
47 -class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
48 - """
49 - Common class for implementing repeated parts
50 -
51 - Main reference: DeepseekV2 paper, and FlashInfer Implementation
52 - (https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551).
53 -
54 - Deepseek's MLA attention works the following way:
55 - * Use a single latent vector to represent the entire KV cache.
56 - * The attention "simulates" a multi-head attention, while the compute is
57 - similar to multi-query attention.
58 - * The dataflow is as follows,
59 -
60 - * B: batch/sequence length
61 - * H: hidden size
62 - * N: number of attention heads
63 - * Lq: latent dimension for Q
64 - * Lkv: latent dimension for K/V
65 - * P: nope dimension, P+R is the actual head_dim in common attention.
66 - * R: rope dimension, this slide of the head_dim goes through rope.
67 - * V: V head dim.
68 - * kv_c: latent/compressed KV
69 - * q_c: latent/compressed Q
70 -
71 - #
72 - # Outside the MLA attention backend
73 - #
74 -
75 - 1. The hidden states (B, H) are projected down into cq (B, Lq) and
76 - kv_c_k_pe (B, Lkv+R).
77 - 2. The kv_c_k_pe is split into kv_c (B, Lkv) and k_pe (B, R). cq
78 - and kv_c are normalized.
79 -
80 - #
81 - # Inside the MLA attention backend
82 - #
83 -
84 - * if prefill:
85 -
86 - 3. The q_c is then projected up into the multi-head version.
87 - * q_c goes from (B, Lq) to (B, N, (P+R)), which is split into q_nope
88 - (B, N, P) and q_pe (B, N, R).
89 - 4. q_pe, k_pe are then passed through rotary embeddings.
90 - 5. kv_c and k_pe are concatenated and inserted into the cache
91 - 6. The kv_c is then projected up into the multi-head version.
92 - * kv_c goes from (B, Lkv) to (B, N, (P+V)) which has the nope
93 - dimensions for K and V, which is split into k_nope (B, N, P)
94 - and v (B, N, V).
95 - 7. q (B, N, (P+R)) and k (B, N, (P+R)) matrices are assembled from
96 - q_nope, q_pe, k_nope, k_pe.
97 - 8. Attention is computued with q, k, v.
98 - 9. The attention computation returns (B, N, V), which is projected back
99 - to (B, H) using out projection.
100 -
101 - * if decode:
102 -
103 - 3. Here's the change, we do not perform up the full up projection for
104 - q_c, and there is no up projection at all for kv_c. This is
105 - achieved by the technique of "weight absorption". The paper says
106 - "Fortunately, due to the associative law of matrix multiplication,
107 - we can absorb WUK into WUQ, and WUV into WO"
108 - * The q up projection turns (B, Lq) into (B, N, (P+R)), we split it
109 - into W_UQ (Lq, N, P) and W_QR (Lq, N, R).
110 - * The kv_c up projection turns (B, Lkv) into (B, N, (P+V)), we split
111 - it into W_UK (Lkv, N, P) and W_UV (Lkv, N, V).
112 - * The out projection shape W_O (N*V, H) turns (B, N, V) into (B, H).
113 - * We can precompute the product of W_UQ and W_UK into
114 - W_UQ_UK (Lq, N, Lkv), which is possible due to QK^T operation in
115 - attention.
116 - * We can precompute the product of W_UV and W_O into
117 - W_UV_O (N, Lkv, H), which is possible due to V@O as the
118 - "epilogue" of attention
119 - 4. We still need to compute q_pe (B, N, R) by applying W_QR to q_latent.
120 - 5. q_pe, k_pe are then passed through rotary embeddings.
121 - 6. kv_c and k_pe are concatenated and inserted into the cache
122 - 7. By applying W_UQ_UK to q_latent, we have the new q_nope of shape
123 - (B, N, Lkv).
124 - 8. q (B, N, (Lkv+R)), k (B, (Lkv+R)) are assembled from q_nope, q_pe,
125 - kv_a, k_pe. v (B, Lkv) is exactly the same vector as kv_a.
126 - 9. The attention is computed with q, k, v. Note that we just performed
127 - a MQA attention with (LKv+R) as our head dim.
128 - 10. The KV cache is updated using the new entries k (B, N, (Lkv+R)),
129 - which included the v and rope values.
130 - 11. The attention computation returns (B, N, Lkv), which is projected
131 - back to (B, H) using W_UV_O.
132 -
133 - From @tsu-bin's calculation, we only want to use the absorption technique
134 - for decode. The prefill algorithm should still use the up-projected MHA
135 - for less flops and memory usage.
136 -
137 - """
138 -
139 - def __init__(
140 - self,
141 - num_heads: int,
142 - head_size: int,
143 - scale: float,
144 - num_kv_heads: int,
145 - alibi_slopes: Optional[List[float]],
146 - sliding_window: Optional[int],
147 - kv_cache_dtype: str,
148 - blocksparse_params: Optional[Dict[str, Any]],
149 - logits_soft_cap: Optional[float],
150 - attn_type: str,
151 - # MLA Specific Arguments
152 - q_lora_rank: Optional[int],
153 - kv_lora_rank: int,
154 - qk_nope_head_dim: int,
155 - qk_rope_head_dim: int,
156 - qk_head_dim: int,
157 - v_head_dim: int,
158 - rotary_emb: RotaryEmbedding,
159 - # q_proj should be q_b_proj if q_lora_rank is not None, but from an
160 - # attention backend perspective we rely on the layer to pass in the
161 - # correct matrix
162 - q_proj: ColumnParallelLinear,
163 - kv_b_proj: ColumnParallelLinear,
164 - o_proj: RowParallelLinear,
165 - ) -> None:
166 - self.num_heads = num_heads
167 - self.head_size = head_size
168 - self.scale = float(scale)
169 - self.num_kv_heads = num_kv_heads
170 - self.kv_cache_dtype = kv_cache_dtype
171 -
172 - self.q_lora_rank = q_lora_rank
173 - self.kv_lora_rank = kv_lora_rank
174 - self.qk_nope_head_dim = qk_nope_head_dim
175 - self.qk_rope_head_dim = qk_rope_head_dim
176 - self.qk_head_dim = qk_head_dim
177 - self.v_head_dim = v_head_dim
178 -
179 - self.rotary_emb = rotary_emb
180 - self.use_yarn_rope = isinstance(rotary_emb,
181 - DeepseekScalingRotaryEmbedding)
182 - self.q_proj = q_proj
183 - self.kv_b_proj = kv_b_proj
184 - self.o_proj = o_proj
185 - self.vllm_flash_attn_version = get_flash_attn_version()
186 -
187 - # Handle the differences between the flash_attn_varlen from flash_attn
188 - # and the one from vllm_flash_attn. The former is used on RoCM and the
189 - # latter has an additional parameter to control FA2 vs FA3
190 - self.flash_attn_varlen_func = flash_attn_varlen_func
191 - if self.vllm_flash_attn_version is not None:
192 - self.flash_attn_varlen_func = \
193 - functools.partial(flash_attn_varlen_func,
194 - fa_version=self.vllm_flash_attn_version)
195 -
196 - def _v_up_proj_and_o_proj(self, x):
197 - if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
198 - if is_fp8(self.W_UV_O):
199 - output_parallel = apply_fp8_linear_generic(
200 - x.flatten(start_dim=1), self.W_UV_O, self.W_UV_O_scales,
201 - self.reqaunt_input_group_shape,
202 - self.reqaunt_weight_group_shape)
203 - else:
204 - output_parallel = torch.matmul(x.flatten(start_dim=1),
205 - self.W_UV_O)
206 - if self.tp_size > 1:
207 - output = tensor_model_parallel_all_reduce(output_parallel)
208 - else:
209 - output = output_parallel
210 - return output
211 - else:
212 - x = torch.einsum("bnl,lnv->bnv", x, self.W_UV)
213 - return self.o_proj(x.reshape(-1,
214 - self.num_heads * self.v_head_dim))[0]
215 -
216 - def _q_proj_and_k_up_proj(self, x):
217 - if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
218 - if is_fp8(self.W_Q_UK):
219 - return apply_fp8_linear_generic(
220 - x, self.W_Q_UK, self.W_Q_UK_scales,
221 - self.reqaunt_input_group_shape,
222 - self.reqaunt_weight_group_shape).view(
223 - -1, self.num_heads, self.kv_lora_rank)
224 - return torch.matmul(x, self.W_Q_UK)\
225 - .view(-1, self.num_heads, self.kv_lora_rank)
226 - else:
227 - x = torch.matmul(x, self.W_Q)\
228 - .view(-1, self.num_heads, self.qk_nope_head_dim)
229 - return torch.einsum("bnp,lnp->bnl", x, self.W_UK)\
230 - .view(-1, self.num_heads, self.kv_lora_rank)
231 -
232 - def process_weights_after_loading(self, act_dtype: torch.dtype):
233 - # TODO(lucas) This is very gross, we need a more wide scale refactor of
234 - # all the FP8 code with a more standard way of
235 - # defining schemes/group-shapes, we should also potentially force
236 - # quant_methods to support a decompress function
237 - #
238 - # returns input_group_shape, weight_group_shape
239 - def get_scale_group_shapes_for_fp8(layer: LinearBase) -> \
240 - Tuple[Tuple[int, int], Tuple[int, int]]:
241 - if isinstance(layer.quant_method, Fp8LinearMethod):
242 - if layer.quant_method.block_quant:
243 - weight_block_size = \
244 - layer.quant_method.quant_config.weight_block_size
245 - # per-token-group (1, X), block-quantized (X, Y)
246 - return (1, weight_block_size[-1]), weight_block_size
247 - else:
248 - return (-1, -1), (-1, -1) # per-tensor, per-tensor
249 - elif isinstance(layer.quant_method, CompressedTensorsLinearMethod)\
250 - and isinstance(layer.scheme, CompressedTensorsW8A8Fp8):
251 - # this is hacky but we always assume the for
252 - # CompressedTensorsW8A8Fp8 the input is dynamic per-token
253 - # we ignore if it is static-per-tensor since we are going to
254 - # requantize after later anyways
255 - strategy = layer.scheme.strategy
256 - if strategy == QuantizationStrategy.TENSOR:
257 - return (1, -1), (-1, -1) # per-token, per-tensor
258 - elif strategy == QuantizationStrategy.CHANNEL:
259 - return (1, -1), (-1, 1) # per-token, per-channel
260 - else:
261 - raise NotImplementedError(
262 - f"QuantizationStrategy.{strategy} is not supported for "
263 - "fp8 MLA, please run with VLLM_MLA_DISABLE=1")
264 - else:
265 - raise NotImplementedError(
266 - "Can't determine scale group shapes for "
267 - f"{layer.quant_method}, please run with VLLM_MLA_DISABLE=1"
268 - )
269 -
270 - def get_layer_weight(layer):
271 - if hasattr(layer, "weight"):
272 - return layer.weight
273 - elif hasattr(layer, "qweight"):
274 - return layer.qweight
275 - else:
276 - raise AttributeError(
277 - f"Layer '{layer}' has neither weight nor qweight")
278 -
279 - def get_and_maybe_dequant_weights(layer: LinearBase):
280 - if not isinstance(layer.quant_method, UnquantizedLinearMethod):
281 - # NOTE: This should only be used offline, since it's O(N^3)
282 - eye = torch.eye(layer.input_size_per_partition,
283 - dtype=act_dtype,
284 - device=get_layer_weight(layer).device)
285 - dequant_weights = layer.quant_method.apply(layer,
286 - eye,
287 - bias=None)
288 - del eye
289 - # standardize to (output, input)
290 - return dequant_weights.T
291 - return layer.weight
292 -
293 - weight_dtype = get_layer_weight(self.kv_b_proj).dtype
294 - assert get_layer_weight(self.o_proj).dtype == weight_dtype
295 - assert get_layer_weight(self.q_proj).dtype == weight_dtype
296 -
297 - kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T
298 - assert kv_b_proj_weight.shape == (
299 - self.kv_lora_rank,
300 - self.num_heads * (self.qk_nope_head_dim + self.v_head_dim)), (
301 - f"{kv_b_proj_weight.shape=}, "
302 - f"{self.kv_lora_rank=}, "
303 - f"{self.num_heads=}, "
304 - f"{self.qk_nope_head_dim=}, "
305 - f"{self.v_head_dim=}")
306 - kv_b_proj_weight = kv_b_proj_weight.view(
307 - self.kv_lora_rank,
308 - self.num_heads,
309 - self.qk_nope_head_dim + self.v_head_dim,
310 - )
311 -
312 - W_UK, W_UV = kv_b_proj_weight.split(
313 - [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
314 -
315 - q_proj_weight = get_and_maybe_dequant_weights(self.q_proj).T\
316 - .view(-1, self.num_heads, self.qk_head_dim)
317 -
318 - # can be W_Q or W_UQ depending q_lora_rank, the former if
319 - # q_lora_rank is None, the latter otherwise. From the Attention backend
320 - # perspective though we call these both W_Q and rely on the layer
321 - # to pass in the correct matrix
322 - W_Q = q_proj_weight[..., :self.qk_nope_head_dim]
323 - self.W_QR = q_proj_weight[..., self.qk_nope_head_dim:]\
324 - .flatten(start_dim=1).contiguous()
325 -
326 - # W_QR is small so for simplicity we dont bother requantizing it
327 - self.W_QR = self.W_QR.to(act_dtype)
328 -
329 - if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
330 - requantization_enabled = not envs.VLLM_MLA_DISABLE_REQUANTIZATION
331 - if is_fp8(weight_dtype) and requantization_enabled:
332 - # This assumes it wise to requantize using the same group shapes
333 - # (i.e. strategy, per-tensor, per-channel, block etc.) that the
334 - # weights were originally quantized
335 - requant_input_group_shape, requant_weight_group_shape = \
336 - get_scale_group_shapes_for_fp8(self.q_proj)
337 - assert (requant_input_group_shape, requant_weight_group_shape)\
338 - == get_scale_group_shapes_for_fp8(self.kv_b_proj)
339 - assert (requant_input_group_shape, requant_weight_group_shape)\
340 - == get_scale_group_shapes_for_fp8(self.o_proj)
341 - self.reqaunt_input_group_shape = requant_input_group_shape
342 - self.reqaunt_weight_group_shape = requant_weight_group_shape
343 -
344 - #
345 - # Perform matrix-absorption following
346 - # https://github.com/flashinfer-ai/flashinfer/pull/551
347 - # for decode, as a result we end up with absorbed weights for decode
348 - # and another copy of raw weights for prefill.
349 - #
350 - self.W_UK, self.W_UV = kv_b_proj_weight.split(
351 - [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
352 - # We absorb `W_UK` into `W_Q` resulting in either W_Q_UK or W_UQ_UK
353 - # depending q_lora_rank, the former if q_lora_rank is None, the
354 - # latter otherwise
355 - # basically if q_lora_rank is none we are absorbing into q_proj
356 - # instead of UQ
357 - W_Q_UK = torch.einsum("qnd,lnd -> qnl", W_Q, W_UK)\
358 - .flatten(start_dim=1).contiguous()
359 -
360 - if is_fp8(weight_dtype) and requantization_enabled:
361 - W_Q_UK, W_Q_UK_scales = scaled_quantize(
362 - W_Q_UK,
363 - self.reqaunt_weight_group_shape,
364 - quant_dtype=current_platform_fp8_dtype)
365 - # For FP8 save the transpose so we can use
366 - # `apply_w8a8_block_fp8_linear` directly
367 - self.W_Q_UK = W_Q_UK.T.contiguous()
368 - self.W_Q_UK_scales = W_Q_UK_scales.T.contiguous()
369 - else:
370 - self.W_Q_UK = W_Q_UK.to(act_dtype)
371 -
372 - W_O = get_and_maybe_dequant_weights(self.o_proj)\
373 - .view(-1, self.num_heads, self.v_head_dim)
374 - W_UV_O = torch.einsum("lnd,hnd -> nlh", W_UV, W_O)\
375 - .flatten(start_dim=0, end_dim=1).contiguous()
376 -
377 - if is_fp8(weight_dtype) and requantization_enabled:
378 - W_UV_O, W_UV_O_scales = scaled_quantize(
379 - W_UV_O,
380 - self.reqaunt_weight_group_shape,
381 - quant_dtype=current_platform_fp8_dtype)
382 - # For FP8 save the transpose so we can use
383 - # `apply_w8a8_block_fp8_linear` directly
384 - self.W_UV_O = W_UV_O.T.contiguous()
385 - self.W_UV_O_scales = W_UV_O_scales.T.contiguous()
386 - else:
387 - self.W_UV_O = W_UV_O.to(act_dtype)
388 -
389 - self.tp_size = get_tensor_model_parallel_world_size()
390 - else:
391 - if is_fp8(weight_dtype):
392 - raise NotImplementedError(
393 - "Currently fp8 requires matrix absorption")
394 -
395 - self.W_UV = W_UV
396 - self.W_UK = W_UK
397 - self.W_Q = W_Q.flatten(start_dim=1)
398 -
399 - @abstractmethod
400 - def _forward_prefill(
401 - self,
402 - q: torch.Tensor,
403 - kv_c_normed: torch.Tensor,
404 - k_pe: torch.Tensor,
405 - attn_metadata: T,
406 - ) -> torch.Tensor:
407 - raise NotImplementedError
408 -
409 - @abstractmethod
410 - def _forward_decode(
411 - self,
412 - q_nope: torch.Tensor,
413 - q_pe: torch.Tensor,
414 - kv_cache: torch.Tensor,
415 - attn_metadata: T,
416 - ) -> torch.Tensor:
417 - raise NotImplementedError
418 -
419 - def forward(
420 - self,
421 - layer: AttentionLayer,
422 - hidden_states_or_q_c: torch.Tensor, # query in unified attn
423 - k_c_normed: torch.Tensor, # key in unified attn
424 - k_pe: torch.Tensor, # value in unified attn
425 - kv_cache: torch.Tensor,
426 - attn_metadata: T,
427 - output: Optional[torch.Tensor] = None,
428 - ) -> torch.Tensor:
429 - if output is not None:
430 - raise NotImplementedError(
431 - "output is not yet supported for MLAImplBase")
432 -
433 - is_decode = attn_metadata.decode_metadata is not None
434 - is_prefill = attn_metadata.prefill_metadata is not None
435 -
436 - if (is_decode and is_prefill):
437 - raise NotImplementedError(
438 - "chunked prefill is not supported for MLAImplBase")
439 -
440 - # Restore head dim (for rotary embedding)
441 - k_pe = k_pe.unsqueeze(1)
442 - assert hasattr(attn_metadata, "input_positions")
443 -
444 - if is_decode:
445 - q_nope = self._q_proj_and_k_up_proj(hidden_states_or_q_c)
446 - q_pe = torch.matmul(hidden_states_or_q_c, self.W_QR)\
447 - .view(-1, self.num_heads, self.qk_rope_head_dim)
448 - q_pe, k_pe = self.rotary_emb(attn_metadata.input_positions, q_pe,
449 - k_pe)
450 - else:
451 - assert is_prefill
452 - q = self.q_proj(hidden_states_or_q_c)[0]\
453 - .view(-1, self.num_heads, self.qk_head_dim)
454 -
455 - # TODO(lucas): there must be a nicer way to write this line
456 - q[..., self.qk_nope_head_dim:], k_pe = \
457 - self.rotary_emb(
458 - attn_metadata.input_positions,
459 - q[..., self.qk_nope_head_dim:], k_pe)
460 -
461 - # write the latent and rope to kv cache
462 - if kv_cache.numel() > 0:
463 - ops.concat_and_cache_mla(
464 - k_c_normed,
465 - k_pe.squeeze(1),
466 - kv_cache,
467 - attn_metadata.slot_mapping.flatten(),
468 - kv_cache_dtype=self.kv_cache_dtype,
469 - scale=layer._k_scale,
470 - )
471 -
472 - if attn_metadata.prefill_metadata is not None:
473 - return self._forward_prefill(q, k_c_normed, k_pe, attn_metadata)
474 -
475 - if attn_metadata.decode_metadata is not None:
476 - return self._forward_decode(q_nope, q_pe, kv_cache, attn_metadata)
477 -
478 - # Optional common flash-attn based prefill
479 - def _forward_prefill_flash(
480 - self,
481 - q: torch.Tensor,
482 - k_c_normed: torch.Tensor,
483 - k_pe: torch.Tensor,
484 - seq_start_loc: torch.Tensor,
485 - max_prefill_seq_len: int,
486 - ) -> torch.Tensor:
487 -
488 - kv_nope = self.kv_b_proj(k_c_normed)[0]\
489 - .view(-1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
490 - k_nope, v = kv_nope\
491 - .split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
492 -
493 - k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
494 -
495 - # For MLA the v head dim is smaller than qk head dim so we pad out
496 - # v with 0s to match the qk head dim
497 - v_padded = torch.nn.functional.pad(v, [0, q.shape[-1] - v.shape[-1]],
498 - value=0)
499 -
500 - attn_output = self.flash_attn_varlen_func(
501 - q=q,
502 - k=k,
503 - v=v_padded,
504 - cu_seqlens_q=seq_start_loc,
505 - cu_seqlens_k=seq_start_loc,
506 - max_seqlen_q=max_prefill_seq_len,
507 - max_seqlen_k=max_prefill_seq_len,
508 - softmax_scale=self.scale,
509 - causal=True,
510 - )
511 - attn_output = attn_output\
512 - .view(-1, self.num_heads, q.shape[-1])[..., :v.shape[-1]]\
513 - .reshape(-1, self.num_heads * v.shape[-1])
514 -
515 - return self.o_proj(attn_output)[0]
vllm/attention/backends/triton_mla.py
+14 -650
@@ -1,40 +1,17 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 -from collections import defaultdict
4 -from contextlib import contextmanager
5 -from dataclasses import dataclass
6 -from itertools import accumulate
7 -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
8 -
9 -from vllm.multimodal import MultiModalPlaceholderMap
10 -
11 -try:
12 - from flashinfer import BatchDecodeMlaWithPagedKVCacheWrapper
13 - FLASHINFER_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024
14 -except ImportError:
15 - BatchDecodeMlaWithPagedKVCacheWrapper = None
16 - FLASHINFER_WORKSPACE_BUFFER_SIZE = 0
3 +from typing import Any, Dict, List, Optional, Type
4
5 import torch
6
20 -from vllm import _custom_ops as ops
21 -from vllm.attention.backends.abstract import (AttentionBackend,
22 - AttentionMetadata,
23 - AttentionMetadataBuilder,
24 - AttentionState, AttentionType)
25 -from vllm.attention.backends.mla.utils import MLACommonImpl, MLACommonMetadata
26 -from vllm.attention.backends.utils import (PAD_SLOT_ID, compute_slot_mapping,
27 - compute_slot_mapping_start_idx,
28 - is_block_tables_empty)
7 +from vllm.attention.backends.abstract import AttentionType
8 +from vllm.attention.backends.mla.common import (MLACommonBackend,
9 + MLACommonImpl,
10 + MLACommonMetadata)
11 from vllm.attention.ops.triton_decode_attention import decode_attention_fwd
30 -from vllm.utils import async_tensor_h2d, make_tensor_with_pad
31 -
32 -if TYPE_CHECKING:
33 - from vllm.worker.model_runner import (ModelInputForGPUBuilder,
34 - ModelInputForGPUWithSamplingMetadata)
12
13
37 -class TritonMLABackend(AttentionBackend):
14 +class TritonMLABackend(MLACommonBackend):
15
16 @staticmethod
17 def get_name() -> str:
@@ -44,610 +21,8 @@ class TritonMLABackend(AttentionBackend):
21 def get_impl_cls() -> Type["TritonMLAImpl"]:
22 return TritonMLAImpl
23
47 - @staticmethod
48 - def get_metadata_cls() -> Type["AttentionMetadata"]:
49 - return TritonMLAMetadata
50 -
51 - @staticmethod
52 - def get_builder_cls() -> Type["TritonMLAMetadataBuilder"]:
53 - return TritonMLAMetadataBuilder
54 -
55 - @staticmethod
56 - def get_state_cls() -> Type["TritonMLAState"]:
57 - return TritonMLAState
58 -
59 - @staticmethod
60 - def get_kv_cache_shape(
61 - num_blocks: int,
62 - block_size: int,
63 - num_kv_heads: int, # assumed to be 1 for MLA
64 - head_size: int,
65 - ) -> Tuple[int, ...]:
66 - return (num_blocks, block_size, head_size)
67 -
68 - @staticmethod
69 - def swap_blocks(
70 - src_kv_cache: torch.Tensor,
71 - dst_kv_cache: torch.Tensor,
72 - src_to_dst: torch.Tensor,
73 - ) -> None:
74 - ops.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)
75 -
76 - @staticmethod
77 - def copy_blocks(
78 - kv_caches: List[torch.Tensor],
79 - src_to_dists: torch.Tensor,
80 - ) -> None:
81 - ops.copy_blocks_mla(kv_caches, src_to_dists)
82 -
83 - @staticmethod
84 - def get_supported_head_sizes() -> List[int]:
85 - return [576]
86 -
87 -
88 -class TritonMLAState(AttentionState):
89 -
90 - def __init__(self, runner):
91 - self.runner = runner
92 - self._is_graph_capturing = False
93 -
94 - @contextmanager
95 - def graph_capture(self, max_batch_size: int):
96 - self._is_graph_capturing = True
97 -
98 - self._graph_slot_mapping = torch.full((max_batch_size, ),
99 - PAD_SLOT_ID,
100 - dtype=torch.long,
101 - device=self.runner.device)
102 - self._graph_seq_lens = torch.ones(max_batch_size,
103 - dtype=torch.int32,
104 - device=self.runner.device)
105 - self._graph_block_tables = torch.from_numpy(
106 - self.runner.graph_block_tables).to(device=self.runner.device)
107 -
108 - self._positions = torch.zeros((max_batch_size, ),
109 - dtype=torch.long,
110 - device=self.runner.device)
111 -
112 - yield
113 -
114 - self._is_graph_capturing = False
115 - del self._graph_slot_mapping
116 - del self._graph_seq_lens
117 - del self._graph_block_tables
118 - del self._positions
119 -
120 - def graph_clone(self, batch_size: int):
121 - assert self._is_graph_capturing
122 - return self.__class__(self.runner)
123 -
124 - def graph_capture_get_metadata_for_batch(
125 - self, batch_size: int, is_encoder_decoder_model: bool = False):
126 - assert self._is_graph_capturing
127 -
128 - attn_metadata = self.runner.attn_backend.make_metadata(
129 - num_prefills=0,
130 - num_prefill_tokens=0,
131 - num_decode_tokens=batch_size,
132 - slot_mapping=self._graph_slot_mapping[:batch_size],
133 - multi_modal_placeholder_index_maps=None,
134 - enable_kv_scales_calculation=True,
135 - seq_lens=None,
136 - seq_lens_tensor=self._graph_seq_lens[:batch_size],
137 - max_query_len=1,
138 - max_decode_query_len=1,
139 - max_prefill_seq_len=0,
140 - max_decode_seq_len=self.runner.max_seq_len_to_capture,
141 - query_start_loc=None,
142 - seq_start_loc=None,
143 - context_lens_tensor=None,
144 - block_tables=self._graph_block_tables[:batch_size],
145 - use_cuda_graph=True,
146 - input_positions=self._positions[:batch_size],
147 - head_dim=self.runner.model_config.get_head_size())
148 -
149 - if is_encoder_decoder_model:
150 - raise NotImplementedError(
151 - "TritonMLAState does not support encoder/decoder yet")
152 -
153 - return attn_metadata
154 -
155 - def get_graph_input_buffers(self,
156 - attn_metadata,
157 - is_encoder_decoder_model: bool = False):
158 - input_buffers = {
159 - "slot_mapping": attn_metadata.slot_mapping,
160 - "seq_lens_tensor": attn_metadata.decode_metadata.seq_lens_tensor,
161 - "block_tables": attn_metadata.decode_metadata.block_tables,
162 - "input_positions": attn_metadata.decode_metadata.input_positions,
163 - }
164 - if is_encoder_decoder_model:
165 - raise NotImplementedError(
166 - "TritonMLAState does not support encoder/decoder yet")
167 -
168 - return input_buffers
169 -
170 - def prepare_graph_input_buffers(self,
171 - input_buffers,
172 - attn_metadata,
173 - is_encoder_decoder_model: bool = False):
174 - input_positions = attn_metadata.input_positions
175 - num_positions = input_positions.shape[0]
176 - input_buffers["seq_lens_tensor"].copy_(
177 - attn_metadata.decode_metadata.seq_lens_tensor, non_blocking=True)
178 - input_buffers["block_tables"].copy_(
179 - attn_metadata.decode_metadata.block_tables, non_blocking=True)
180 - # CUDA graph buffer is padded so only perform a partial copy based on
181 - # num_positions
182 - input_buffers["input_positions"][:num_positions].copy_(
183 - input_positions, non_blocking=True)
184 - if is_encoder_decoder_model:
185 - raise NotImplementedError(
186 - "TritonMLAState does not support encoder/decoder yet")
187 -
188 - def begin_forward(self, model_input):
189 - return
190 -
191 -
192 -@dataclass
193 -class TritonMLAMetadata(MLACommonMetadata):
194 - """Metadata for TritonMLAMetadata.
195 -
196 - NOTE: Any python object stored here is not updated when it is
197 - cuda-graph replayed. If you have values that need to be changed
198 - dynamically, it should be stored in tensor. The tensor has to be
199 - updated from `CUDAGraphRunner.forward` API.
200 - """
201 - # (batch_size,). The sequence length per sequence. Sequence length means
202 - # the computed tokens + new tokens None if it is a decoding.
203 - seq_lens: Optional[List[int]]
204 - # seq_lens stored as a tensor.
205 - seq_lens_tensor: Optional[torch.Tensor]
206 -
207 - # NOTE(sang): Definition of context_len, query_len, and seq_len.
208 - # |---------- N-1 iteration --------|
209 - # |---------------- N iteration ---------------------|
210 - # |- tokenA -|......................|-- newTokens ---|
211 - # |---------- context_len ----------|
212 - # |-------------------- seq_len ---------------------|
213 - # |-- query_len ---|
214 -
215 - # Maximum sequence length among prefill batch. 0 if there are decoding
216 - # requests only.
217 - max_prefill_seq_len: int
218 - # Maximum sequence length among decode batch. 0 if there are prefill
219 - # requests only.
220 - max_decode_seq_len: int
221 - # (batch_size,) A tensor of context lengths (tokens that are computed
222 - # so far).
223 - context_lens_tensor: Optional[torch.Tensor]
224 -
225 - # (batch_size, max_blocks_per_seq).
226 - # Block addresses per sequence. (Seq id -> list of physical block)
227 - # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks
228 - # in the kv cache. Each block can contain up to block_size tokens.
229 - # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph
230 - # captured.
231 - block_tables: Optional[torch.Tensor]
232 -
233 - # Whether or not if cuda graph is enabled.
234 - # Cuda-graph is currently enabled for decoding only.
235 - # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.
236 -
237 - use_cuda_graph: bool
238 -
239 - # Maximum query length in the batch.
240 - max_query_len: Optional[int] = None
241 -
242 - # Max number of query tokens among request in the batch.
243 - max_decode_query_len: Optional[int] = None
244 -
245 - # (batch_size + 1,). The cumulative subquery lengths of the sequences in
246 - # the batch, used to index into subquery. E.g., if the subquery length
247 - # is [4, 6], it is [0, 4, 10].
248 - query_start_loc: Optional[torch.Tensor] = None
249 - # (batch_size + 1,). The cumulative sequence lengths of the sequences in
250 - # the batch, used to index into sequence. E.g., if the sequence length is
251 - # [4, 6], it is [0, 4, 10].
252 - seq_start_loc: Optional[torch.Tensor] = None
253 -
254 - _cached_prefill_metadata: Optional["TritonMLAMetadata"] = None
255 - _cached_decode_metadata: Optional["TritonMLAMetadata"] = None
24
257 - num_prefill_tokens: int
258 -
259 - num_kv_splits: int = 4 # TODO(lucas) add heuristic
260 - attn_logits: Optional[torch.Tensor] = None
261 - req_idx: Optional[torch.Tensor] = None
262 -
263 - # The dimension of the attention heads
264 - head_dim: Optional[int] = None
265 -
266 - def __post_init__(self):
267 - supported_head_sizes = TritonMLABackend.get_supported_head_sizes()
268 - if self.head_dim is not None and self.head_dim \
269 - not in supported_head_sizes:
270 - raise ValueError(
271 - f"Only {supported_head_sizes} are supported for head_dim,",
272 - f"received {self.head_dim}.")
273 -
274 - @property
275 - def prefill_metadata(self) -> Optional["TritonMLAMetadata"]:
276 - if self.num_prefills == 0:
277 - return None
278 -
279 - if self._cached_prefill_metadata is not None:
280 - return self._cached_prefill_metadata
281 -
282 - assert self.seq_lens is not None
283 - assert self.seq_lens_tensor is not None
284 -
285 - # Compute some attn_metadata fields which default to None
286 - query_start_loc = (None if self.query_start_loc is None else
287 - self.query_start_loc[:self.num_prefills + 1])
288 - slot_mapping = (None if self.slot_mapping is None else
289 - self.slot_mapping[:self.num_prefill_tokens])
290 - seq_lens = (None if self.seq_lens is None else
291 - self.seq_lens[:self.num_prefills])
292 - seq_lens_tensor = (None if self.seq_lens_tensor is None else
293 - self.seq_lens_tensor[:self.num_prefills])
294 - seq_start_loc = (None if self.seq_start_loc is None else
295 - self.seq_start_loc[:self.num_prefills + 1])
296 - context_lens_tensor = (None if self.context_lens_tensor is None else
297 - self.context_lens_tensor[:self.num_prefills])
298 - block_tables = (None if self.block_tables is None else
299 - self.block_tables[:self.num_prefills])
300 - input_positions = (None if self.input_positions is None else
301 - self.input_positions[:self.num_prefill_tokens])
302 -
303 - self._cached_prefill_metadata = TritonMLAMetadata(
304 - num_prefills=self.num_prefills,
305 - num_prefill_tokens=self.num_prefill_tokens,
306 - num_decode_tokens=0,
307 - slot_mapping=slot_mapping,
308 - multi_modal_placeholder_index_maps=self.
309 - multi_modal_placeholder_index_maps,
310 - enable_kv_scales_calculation=self.enable_kv_scales_calculation,
311 - input_positions=input_positions,
312 - seq_lens=seq_lens,
313 - seq_lens_tensor=seq_lens_tensor,
314 - max_query_len=self.max_query_len,
315 - max_prefill_seq_len=self.max_prefill_seq_len,
316 - max_decode_query_len=0,
317 - max_decode_seq_len=0,
318 - query_start_loc=query_start_loc,
319 - seq_start_loc=seq_start_loc,
320 - context_lens_tensor=context_lens_tensor,
321 - block_tables=block_tables,
322 - use_cuda_graph=False,
323 - head_dim=self.head_dim)
324 - return self._cached_prefill_metadata
325 -
326 - @property
327 - def decode_metadata(self) -> Optional["TritonMLAMetadata"]:
328 - if self.num_decode_tokens == 0:
329 - return None
330 -
331 - if self._cached_decode_metadata is not None:
332 - return self._cached_decode_metadata
333 - assert self.seq_lens_tensor is not None
334 -
335 - # Compute some attn_metadata fields which default to None
336 - slot_mapping = (None if self.slot_mapping is None else
337 - self.slot_mapping[self.num_prefill_tokens:])
338 - seq_lens_tensor = (None if self.seq_lens_tensor is None else
339 - self.seq_lens_tensor[self.num_prefills:])
340 - block_tables = (None if self.block_tables is None else
341 - self.block_tables[self.num_prefills:])
342 - input_positions = (None if self.input_positions is None else
343 - self.input_positions[self.num_prefill_tokens:])
344 -
345 - self._cached_decode_metadata = TritonMLAMetadata(
346 - num_prefills=0,
347 - num_prefill_tokens=0,
348 - num_decode_tokens=self.num_decode_tokens,
349 - slot_mapping=slot_mapping,
350 - multi_modal_placeholder_index_maps=None,
351 - enable_kv_scales_calculation=True,
352 - seq_lens=None,
353 - seq_lens_tensor=seq_lens_tensor,
354 - max_decode_query_len=self.max_decode_query_len,
355 - max_query_len=self.max_query_len,
356 - max_prefill_seq_len=0,
357 - max_decode_seq_len=self.max_decode_seq_len,
358 - # Batch may be composed of prefill|decodes, adjust query start
359 - # indices to refer to the start of decodes. E.g.
360 - # in tokens:[3 prefills|6 decodes], query_start_loc=[3,9] => [0,6].
361 - query_start_loc=(self.query_start_loc[self.num_prefills:] -
362 - self.query_start_loc[self.num_prefills])
363 - if self.query_start_loc is not None else None,
364 - seq_start_loc=self.seq_start_loc[self.num_prefills:]
365 - if self.seq_start_loc is not None else None,
366 - context_lens_tensor=None,
367 - block_tables=block_tables,
368 - use_cuda_graph=self.use_cuda_graph,
369 - input_positions=input_positions,
370 - head_dim=self.head_dim)
371 - return self._cached_decode_metadata
372 -
373 - def advance_step(self,
374 - model_input: "ModelInputForGPUWithSamplingMetadata",
375 - sampled_token_ids: Optional[torch.Tensor],
376 - block_size: int,
377 - num_seqs: int,
378 - num_queries: int,
379 - turn_prefills_into_decodes: bool = False):
380 - """
381 - Update metadata in-place to advance one decode step.
382 - """
383 - # When using cudagraph, the num_seqs is padded to the next captured
384 - # batch sized, but num_queries tracks the actual number of requests in
385 - # the batch. For --enforce-eager mode, num_seqs == num_queries
386 - if num_seqs != num_queries:
387 - assert num_seqs > num_queries
388 - assert self.use_cuda_graph
389 -
390 - if turn_prefills_into_decodes:
391 - # When Mutli-Step is enabled with Chunked-Prefill, prefills and
392 - # decodes are scheduled together. In the first step, all the
393 - # prefills turn into decodes. This update reflects that
394 - # conversion.
395 - assert self.num_decode_tokens + self.num_prefills == num_seqs
396 - self.num_decode_tokens += self.num_prefills
397 - self.num_prefills = 0
398 - self.num_prefill_tokens = 0
399 - self.max_prefill_seq_len = 0
400 - self.max_query_len = 1
401 -
402 - self.slot_mapping = self.slot_mapping[:num_seqs]
403 - else:
404 - assert self.seq_lens is not None
405 - assert self.max_decode_seq_len == max(self.seq_lens)
406 -
407 - assert self.num_prefills == 0
408 - assert self.num_prefill_tokens == 0
409 - assert self.num_decode_tokens == num_seqs
410 - assert self.slot_mapping.shape == (num_seqs, )
411 -
412 - assert self.seq_lens is not None
413 - assert len(self.seq_lens) == num_seqs
414 - assert self.seq_lens_tensor is not None
415 - assert self.seq_lens_tensor.shape == (num_seqs, )
416 - assert self.max_query_len == 1
417 - assert self.max_prefill_seq_len == 0
418 -
419 - assert self.query_start_loc is not None
420 - assert self.query_start_loc.shape == (num_queries + 1, )
421 - assert self.seq_start_loc is not None
422 - assert self.seq_start_loc.shape == (num_seqs + 1, )
423 -
424 - assert self.context_lens_tensor is not None
425 - assert self.context_lens_tensor.shape == (num_queries, )
426 -
427 - assert self.block_tables is not None
428 - assert self.block_tables.shape[0] == num_seqs
429 -
430 - # Update query lengths. Note that we update only queries and not seqs,
431 - # since tensors may be padded due to captured cuda graph batch size
432 - for i in range(num_queries):
433 - self.seq_lens[i] += 1
434 - self.max_decode_seq_len = max(self.seq_lens)
435 -
436 - ops.advance_step_flashattn(num_seqs=num_seqs,
437 - num_queries=num_queries,
438 - block_size=block_size,
439 - input_tokens=model_input.input_tokens,
440 - sampled_token_ids=sampled_token_ids,
441 - input_positions=model_input.input_positions,
442 - seq_lens=self.seq_lens_tensor,
443 - slot_mapping=self.slot_mapping,
444 - block_tables=self.block_tables)
445 -
446 -
447 -class TritonMLAMetadataBuilder(AttentionMetadataBuilder[TritonMLAMetadata]):
448 -
449 - def __init__(self, input_builder: "ModelInputForGPUBuilder"):
450 - self.input_builder = input_builder
451 - self.runner = input_builder.runner
452 - self.sliding_window = input_builder.sliding_window
453 - self.block_size = input_builder.block_size
454 -
455 - def prepare(self):
456 - self.slot_mapping: List[int] = []
457 - self.prefill_seq_lens: List[int] = []
458 - self.context_lens: List[int] = []
459 - self.block_tables: List[List[int]] = []
460 - self.curr_seq_lens: List[int] = []
461 - self.input_positions: List[int] = []
462 - self.multimodal_placeholder_maps: Dict[
463 - str,
464 - MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap)
465 - self.num_prefills = 0
466 - self.num_prefill_tokens = 0
467 - self.num_decode_tokens = 0
468 - self.has_prefix_cache_hit = False
469 -
470 - def _add_seq_group(
471 - self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup",
472 - chunked_prefill_enabled: bool, prefix_cache_hit: bool):
473 - """Add a sequence group to the metadata. Specifically update/append
474 - 1. context length.
475 - 2. block table.
476 - 3. slot mapping.
477 - """
478 - is_prompt = inter_data.is_prompt
479 - block_tables = inter_data.block_tables
480 -
481 - for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,
482 - curr_sliding_window_block, input_positions) in zip(
483 - inter_data.seq_ids, [len(t) for t in inter_data.input_tokens],
484 - inter_data.orig_seq_lens, inter_data.seq_lens,
485 - inter_data.query_lens, inter_data.context_lens,
486 - inter_data.curr_sliding_window_blocks,
487 - inter_data.input_positions):
488 - self.input_positions.extend(input_positions)
489 - self.context_lens.append(context_len)
490 - if is_prompt:
491 - mm_maps = inter_data.multi_modal_placeholder_maps
492 - if mm_maps:
493 - for modality, placeholders in mm_maps.items():
494 - self.multimodal_placeholder_maps[modality].extend(
495 - placeholders)
496 -
497 - self.num_prefills += 1
498 - self.num_prefill_tokens += token_len
499 - self.prefill_seq_lens.append(seq_len)
500 - else:
501 - self.num_decode_tokens += query_len
502 - self.curr_seq_lens.append(curr_seq_len)
503 -
504 - # Compute block table.
505 - # TODO(sang): Combine chunked prefill and prefix caching by
506 - # only allowing multiple of block_size chunk size.
507 - # NOTE: This only works for oooooooxxx style attention.
508 - block_table = []
509 - if prefix_cache_hit:
510 - # NOTE(woosuk): For flash-attn, the block table should
511 - # include the entries for the incoming prefill tokens.
512 - block_table = block_tables[seq_id]
513 - elif ((chunked_prefill_enabled or not is_prompt)
514 - and block_tables is not None):
515 - if curr_sliding_window_block == 0:
516 - block_table = block_tables[seq_id]
517 - else:
518 - block_table = block_tables[seq_id][
519 - -curr_sliding_window_block:]
520 - self.block_tables.append(block_table)
521 -
522 - # Compute slot mapping.
523 - is_profile_run = is_block_tables_empty(block_tables)
524 - start_idx = compute_slot_mapping_start_idx(is_prompt, query_len,
525 - context_len,
526 - self.sliding_window)
527 - compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,
528 - seq_len, context_len, start_idx,
529 - self.block_size, inter_data.block_tables)
530 -
531 - def _get_graph_runner_block_tables(
532 - self, num_seqs: int,
533 - block_tables: List[List[int]]) -> torch.Tensor:
534 - # The shape of graph_block_tables is
535 - # [max batch size, max context len // block size].
536 - max_batch_size, max_blocks = self.runner.graph_block_tables.shape
537 - assert max_batch_size >= num_seqs
538 -
539 - graph_block_tables = self.runner.graph_block_tables[:num_seqs]
540 - for i, block_table in enumerate(block_tables):
541 - if block_table:
542 - num_blocks = len(block_table)
543 - if num_blocks <= max_blocks:
544 - graph_block_tables[i, :num_blocks] = block_table
545 - else:
546 - # It may be possible to have more blocks allocated due
547 - # to lookahead slots of multi-step, however, they are
548 - # not used anyway, so can be safely ignored.
549 - graph_block_tables[
550 - i, :max_blocks] = block_table[:max_blocks]
551 -
552 - return torch.from_numpy(graph_block_tables).to(
553 - device=self.runner.device, non_blocking=True)
554 -
555 - def build(self, seq_lens: List[int], query_lens: List[int],
556 - cuda_graph_pad_size: int, batch_size: int):
557 - """Build attention metadata with on-device tensors.
558 -
559 - Args:
560 - seq_lens: The maybe padded sequence lengths of the input sequences.
561 - query_lens: The query lengths of the input sequences.
562 - cuda_graph_pad_size: The padding size for cuda graph.
563 - -1 if cuda graph is not used.
564 - batch_size: The maybe padded batch size.
565 - """
566 - prefix_cache_hit = any([
567 - inter_data.prefix_cache_hit
568 - for inter_data in self.input_builder.inter_data_list
569 - ])
570 - for inter_data in self.input_builder.inter_data_list:
571 - self._add_seq_group(inter_data,
572 - self.input_builder.chunked_prefill_enabled,
573 - prefix_cache_hit)
574 -
575 - device = self.runner.device
576 - use_captured_graph = cuda_graph_pad_size != -1
577 -
578 - max_query_len = max(query_lens)
579 - decode_query_lens = query_lens[self.num_prefills:]
580 - if len(decode_query_lens) > 0:
581 - max_decode_query_len = max(decode_query_lens)
582 - else:
583 - max_decode_query_len = 1
584 - max_prefill_seq_len = max(self.prefill_seq_lens, default=0)
585 - max_decode_seq_len = max(self.curr_seq_lens, default=0)
586 - num_decode_tokens = self.num_decode_tokens
587 - query_start_loc = list(accumulate(query_lens, initial=0))
588 - seq_start_loc = list(accumulate(seq_lens, initial=0))
589 -
590 - num_seqs = len(seq_lens)
591 - if use_captured_graph:
592 - self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size)
593 - self.block_tables.extend([] * cuda_graph_pad_size)
594 - num_decode_tokens = batch_size - self.num_prefill_tokens
595 - block_tables = self._get_graph_runner_block_tables(
596 - num_seqs, self.block_tables)
597 - else:
598 - block_tables = make_tensor_with_pad(
599 - self.block_tables,
600 - pad=0,
601 - dtype=torch.int,
602 - device=device,
603 - )
604 - assert max_query_len > 0, ("query_lens: {}".format(query_lens))
605 -
606 - assert device is not None
607 - context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int,
608 - device, self.runner.pin_memory)
609 - seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device,
610 - self.runner.pin_memory)
611 - input_positions = async_tensor_h2d(self.input_positions, torch.long,
612 - device, self.runner.pin_memory)
613 - slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.long,
614 - device, self.runner.pin_memory)
615 - query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32,
616 - device,
617 - self.runner.pin_memory)
618 - seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32,
619 - device, self.runner.pin_memory)
620 - placeholder_index_maps = {
621 - modality: placeholder_map.index_map()
622 - for modality, placeholder_map in
623 - self.multimodal_placeholder_maps.items()
624 - }
625 -
626 - return TritonMLAMetadata(
627 - num_prefills=self.num_prefills,
628 - slot_mapping=slot_mapping_tensor,
629 - num_prefill_tokens=self.num_prefill_tokens,
630 - num_decode_tokens=num_decode_tokens,
631 - seq_lens=seq_lens,
632 - multi_modal_placeholder_index_maps=placeholder_index_maps,
633 - enable_kv_scales_calculation=True,
634 - input_positions=input_positions,
635 - seq_lens_tensor=seq_lens_tensor,
636 - max_query_len=max_query_len,
637 - max_decode_query_len=max_decode_query_len,
638 - max_prefill_seq_len=max_prefill_seq_len,
639 - max_decode_seq_len=max_decode_seq_len,
640 - query_start_loc=query_start_loc_tensor,
641 - seq_start_loc=seq_start_loc_tensor,
642 - context_lens_tensor=context_lens_tensor,
643 - block_tables=block_tables,
644 - use_cuda_graph=use_captured_graph,
645 - num_kv_splits=4, # TODO(lucas) add heuristic
646 - head_dim=self.runner.model_config.get_head_size(),
647 - )
648 -
649 -
650 -class TritonMLAImpl(MLACommonImpl[TritonMLAMetadata]):
25 +class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
26
27 def __init__(
28 self,
@@ -662,11 +37,11 @@ class TritonMLAImpl(MLACommonImpl[TritonMLAMetadata]):
37 logits_soft_cap: Optional[float],
38 attn_type: str,
39 # MLA Specific Arguments
665 - **kwargs) -> None:
40 + **mla_args) -> None:
41 super().__init__(num_heads, head_size, scale, num_kv_heads,
42 alibi_slopes, sliding_window, kv_cache_dtype,
43 blocksparse_params, logits_soft_cap, attn_type,
669 - **kwargs)
44 + **mla_args)
45
46 unsupported_features = [
47 alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
@@ -683,24 +58,12 @@ class TritonMLAImpl(MLACommonImpl[TritonMLAMetadata]):
58 "are not implemented for "
59 "TritonMLAImpl")
60
686 - def _forward_prefill(
687 - self,
688 - q: torch.Tensor,
689 - kv_c_normed: torch.Tensor,
690 - k_pe: torch.Tensor,
691 - attn_metadata: TritonMLAMetadata,
692 - ) -> torch.Tensor:
693 - assert isinstance(attn_metadata, TritonMLAMetadata)
694 - return self._forward_prefill_flash(q, kv_c_normed, k_pe,
695 - attn_metadata.seq_start_loc,
696 - attn_metadata.max_prefill_seq_len)
697 -
61 def _forward_decode(
62 self,
63 q_nope: torch.Tensor,
64 q_pe: torch.Tensor,
65 kv_c_and_k_pe_cache: torch.Tensor,
703 - attn_metadata: TritonMLAMetadata,
66 + attn_metadata: MLACommonMetadata,
67 ) -> torch.Tensor:
68 assert kv_c_and_k_pe_cache.numel() > 0
69 if self.kv_cache_dtype.startswith("fp8"):
@@ -717,12 +80,14 @@ class TritonMLAImpl(MLACommonImpl[TritonMLAMetadata]):
80 dtype=q.dtype,
81 device=q.device)
82
83 + num_kv_splits = 4 # TODO: heuristic
84 +
85 # TODO(lucas) Allocate ahead of time
86 attn_logits = torch.empty(
87 (
88 B,
89 self.num_heads,
725 - attn_metadata.num_kv_splits,
90 + num_kv_splits,
91 # NOTE(lucas) idk why the +1 is here but sglang has it so we
92 # just mirror that
93 self.kv_lora_rank + 1,
@@ -740,7 +105,6 @@ class TritonMLAImpl(MLACommonImpl[TritonMLAMetadata]):
105 decode_attention_fwd(q, kv_c_and_k_pe_cache, kv_c_cache, o,
106 decode_meta.block_tables,
107 decode_meta.seq_lens_tensor, attn_logits,
743 - attn_metadata.num_kv_splits, self.scale,
744 - PAGE_SIZE)
108 + num_kv_splits, self.scale, PAGE_SIZE)
109
110 return self._v_up_proj_and_o_proj(o)
vllm/attention/ops/triton_merge_attn_states.py new
+84
@@ -0,0 +1,84 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Optional
3 +
4 +import torch
5 +import triton
6 +import triton.language as tl
7 +
8 +
9 +# Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
10 +# can be used to combine partial attention results (in the split-KV case)
11 +def merge_attn_states(
12 + output: torch.Tensor,
13 + prefix_output: torch.Tensor,
14 + prefix_lse: torch.Tensor,
15 + suffix_output: torch.Tensor,
16 + suffix_lse: torch.Tensor,
17 + output_lse: Optional[torch.Tensor] = None,
18 +) -> None:
19 + num_tokens = output.shape[0]
20 + num_query_heads = output.shape[1]
21 + head_size = output.shape[2]
22 + padded_head_size = triton.next_power_of_2(head_size)
23 +
24 + # TODO(woosuk): Use CUDA kernel instead of Triton to minimize CPU overhead.
25 + merge_attn_states_kernel[(num_tokens, num_query_heads)](
26 + output,
27 + output_lse,
28 + prefix_output,
29 + prefix_lse,
30 + suffix_output,
31 + suffix_lse,
32 + head_size,
33 + padded_head_size,
34 + output_lse is not None,
35 + )
36 +
37 +
38 +@triton.jit
39 +def merge_attn_states_kernel(
40 + output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
41 + output_lse, # [NUM_HEADS, NUM_TOKENS]
42 + prefix_output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
43 + prefix_lse, # [NUM_HEADS, NUM_TOKENS]
44 + suffix_output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
45 + suffix_lse, # [NUM_HEADS, NUM_TOKENS]
46 + HEAD_SIZE: tl.constexpr,
47 + PADDED_HEAD_SIZE: tl.constexpr,
48 + OUTPUT_LSE: tl.constexpr,
49 +):
50 + token_idx = tl.program_id(0)
51 + num_tokens = tl.num_programs(0)
52 + head_idx = tl.program_id(1)
53 + num_heads = tl.num_programs(1)
54 +
55 + p_lse = tl.load(prefix_lse + head_idx * num_tokens + token_idx)
56 + s_lse = tl.load(suffix_lse + head_idx * num_tokens + token_idx)
57 + max_lse = tl.maximum(p_lse, s_lse)
58 + p_lse = p_lse - max_lse
59 + s_lse = s_lse - max_lse
60 + out_se = (tl.exp(p_lse) + tl.exp(s_lse))
61 +
62 + if OUTPUT_LSE:
63 + out_lse = tl.log(out_se) + max_lse
64 + tl.store(output_lse + head_idx * num_tokens + token_idx, out_lse)
65 +
66 + head_arange = tl.arange(0, PADDED_HEAD_SIZE)
67 + head_mask = head_arange < HEAD_SIZE
68 + p_out = tl.load(prefix_output + token_idx * num_heads * HEAD_SIZE +
69 + head_idx * HEAD_SIZE + head_arange,
70 + mask=head_mask)
71 + s_out = tl.load(suffix_output + token_idx * num_heads * HEAD_SIZE +
72 + head_idx * HEAD_SIZE + head_arange,
73 + mask=head_mask)
74 +
75 + # NOTE(woosuk): Be careful with the numerical stability.
76 + # We should compute the scale first, and then multiply it with the output.
77 + # Do not multiply the output with tl.exp(p_lse) or tl.exp(s_lse) directly.
78 + p_scale = tl.exp(p_lse) / out_se
79 + s_scale = tl.exp(s_lse) / out_se
80 + out = p_out * p_scale + s_out * s_scale
81 + tl.store(output + token_idx * num_heads * HEAD_SIZE +
82 + head_idx * HEAD_SIZE + head_arange,
83 + out,
84 + mask=head_mask)
vllm/config.py
-13
@@ -3332,19 +3332,6 @@ class VllmConfig:
3332
3333 current_platform.check_and_update_config(self)
3334
3335 - # If MLA is enabled, force disable chunked prefill and prefix caching
3336 - if self.model_config and self.model_config.use_mla:
3337 - logger.info("MLA is enabled; forcing chunked prefill and prefix "
3338 - "caching to be disabled.")
3339 - self.scheduler_config.enable_chunked_prefill = False
3340 - self.scheduler_config.chunked_prefill_enabled = False
3341 - self.scheduler_config.max_num_batched_tokens = max(
3342 - self.scheduler_config.max_model_len,
3343 - _DEFAULT_MAX_NUM_BATCHED_TOKENS)
3344 -
3345 - if self.cache_config is not None:
3346 - self.cache_config.enable_prefix_caching = False
3347 -
3335 if not self.instance_id:
3336 self.instance_id = random_uuid()[:5]
3337
vllm/engine/arg_utils.py
+3 -4
@@ -1170,9 +1170,9 @@ class EngineArgs:
1170 # long context (> 32K) models. This is to avoid OOM errors in the
1171 # initial memory profiling phase.
1172
1173 - # For multimodal models, chunked prefill is disabled by default in
1174 - # V0, but enabled by design in V1
1175 - if model_config.is_multimodal_model:
1173 + # For multimodal models and models with MLA, chunked prefill is
1174 + # disabled by default in V0, but enabled by design in V1
1175 + if model_config.is_multimodal_model or model_config.use_mla:
1176 self.enable_chunked_prefill = bool(envs.VLLM_USE_V1)
1177
1178 elif use_long_context:
@@ -1207,7 +1207,6 @@ class EngineArgs:
1207 msg = "Chunked prefill is not supported for pooling models"
1208 raise ValueError(msg)
1209
1210 -
1210 speculative_config = SpeculativeConfig.maybe_create_spec_config(
1211 target_model_config=model_config,
1212 target_parallel_config=parallel_config,
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+20 -3
@@ -162,6 +162,9 @@ def _per_token_group_quant_fp8(
162 y_q_ptr,
163 y_s_ptr,
164 group_size,
165 + # Num columns of y
166 + y_num_columns,
167 + y_row_stride,
168 # Avoid to divide zero
169 eps,
170 # Information for float8
@@ -174,9 +177,14 @@ def _per_token_group_quant_fp8(
177 quantization on a tensor.
178 This function converts the tensor values into float8 values.
179 """
180 + groups_per_row = y_num_columns // group_size
181 +
182 # Map the program id to the row of X and Y it should compute.
183 g_id = tl.program_id(0)
179 - y_ptr += g_id * group_size
184 + row = g_id // groups_per_row
185 + row_g_id = g_id % groups_per_row
186 +
187 + y_ptr += (row * y_row_stride) + (row_g_id * group_size)
188 y_q_ptr += g_id * group_size
189 y_s_ptr += g_id
190
@@ -202,6 +210,7 @@ def _per_token_group_quant_fp8_colmajor(
210 group_size,
211 # Num columns of y
212 y_num_columns,
213 + y_row_stride,
214 # Stride from one column to the next of y_s
215 y_s_col_stride,
216 # Avoid to divide zero
@@ -216,9 +225,14 @@ def _per_token_group_quant_fp8_colmajor(
225 quantization on a tensor.
226 This function converts the tensor values into float8 values.
227 """
228 + groups_per_row = y_num_columns // group_size
229 +
230 # Map the program id to the row of X and Y it should compute.
231 g_id = tl.program_id(0)
221 - y_ptr += g_id * group_size
232 + row = g_id // groups_per_row
233 + row_g_id = g_id % groups_per_row
234 +
235 + y_ptr += (row * y_row_stride) + (row_g_id * group_size)
236 y_q_ptr += g_id * group_size
237
238 # Convert g_id the flattened block coordinate to 2D so we can index
@@ -267,7 +281,7 @@ def per_token_group_quant_fp8(
281 assert (x.shape[-1] % group_size == 0), (
282 f"the last dimension of `x` {x.shape[-1]} must be divisible "
283 f"by `group_size` {group_size}")
270 - assert x.is_contiguous(), "`x` must be contiguous"
284 + assert x.stride(-1) == 1, "`x` groups must be contiguous"
285
286 finfo = torch.finfo(dtype)
287 fp8_min = finfo.min
@@ -295,6 +309,7 @@ def per_token_group_quant_fp8(
309 x_s,
310 group_size,
311 x.shape[1],
312 + x.stride(0),
313 x_s.stride(1),
314 eps,
315 fp8_min=fp8_min,
@@ -309,6 +324,8 @@ def per_token_group_quant_fp8(
324 x_q,
325 x_s,
326 group_size,
327 + x.shape[1],
328 + x.stride(0),
329 eps,
330 fp8_min=fp8_min,
331 fp8_max=fp8_max,
vllm/utils.py
+4
@@ -565,6 +565,10 @@ def round_up(x: int, y: int) -> int:
565 return ((x + y - 1) // y) * y
566
567
568 +def round_down(x: int, y: int) -> int:
569 + return (x // y) * y
570 +
571 +
572 def _generate_random_fp8(
573 tensor: torch.Tensor,
574 low: float,
vllm/v1/attention/backends/flash_attn.py
+2 -69
@@ -5,12 +5,11 @@ from typing import Any, Dict, List, Optional, Tuple, Type
5
6 import numpy as np
7 import torch
8 -import triton
9 -import triton.language as tl
8
9 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
10 AttentionMetadata, AttentionType)
11 from vllm.attention.backends.utils import get_flash_attn_version
12 +from vllm.attention.ops.triton_merge_attn_states import merge_attn_states
13 from vllm.logger import init_logger
14 from vllm.platforms import current_platform
15 from vllm.utils import cdiv
@@ -372,70 +371,4 @@ def cascade_attention(
371
372 # Merge prefix and suffix outputs, and store the result in output.
373 merge_attn_states(output, prefix_output, prefix_lse, suffix_output,
375 - suffix_lse)
376 -
377 -
378 -def merge_attn_states(
379 - output: torch.Tensor,
380 - prefix_output: torch.Tensor,
381 - prefix_lse: torch.Tensor,
382 - suffix_output: torch.Tensor,
383 - suffix_lse: torch.Tensor,
384 -) -> None:
385 - num_tokens = output.shape[0]
386 - num_query_heads = output.shape[1]
387 - head_size = output.shape[2]
388 - padded_head_size = triton.next_power_of_2(head_size)
389 -
390 - # TODO(woosuk): Use CUDA kernel instead of Triton to minimize CPU overhead.
391 - merge_attn_states_kernel[(num_tokens, num_query_heads)](
392 - output,
393 - prefix_output,
394 - prefix_lse,
395 - suffix_output,
396 - suffix_lse,
397 - head_size,
398 - padded_head_size,
399 - )
400 -
401 -
402 -@triton.jit
403 -def merge_attn_states_kernel(
404 - output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
405 - prefix_output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
406 - prefix_lse, # [NUM_HEADS, NUM_TOKENS]
407 - suffix_output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
408 - suffix_lse, # [NUM_HEADS, NUM_TOKENS]
409 - HEAD_SIZE: tl.constexpr,
410 - PADDED_HEAD_SIZE: tl.constexpr,
411 -):
412 - token_idx = tl.program_id(0)
413 - num_tokens = tl.num_programs(0)
414 - head_idx = tl.program_id(1)
415 - num_heads = tl.num_programs(1)
416 -
417 - p_lse = tl.load(prefix_lse + head_idx * num_tokens + token_idx)
418 - s_lse = tl.load(suffix_lse + head_idx * num_tokens + token_idx)
419 - max_lse = tl.maximum(p_lse, s_lse)
420 - p_lse = p_lse - max_lse
421 - s_lse = s_lse - max_lse
422 -
423 - head_arange = tl.arange(0, PADDED_HEAD_SIZE)
424 - head_mask = head_arange < HEAD_SIZE
425 - p_out = tl.load(prefix_output + token_idx * num_heads * HEAD_SIZE +
426 - head_idx * HEAD_SIZE + head_arange,
427 - mask=head_mask)
428 - s_out = tl.load(suffix_output + token_idx * num_heads * HEAD_SIZE +
429 - head_idx * HEAD_SIZE + head_arange,
430 - mask=head_mask)
431 -
432 - # NOTE(woosuk): Be careful with the numerical stability.
433 - # We should compute the scale first, and then multiply it with the output.
434 - # Do not multiply the output with tl.exp(p_lse) or tl.exp(s_lse) directly.
435 - p_scale = tl.exp(p_lse) / (tl.exp(p_lse) + tl.exp(s_lse))
436 - s_scale = tl.exp(s_lse) / (tl.exp(p_lse) + tl.exp(s_lse))
437 - out = p_out * p_scale + s_out * s_scale
438 - tl.store(output + token_idx * num_heads * HEAD_SIZE +
439 - head_idx * HEAD_SIZE + head_arange,
440 - out,
441 - mask=head_mask)
374 + suffix_lse)
\ No newline at end of file