@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 19fe1a0510

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 14, 05:30 UTC

[Kernel] Add FP8 support with FlashMLA backend (#22668)

Signed-off-by: Matthew Bonanni <mbonanni001@gmail.com>

Matthew Bonanni committed Aug 21, 2025 at 22:26 UTC 19fe1a0510852b5a892932653a5881e58e359660

19 files changed +235 -109

cmake/external_projects/flashmla.cmake

+5 -4

@@ -19,7 +19,7 @@ else()

19 FetchContent_Declare(

20 flashmla

21 GIT_REPOSITORY https://github.com/vllm-project/FlashMLA.git

22 - GIT_TAG 0e43e774597682284358ff2c54530757b654b8d1

22 + GIT_TAG a757314c04eedd166e329e846c820eb1bdd702de

23 GIT_PROGRESS TRUE

24 CONFIGURE_COMMAND ""

25 BUILD_COMMAND ""

@@ -37,13 +37,14 @@ cuda_archs_loose_intersection(FLASH_MLA_ARCHS "9.0a" "${CUDA_ARCHS}")

37 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.3 AND FLASH_MLA_ARCHS)

38 set(FlashMLA_SOURCES

39 ${flashmla_SOURCE_DIR}/csrc/flash_api.cpp

40 - ${flashmla_SOURCE_DIR}/csrc/kernels/splitkv_mla.cu

40 + ${flashmla_SOURCE_DIR}/csrc/kernels/get_mla_metadata.cu

41 ${flashmla_SOURCE_DIR}/csrc/kernels/mla_combine.cu

42 - ${flashmla_SOURCE_DIR}/csrc/kernels/get_mla_metadata.cu)

42 + ${flashmla_SOURCE_DIR}/csrc/kernels/splitkv_mla.cu

43 + ${flashmla_SOURCE_DIR}/csrc/kernels_fp8/flash_fwd_mla_fp8_sm90.cu)

45 set(FlashMLA_INCLUDES

46 ${flashmla_SOURCE_DIR}/csrc/cutlass/include

46 - ${flashmla_SOURCE_DIR}/csrc/include)

47 + ${flashmla_SOURCE_DIR}/csrc)

49 set_gencode_flags_for_srcs(

50 SRCS "${FlashMLA_SOURCES}"

csrc/cache.h

+4 -2

@@ -40,9 +40,11 @@ void concat_and_cache_mla(torch::Tensor& kv_c, torch::Tensor& k_pe,

40 void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,

41 const double scale, const std::string& kv_cache_dtype);

43 -void gather_cache(

43 +void gather_and_maybe_dequant_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

48 + int64_t batch_size, const std::string& kv_cache_dtype,

49 + torch::Tensor const& scale,

50 + std::optional<torch::Tensor> seq_starts = std::nullopt);

\ No newline at end of file

csrc/cache_kernels.cu

+29 -28

@@ -624,9 +624,9 @@ void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,

624 namespace vllm {

625

626 // grid is launched with dimensions (batch, num_splits)

627 -template <typename scalar_t>

628 -__global__ void gather_cache(

629 - const scalar_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE,

627 +template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>

628 +__global__ void gather_and_maybe_dequant_cache(

629 + const cache_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE,

630 // ENTRIES...]

631 scalar_t* __restrict__ dst, // [TOT_TOKENS, ENTRIES...]

632 const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES]

@@ -634,6 +634,7 @@ __global__ void gather_cache(

634 const int32_t block_size, const int32_t entry_size,

635 const int64_t block_table_stride, const int64_t cache_block_stride,

636 const int64_t cache_entry_stride, const int64_t dst_entry_stride,

637 + const float* __restrict__ scale,

638 const int32_t* __restrict__ seq_starts) { // Optional: starting offsets per

639 // batch

640

@@ -675,10 +676,16 @@ __global__ void gather_cache(

676 if (partial_block_size) full_blocks_end -= 1;

677 }

678

678 - auto copy_entry = [&](const scalar_t* __restrict__ _src,

679 + auto copy_entry = [&](const cache_t* __restrict__ _src,

680 scalar_t* __restrict__ _dst) {

680 - for (int i = threadIdx.x; i < entry_size; i += blockDim.x)

681 - _dst[i] = _src[i];

681 + for (int i = threadIdx.x; i < entry_size; i += blockDim.x) {

682 + if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {

683 + _dst[i] = static_cast<scalar_t>(_src[i]);

684 + } else {

685 + _dst[i] =

686 + fp8::scaled_convert<scalar_t, cache_t, kv_dt>(_src[i], *scale);

687 + }

688 + }

689 };

690

691 for (int pid = split_start; pid < full_blocks_end; ++pid) {

@@ -705,25 +712,31 @@ __global__ void gather_cache(

712 } // namespace vllm

713

714 // Macro to dispatch the kernel based on the data type.

708 -#define CALL_GATHER_CACHE(CPY_DTYPE) \

709 - vllm::gather_cache<CPY_DTYPE><<<grid, block, 0, stream>>>( \

710 - reinterpret_cast<CPY_DTYPE*>(src_cache.data_ptr()), \

711 - reinterpret_cast<CPY_DTYPE*>(dst.data_ptr()), \

712 - block_table.data_ptr<int32_t>(), cu_seq_lens.data_ptr<int32_t>(), \

713 - block_size, entry_size, block_table_stride, cache_block_stride, \

714 - cache_entry_stride, dst_entry_stride, seq_starts_ptr);

715 +// SCALAR_T is the data type of the destination tensor.

716 +// CACHE_T is the stored data type of kv-cache.

717 +// KV_DTYPE is the real data type of kv-cache.

718 +#define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE) \

719 + vllm::gather_and_maybe_dequant_cache<SCALAR_T, CACHE_T, KV_DTYPE> \

720 + <<<grid, block, 0, stream>>>( \

721 + reinterpret_cast<CACHE_T*>(src_cache.data_ptr()), \

722 + reinterpret_cast<SCALAR_T*>(dst.data_ptr()), \

723 + block_table.data_ptr<int32_t>(), cu_seq_lens.data_ptr<int32_t>(), \

724 + block_size, entry_size, block_table_stride, cache_block_stride, \

725 + cache_entry_stride, dst_entry_stride, \

726 + reinterpret_cast<const float*>(scale.data_ptr()), seq_starts_ptr);

727

728 // Gather sequences from the cache into the destination tensor.

729 // - cu_seq_lens contains the cumulative sequence lengths for each batch

730 // - block_table contains the cache block indices for each sequence

731 // - Optionally, seq_starts (if provided) offsets the starting block index by

732 // (seq_starts[bid] / page_size)

721 -void gather_cache(

733 +void gather_and_maybe_dequant_cache(

734 torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]

735 torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]

736 torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]

737 torch::Tensor const& cu_seq_lens, // [BATCH+1]

726 - int64_t batch_size,

738 + int64_t batch_size, const std::string& kv_cache_dtype,

739 + torch::Tensor const& scale,

740 std::optional<torch::Tensor> seq_starts = std::nullopt) {

741 at::cuda::OptionalCUDAGuard device_guard(src_cache.device());

742 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

@@ -761,20 +774,8 @@ void gather_cache(

774 dim3 grid(batch_size, num_splits);

775 dim3 block(1024);

776

764 - TORCH_CHECK(src_cache.dtype() == dst.dtype(),

765 - "src_cache and dst must have the same dtype");

766 -

767 - const int dtype_bits = src_cache.element_size() * 8;

777 const int32_t* seq_starts_ptr =

778 seq_starts.has_value() ? seq_starts.value().data_ptr<int32_t>() : nullptr;

779

771 - if (dtype_bits == 32) {

772 - CALL_GATHER_CACHE(uint32_t);

773 - } else if (dtype_bits == 16) {

774 - CALL_GATHER_CACHE(uint16_t);

775 - } else if (dtype_bits == 8) {

776 - CALL_GATHER_CACHE(uint8_t);

777 - } else {

778 - TORCH_CHECK(false, "Unsupported data type width: ", dtype_bits);

779 - }

780 + DISPATCH_BY_KV_CACHE_DTYPE(dst.dtype(), kv_cache_dtype, CALL_GATHER_CACHE);

781 }

csrc/torch_bindings.cpp

+9 -4

@@ -672,11 +672,16 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {

672 "str kv_cache_dtype) -> ()");

673 cache_ops.impl("convert_fp8", torch::kCUDA, &convert_fp8);

674

675 - // Gather cache blocks from src_cache to dst.

675 + // Gather cache blocks from src_cache to dst, dequantizing from

676 + // src_cache's dtype to dst's dtype if necessary.

677 cache_ops.def(

677 - "gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, "

678 - "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()");

679 - cache_ops.impl("gather_cache", torch::kCUDA, &gather_cache);

678 + "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, "

679 + " Tensor block_table, Tensor cu_seq_lens, "

680 + " int batch_size, "

681 + " str kv_cache_dtype, "

682 + " Tensor scale, Tensor? seq_starts) -> ()");

683 + cache_ops.impl("gather_and_maybe_dequant_cache", torch::kCUDA,

684 + &gather_and_maybe_dequant_cache);

685 }

686

687 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {

tests/kernels/attention/test_cache.py

+29 -13

@@ -709,14 +709,15 @@ def test_swap_blocks_mla(

709 @pytest.mark.parametrize("max_seq_len", [512])

710 @pytest.mark.parametrize("batch_size", [8])

711 @pytest.mark.parametrize("dtype", [torch.float32])

712 -@pytest.mark.parametrize("kv_cache_dtype",

713 - ["auto"]) # You can also test "fp8" if needed.

712 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])

713 @pytest.mark.parametrize("device", CUDA_DEVICES)

714 @torch.inference_mode()

716 -def test_gather_cache_mla(kv_lora_rank, qk_rope_head_dim, block_size,

717 - num_blocks, max_seq_len, batch_size, dtype,

718 - kv_cache_dtype, device):

715 +def test_gather_and_maybe_dequant_cache_mla(kv_lora_rank, qk_rope_head_dim,

716 + block_size, num_blocks,

717 + max_seq_len, batch_size, dtype,

718 + kv_cache_dtype, device):

719 entry_size = kv_lora_rank + qk_rope_head_dim

720 + scale = torch.tensor(0.1, dtype=torch.float32, device=device)

721 src_cache = _create_mla_cache(num_blocks, block_size, entry_size, dtype,

722 kv_cache_dtype, device)

723 _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype)

@@ -742,9 +743,7 @@ def test_gather_cache_mla(kv_lora_rank, qk_rope_head_dim, block_size,

743 perm = torch.randperm(num_blocks, device=device)

744 block_table[b, :] = perm

745

745 - dst = torch.zeros((total_tokens, entry_size),

746 - dtype=src_cache.dtype,

747 - device=device)

746 + dst = torch.zeros((total_tokens, entry_size), dtype=dtype, device=device)

747

748 expected_batches = []

749 for b in range(batch_size):

@@ -756,21 +755,38 @@ def test_gather_cache_mla(kv_lora_rank, qk_rope_head_dim, block_size,

755

756 gathered_rows = []

757 for i in range(tot - 1):

759 - gathered_rows.append(src_cache[blocks[i]])

758 + block_data = src_cache[blocks[i]]

759 + if kv_cache_dtype == "fp8":

760 + dequantized_block = torch.empty_like(block_data, dtype=dtype)

761 + ops.convert_fp8(dequantized_block, block_data, scale.item())

762 + gathered_rows.append(dequantized_block)

763 + else:

764 + gathered_rows.append(block_data)

765 remaining = s - (tot - 1) * block_size

761 - gathered_rows.append(src_cache[blocks[-1], :remaining, :])

766 + last_block_data = src_cache[blocks[-1], :remaining, :]

767 + if kv_cache_dtype == "fp8":

768 + dequantized_last_block = torch.empty_like(last_block_data,

769 + dtype=dtype)

770 + ops.convert_fp8(dequantized_last_block, last_block_data,

771 + scale.item())

772 + gathered_rows.append(dequantized_last_block)

773 + else:

774 + gathered_rows.append(last_block_data)

775

776 batch_expected = torch.cat(gathered_rows, dim=0)

777 expected_batches.append(batch_expected)

778 expected = torch.cat(expected_batches, dim=0)

779

780 opcheck(

768 - torch.ops._C_cache_ops.gather_cache,

769 - (src_cache, dst, block_table, cu_seq_lens, batch_size, None),

781 + torch.ops._C_cache_ops.gather_and_maybe_dequant_cache,

782 + (src_cache, dst, block_table, cu_seq_lens, batch_size, kv_cache_dtype,

783 + scale, None),

784 test_utils=DEFAULT_OPCHECK_TEST_UTILS,

785 )

786

773 - ops.gather_cache(src_cache, dst, block_table, cu_seq_lens, batch_size)

787 + ops.gather_and_maybe_dequant_cache(src_cache, dst, block_table,

788 + cu_seq_lens, batch_size, kv_cache_dtype,

789 + scale, None)

790 torch.testing.assert_close(dst, expected)

791

792

tests/kernels/attention/test_flashmla.py

+51 -18

@@ -13,11 +13,17 @@ from vllm.attention.ops.flashmla import (flash_mla_with_kvcache,

13 from vllm.triton_utils import triton

16 -def cal_diff(x: torch.Tensor, y: torch.Tensor, name: str) -> None:

16 +def cal_diff(x: torch.Tensor,

17 + y: torch.Tensor,

18 + name: str,

19 + use_fp8: bool = False) -> None:

20 x, y = x.double(), y.double()

21 cos_diff = 1 - 2 * (x * y).sum().item() / max(

22 (x * x + y * y).sum().item(), 1e-12)

20 - assert cos_diff < 1e-5

23 + if (use_fp8):

24 + assert cos_diff < 1e-4

25 + else:

26 + assert cos_diff < 1e-5

28 FLASH_MLA_UNSUPPORTED_REASON = is_flashmla_supported()[1] \

29 if not is_flashmla_supported()[0] else "FlashMLA is supported"

@@ -27,7 +33,7 @@ FLASH_MLA_UNSUPPORTED_REASON = is_flashmla_supported()[1] \

33 reason=FLASH_MLA_UNSUPPORTED_REASON)

34 @pytest.mark.parametrize("b", [128])

35 @pytest.mark.parametrize("s_q", [1, 2])

30 -@pytest.mark.parametrize("mean_sk", [4096, 8192])

36 +@pytest.mark.parametrize("mean_sk", [4096, 8192, 16384])

37 @pytest.mark.parametrize("h_q", [16, 32, 64, 128])

38 @pytest.mark.parametrize("h_kv", [1])

39 @pytest.mark.parametrize("d", [576])

@@ -35,20 +41,26 @@ FLASH_MLA_UNSUPPORTED_REASON = is_flashmla_supported()[1] \

41 @pytest.mark.parametrize("block_size", [64])

42 @pytest.mark.parametrize("causal", [True])

43 @pytest.mark.parametrize("varlen", [False, True])

38 -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])

44 +@pytest.mark.parametrize("torch_dtype",

45 + [torch.bfloat16, torch.float16, torch.float8_e4m3fn])

46 @torch.inference_mode()

47 def test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal,

41 - varlen, dtype):

48 + varlen, torch_dtype):

49 device = torch.device("cuda:0")

43 - torch.set_default_dtype(dtype)

50 + if torch_dtype == torch.float8_e4m3fn:

51 + init_dtype = torch.bfloat16

52 + else:

53 + init_dtype = torch_dtype

54 + torch.set_default_dtype(init_dtype)

55 torch.set_default_device(device)

56 torch.cuda.set_device(device)

57 torch.manual_seed(0)

58 random.seed(0)

60 print(f"{b=}, {s_q=}, {mean_sk=}, {h_q=}, {h_kv=}, "

50 - f"{d=}, {dv=}, {causal=}, {varlen=}, {dtype=}")

61 + f"{d=}, {dv=}, {causal=}, {varlen=}, {torch_dtype=}")

63 + use_fp8 = torch_dtype == torch.float8_e4m3fn

64 cache_seqlens = torch.full((b, ), mean_sk, dtype=torch.int32)

65 if varlen:

66 for i in range(b):

@@ -71,6 +83,19 @@ def test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal,

83 tile_scheduler_metadata, num_splits = get_mla_metadata(

84 cache_seqlens, s_q * h_q // h_kv, h_kv)

86 + init_dtype = q.dtype

87 + if use_fp8:

88 + fp8_dtype = torch.float8_e4m3fn

89 + descale_q = torch.ones((1), dtype=torch.float32)

90 + descale_k = torch.ones((1), dtype=torch.float32)

91 +

92 + q = q.to(fp8_dtype)

93 + blocked_k = blocked_k.to(fp8_dtype)

94 + blocked_v = blocked_v.to(fp8_dtype)

95 + else:

96 + descale_q = None

97 + descale_k = None

98 +

99 def flash_mla():

100 return flash_mla_with_kvcache(

101 q,

@@ -81,6 +106,8 @@ def test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal,

106 tile_scheduler_metadata,

107 num_splits,

108 causal=causal,

109 + descale_q=descale_q,

110 + descale_k=descale_k,

111 )

112

113 def scaled_dot_product_attention(query, key, value, is_causal=False):

@@ -104,29 +131,35 @@ def test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal,

131 return attn_weight @ value, lse

132

133 def ref_mla():

134 + q_ = (q.to(torch.float) * descale_q).to(init_dtype) if use_fp8 else q

135 + blocked_k_ = (blocked_k.to(torch.float) *

136 + descale_k).to(init_dtype) if use_fp8 else blocked_k

137 + blocked_v_ = (blocked_v.to(torch.float) *

138 + descale_k).to(init_dtype) if use_fp8 else blocked_v

139 out = torch.empty(b, s_q, h_q, dv, dtype=torch.float32)

140 lse = torch.empty(b, h_q, s_q, dtype=torch.float32)

141 for i in range(b):

142 begin = i * max_seqlen_pad

143 end = begin + cache_seqlens[i]

112 - ref_O, LSE = scaled_dot_product_attention(

113 - q[i].transpose(0, 1),

114 - blocked_k.view(-1, h_kv, d)[begin:end].transpose(0, 1),

115 - blocked_v.view(-1, h_kv, dv)[begin:end].transpose(0, 1),

144 + out_i, lse_i = scaled_dot_product_attention(

145 + q_[i].transpose(0, 1),

146 + blocked_k_.view(-1, h_kv, d)[begin:end].transpose(0, 1),

147 + blocked_v_.view(-1, h_kv, dv)[begin:end].transpose(0, 1),

148 is_causal=causal,

149 )

118 - out[i] = ref_O.transpose(0, 1)

119 - lse[i] = LSE

150 + out[i] = out_i.transpose(0, 1)

151 + lse[i] = lse_i

152 return out, lse

153

154 out_flash, lse_flash = flash_mla()

155 out_torch, lse_torch = ref_mla()

124 - cal_diff(out_flash, out_torch, "out")

156 + cal_diff(out_flash, out_torch, "out", use_fp8)

157 cal_diff(lse_flash, lse_torch, "lse")

158

159 t = triton.testing.do_bench(flash_mla)

160 FLOPS = s_q * total_seqlens * h_q * (d + dv) * 2

129 - bytes = (total_seqlens * h_kv * d + b * s_q * h_q * d +

130 - b * s_q * h_q * dv) * (torch.finfo(dtype).bits // 8)

131 - print(f"{t:.3f} ms, {FLOPS / 10 ** 9 / t:.0f} "

132 - f"TFLOPS, {bytes / 10 ** 6 / t:.0f} GB/s")

161 + bytes = (total_seqlens * h_kv * d +

162 + b * s_q * h_q * d) * (torch.finfo(torch_dtype).bits // 8) + (

163 + b * s_q * h_q * dv) * (torch.finfo(init_dtype).bits // 8)

164 + print(f"{t:.3f} ms, {FLOPS / 10 ** 9 / t:.0f} TFLOPS,",

165 + f"{bytes / 10 ** 6 / t:.0f} GB/s")

vllm/_custom_ops.py

+12 -8

@@ -1589,14 +1589,18 @@ def convert_fp8(output: torch.Tensor,

1589 torch.ops._C_cache_ops.convert_fp8(output, input, scale, kv_dtype)

1590

1591

1592 -def gather_cache(src_cache: torch.Tensor,

1593 - dst: torch.Tensor,

1594 - block_table: torch.Tensor,

1595 - cu_seq_lens: torch.Tensor,

1596 - batch_size: int,

1597 - seq_starts: Optional[torch.Tensor] = None) -> None:

1598 - torch.ops._C_cache_ops.gather_cache(src_cache, dst, block_table,

1599 - cu_seq_lens, batch_size, seq_starts)

1592 +def gather_and_maybe_dequant_cache(

1593 + src_cache: torch.Tensor,

1594 + dst: torch.Tensor,

1595 + block_table: torch.Tensor,

1596 + cu_seq_lens: torch.Tensor,

1597 + batch_size: int,

1598 + kv_cache_dtype: str,

1599 + scale: torch.Tensor,

1600 + seq_starts: Optional[torch.Tensor] = None) -> None:

1601 + torch.ops._C_cache_ops.gather_and_maybe_dequant_cache(

1602 + src_cache, dst, block_table, cu_seq_lens, batch_size, kv_cache_dtype,

1603 + scale, seq_starts)

1604

1605

1606 def get_device_attribute(attribute: int, device: int) -> int:

vllm/attention/backends/mla/common.py

+9 -5

@@ -837,8 +837,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):

837 self.context_chunk_workspace_size // num_prefills_with_context

838

839 # align max_context_chunk to page_size by rounding down,

840 - # currently the `gather_cache` kernel cannot handle

841 - # `context_chunk_starts` that are not aligned to page_size

840 + # currently the `gather_and_maybe_dequant_cache` kernel cannot

841 + # handle `context_chunk_starts` that are not aligned to page_size

842 max_context_chunk = round_down(max_context_chunk, self.page_size)

843 assert max_context_chunk > 0

844 num_chunks = cdiv(context_lens_tensor.max(), max_context_chunk)

@@ -1082,6 +1082,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1082 q: torch.Tensor,

1083 kv_c_and_k_pe_cache: torch.Tensor,

1084 attn_metadata: MLACommonMetadata,

1085 + k_scale: torch.Tensor,

1086 ):

1087 prefill_metadata = attn_metadata.prefill_metadata

1088 assert prefill_metadata is not None

@@ -1103,12 +1104,14 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1104 for i in range(iters):

1105 toks = prefill_metadata.context_chunk_seq_tot[i]

1106

1106 - ops.gather_cache(

1107 + ops.gather_and_maybe_dequant_cache(

1108 src_cache=kv_c_and_k_pe_cache,

1109 dst=workspace,

1110 block_table=prefill_metadata.block_tables,

1111 cu_seq_lens=prefill_metadata.context_chunk_cu_seq_lens[i],

1112 batch_size=prefill_metadata.num_prefills,

1113 + kv_cache_dtype=self.kv_cache_dtype,

1114 + scale=k_scale,

1115 seq_starts=prefill_metadata.context_chunk_starts[i],

1116 )

1117

@@ -1165,6 +1168,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1168 k_pe: torch.Tensor,

1169 kv_c_and_k_pe_cache: torch.Tensor,

1170 attn_metadata: MLACommonMetadata,

1171 + k_scale: torch.Tensor,

1172 ) -> torch.Tensor:

1173

1174 prefill_metadata = attn_metadata.prefill_metadata

@@ -1197,7 +1201,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1201 # ROCm flash_attn_varlen_func will return 3 objects instead of 2

1202 suffix_output, suffix_lse = output

1203 context_output, context_lse = self._compute_prefill_context( \

1200 - q, kv_c_and_k_pe_cache, attn_metadata)

1204 + q, kv_c_and_k_pe_cache, attn_metadata, k_scale)

1205

1206 output = torch.empty_like(suffix_output)

1207 merge_attn_states(

@@ -1287,7 +1291,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1291 if has_prefill:

1292 output[:num_prefill_tokens] = self._forward_prefill(

1293 prefill_q, prefill_k_c_normed, prefill_k_pe, kv_cache,

1290 - attn_metadata)

1294 + attn_metadata, layer._k_scale)

1295

1296 if has_decode:

1297 decode_q_nope, decode_q_pe = decode_q.split(

vllm/attention/ops/flashmla.py

@@ -67,6 +67,8 @@ def flash_mla_with_kvcache(

67 num_splits: torch.Tensor,

68 softmax_scale: Optional[float] = None,

69 causal: bool = False,

70 + descale_q: Optional[torch.Tensor] = None,

71 + descale_k: Optional[torch.Tensor] = None,

72 ) -> Tuple[torch.Tensor, torch.Tensor]:

73 """

74 Arguments:

@@ -81,6 +83,8 @@ def flash_mla_with_kvcache(

83 softmax_scale: float. The scaling of QK^T before applying softmax.

84 Default to 1 / sqrt(head_dim).

85 causal: bool. Whether to apply causal attention mask.

86 + descale_q: (batch_size), torch.float32. Descaling factors for Q.

87 + descale_k: (batch_size), torch.float32. Descaling factors for K.

89 Return:

90 out: (batch_size, seq_len_q, num_heads_q, head_dim_v).

@@ -98,6 +102,8 @@ def flash_mla_with_kvcache(

102 causal,

103 tile_scheduler_metadata,

104 num_splits,

105 + descale_q,

106 + descale_k,

107 )

108 return out, softmax_lse

109

vllm/engine/arg_utils.py

+1 -2

@@ -1445,10 +1445,9 @@ class EngineArgs:

1445 recommend_to_remove=False)

1446 return False

1447

1448 - # No Fp8 KV cache so far.

1448 if self.kv_cache_dtype != "auto":

1449 supported = current_platform.is_kv_cache_dtype_supported(

1451 - self.kv_cache_dtype)

1450 + self.kv_cache_dtype, model_config)

1451 if not supported:

1452 _raise_or_fallback(feature_name="--kv-cache-dtype",

1453 recommend_to_remove=False)

vllm/platforms/cuda.py

+33 -8

@@ -481,16 +481,41 @@ class CudaPlatformBase(Platform):

481 return cuda_device_count_stateless()

482

483 @classmethod

484 - def is_kv_cache_dtype_supported(cls, kv_cache_dtype: str) -> bool:

484 + def is_kv_cache_dtype_supported(cls, kv_cache_dtype: str,

485 + model_config: "ModelConfig") -> bool:

486 fp8_attention = kv_cache_dtype.startswith("fp8")

486 - will_use_fa = (not envs.is_set("VLLM_ATTENTION_BACKEND")

487 - ) or envs.VLLM_ATTENTION_BACKEND == "FLASH_ATTN_VLLM_V1"

487 + attention_backend = envs.VLLM_ATTENTION_BACKEND

488 +

489 supported = False

489 - if cls.is_device_capability(100):

490 - supported = True

491 - elif fp8_attention and will_use_fa:

492 - from vllm.attention.utils.fa_utils import flash_attn_supports_fp8

493 - supported = flash_attn_supports_fp8()

490 + if model_config is not None and model_config.use_mla:

491 + # Default to CutlassMLA for blackwell,

492 + # FlashMLA otherwise

493 + if attention_backend is None:

494 + if cls.is_device_capability(100):

495 + attention_backend = "CUTLASS_MLA"

496 + else:

497 + attention_backend = "FLASHMLA"

498 +

499 + # Only FlashMLA supports fp8

500 + if attention_backend == "FLASHMLA":

501 + supported = True

502 + else:

503 + supported = (not fp8_attention)

504 + else:

505 + # Default to FlashAttention

506 + if attention_backend is None:

507 + attention_backend = "FLASH_ATTN_VLLM_V1"

508 +

509 + # All Blackwell backends support fp8

510 + if cls.is_device_capability(100):

511 + supported = True

512 + elif attention_backend == "FLASH_ATTN_VLLM_V1":

513 + if fp8_attention:

514 + from vllm.attention.utils.fa_utils import (

515 + flash_attn_supports_fp8)

516 + supported = flash_attn_supports_fp8()

517 + else:

518 + supported = True

519 return supported

520

521

vllm/platforms/interface.py

+2 -1

@@ -565,7 +565,8 @@ class Platform:

565 raise RuntimeError(f"Unsupported torch distributed backend: {backend}")

566

567 @classmethod

568 - def is_kv_cache_dtype_supported(cls, kv_cache_dtype: str) -> bool:

568 + def is_kv_cache_dtype_supported(cls, kv_cache_dtype: str,

569 + model_config: "ModelConfig") -> bool:

570 """

571 Returns if the kv_cache_dtype is supported by the current platform.

572 """

vllm/platforms/rocm.py

+2 -1

@@ -459,5 +459,6 @@ class RocmPlatform(Platform):

459 return cuda_device_count_stateless()

460

461 @classmethod

462 - def is_kv_cache_dtype_supported(cls, kv_cache_dtype: str) -> bool:

462 + def is_kv_cache_dtype_supported(cls, kv_cache_dtype: str,

463 + model_config: "ModelConfig") -> bool:

464 return True

vllm/platforms/tpu.py

+2 -1

@@ -196,7 +196,8 @@ class TpuPlatform(Platform):

196 raise ValueError("Torch XLA does not support per-request seed.")

197

198 @classmethod

199 - def is_kv_cache_dtype_supported(cls, kv_cache_dtype: str) -> bool:

199 + def is_kv_cache_dtype_supported(cls, kv_cache_dtype: str,

200 + model_config: "ModelConfig") -> bool:

201 return True

202

203

vllm/v1/attention/backends/mla/common.py

+31 -6

@@ -631,8 +631,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

631

632 if self.aot_schedule:

633 # align max_context_chunk to page_size by rounding down,

634 - # currently the `gather_cache` kernel cannot handle

635 - # `context_chunk_starts` that are not aligned to page_size

634 + # currently the `gather_and_maybe_dequant_cache` kernel

635 + # cannot handle `context_chunk_starts` that are not aligned

636 + # to page_size

637 max_context_chunk = round_down(max_context_chunk,

638 self.page_size)

639

@@ -1005,6 +1006,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

1006 q: torch.Tensor,

1007 kv_c_and_k_pe_cache: torch.Tensor,

1008 attn_metadata: MLACommonMetadata,

1009 + k_scale: torch.Tensor,

1010 ):

1011 assert attn_metadata.prefill is not None

1012 prefill_metadata = attn_metadata.prefill

@@ -1017,12 +1019,14 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

1019 for i in range(iters):

1020 toks = prefill_metadata.chunked_context.seq_tot[i]

1021

1020 - ops.gather_cache(

1022 + ops.gather_and_maybe_dequant_cache(

1023 src_cache=kv_c_and_k_pe_cache,

1024 dst=workspace,

1025 block_table=prefill_metadata.block_table,

1026 cu_seq_lens=prefill_metadata.chunked_context.cu_seq_lens[i],

1027 batch_size=attn_metadata.num_prefills,

1028 + kv_cache_dtype=self.kv_cache_dtype,

1029 + scale=k_scale,

1030 seq_starts=prefill_metadata.chunked_context.starts[i],

1031 )

1032

@@ -1073,6 +1077,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

1077 k_pe: torch.Tensor,

1078 kv_c_and_k_pe_cache: torch.Tensor,

1079 attn_metadata: MLACommonMetadata,

1080 + k_scale: torch.Tensor,

1081 ) -> torch.Tensor:

1082 assert attn_metadata.prefill is not None

1083

@@ -1095,7 +1100,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

1100 if has_context:

1101 suffix_output, suffix_lse = output

1102 context_output, context_lse = self._compute_prefill_context( \

1098 - q, kv_c_and_k_pe_cache, attn_metadata)

1103 + q, kv_c_and_k_pe_cache, attn_metadata, k_scale)

1104

1105 output = torch.empty_like(suffix_output)

1106 merge_attn_states(

@@ -1119,6 +1124,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

1124 q_pe: torch.Tensor,

1125 kv_c_and_k_pe_cache: torch.Tensor,

1126 attn_metadata: M,

1127 + layer: AttentionLayer,

1128 ) -> torch.Tensor:

1129 raise NotImplementedError

1130

@@ -1146,6 +1152,8 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

1152 # same expert outputs.

1153 return output.fill_(0)

1154

1155 + fp8_attention = self.kv_cache_dtype.startswith("fp8")

1156 +

1157 num_actual_toks = attn_metadata.num_actual_tokens

1158

1159 # Inputs and outputs may be padded for CUDA graphs

@@ -1180,10 +1188,13 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

1188 scale=layer._k_scale,

1189 )

1190

1191 + if fp8_attention:

1192 + kv_cache = kv_cache.view(current_platform.fp8_dtype())

1193 +

1194 if has_prefill:

1195 output[num_decode_tokens:] = self._forward_prefill(

1196 prefill_q, prefill_k_c_normed, prefill_k_pe, kv_cache,

1186 - attn_metadata)

1197 + attn_metadata, layer._k_scale)

1198

1199 if has_decode:

1200 assert attn_metadata.decode is not None

@@ -1196,7 +1207,21 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

1207 # Convert from (N, B, L) to (B, N, L)

1208 decode_ql_nope = decode_ql_nope.transpose(0, 1)

1209

1210 + if fp8_attention:

1211 + ql_nope_shape = decode_ql_nope.shape

1212 + decode_ql_nope, _ = ops.scaled_fp8_quant(

1213 + decode_ql_nope.reshape([

1214 + ql_nope_shape[0], ql_nope_shape[1] * ql_nope_shape[2]

1215 + ]), layer._q_scale)

1216 + decode_ql_nope = decode_ql_nope.reshape(ql_nope_shape)

1217 + q_pe_shape = decode_q_pe.shape

1218 + decode_q_pe, _ = ops.scaled_fp8_quant(

1219 + decode_q_pe.reshape(

1220 + [q_pe_shape[0], q_pe_shape[1] * q_pe_shape[2]]),

1221 + layer._q_scale)

1222 + decode_q_pe = decode_q_pe.reshape(q_pe_shape)

1223 +

1224 output[:num_decode_tokens] = self._forward_decode(

1200 - decode_ql_nope, decode_q_pe, kv_cache, attn_metadata)

1225 + decode_ql_nope, decode_q_pe, kv_cache, attn_metadata, layer)

1226

1227 return output_padded

vllm/v1/attention/backends/mla/cutlass_mla.py

+2 -1

@@ -7,7 +7,7 @@ from typing import ClassVar, Optional

7 import torch

9 import vllm._custom_ops as ops

10 -from vllm.attention.backends.abstract import (AttentionType,

10 +from vllm.attention.backends.abstract import (AttentionLayer, AttentionType,

11 is_quantized_kv_cache)

12 from vllm.logger import init_logger

13 from vllm.v1.attention.backends.mla.common import (MLACommonBackend,

@@ -278,6 +278,7 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):

278 q_pe: torch.Tensor,

279 kv_c_and_k_pe_cache: torch.Tensor,

280 attn_metadata: MLACommonMetadata,

281 + layer: AttentionLayer,

282 ) -> torch.Tensor:

283 if self._use_old_cutlass_mla:

284 # TODO: Remove the old cutlass MLA kernel after more extensive

vllm/v1/attention/backends/mla/flashmla.py

+4 -6

@@ -6,8 +6,7 @@ from typing import ClassVar, Optional

7 import torch

9 -from vllm.attention.backends.abstract import (AttentionType,

10 - is_quantized_kv_cache)

9 +from vllm.attention.backends.abstract import AttentionLayer, AttentionType

10 from vllm.attention.ops.flashmla import (flash_mla_with_kvcache,

11 get_mla_metadata,

12 is_flashmla_supported)

@@ -166,16 +165,13 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

165 "are not implemented for "

166 "FlashMLAImpl")

167

169 - if is_quantized_kv_cache(self.kv_cache_dtype):

170 - raise NotImplementedError(

171 - "FlashMLA V1 with FP8 KV cache not yet supported")

172 -

168 def _forward_decode(

169 self,

170 q_nope: torch.Tensor,

171 q_pe: torch.Tensor,

172 kv_c_and_k_pe_cache: torch.Tensor,

173 attn_metadata: FlashMLAMetadata,

174 + layer: AttentionLayer,

175 ) -> torch.Tensor:

176 assert kv_c_and_k_pe_cache.numel() > 0

177 assert attn_metadata.decode is not None

@@ -194,6 +190,8 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

190 num_splits=attn_metadata.decode.num_splits,

191 softmax_scale=self.scale,

192 causal=True,

193 + descale_q=layer._q_scale.reshape(1),

194 + descale_k=layer._k_scale.reshape(1),

195 )

196

197 return self._v_up_proj(o)

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

@@ -7,6 +7,7 @@ from typing import ClassVar, Optional

7 import torch

9 import vllm.envs as envs

10 +from vllm.attention.backends.abstract import AttentionLayer

11 from vllm.attention.ops.rocm_aiter_mla import aiter_mla_decode_fwd

12 from vllm.config import VllmConfig

13 from vllm.utils import cdiv

@@ -221,6 +222,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

222 q_pe: torch.Tensor,

223 kv_c_and_k_pe_cache: torch.Tensor,

224 attn_metadata: AiterMLAMetadata,

225 + layer: AttentionLayer,

226 ) -> torch.Tensor:

227 assert kv_c_and_k_pe_cache.numel() > 0

228 assert attn_metadata.decode is not None

vllm/v1/attention/backends/mla/triton_mla.py

+2 -1

@@ -6,7 +6,7 @@ from typing import Optional

6 import torch

8 from vllm import envs

9 -from vllm.attention.backends.abstract import (AttentionType,

9 +from vllm.attention.backends.abstract import (AttentionLayer, AttentionType,

10 is_quantized_kv_cache)

11 from vllm.attention.ops.triton_decode_attention import decode_attention_fwd

12 from vllm.attention.ops.triton_flash_attention import triton_attention

@@ -127,6 +127,7 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

127 q_pe: torch.Tensor,

128 kv_c_and_k_pe_cache: torch.Tensor,

129 attn_metadata: MLACommonMetadata,

130 + layer: AttentionLayer,

131 ) -> torch.Tensor:

132 assert kv_c_and_k_pe_cache.numel() > 0

133 assert attn_metadata.decode is not None