@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / fa7e254a7f
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

[New Model] DeepSeek-V3.2 (Rebased to Main) (#25896)

Signed-off-by: Chen Zhang <zhangch99@outlook.com> Signed-off-by: youkaichao <youkaichao@gmail.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: NickLucche <nlucches@redhat.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Signed-off-by: Lucia Fang <fanglu@meta.com> Co-authored-by: Chen Zhang <zhangch99@outlook.com> Co-authored-by: youkaichao <youkaichao@gmail.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: mgoin <mgoin64@gmail.com> Co-authored-by: Lucia Fang <116399278+luccafong@users.noreply.github.com> Co-authored-by: Lucia Fang <fanglu@meta.com> Co-authored-by: NickLucche <nlucches@redhat.com> Co-authored-by: Siyuan Fu <siyuanf@nvidia.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Xiaozhu Meng <mxz297@gmail.com> Co-authored-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>

Yongye Zhu committed Sep 30, 2025 at 05:14 UTC fa7e254a7f3ea7b28f1c5908c83ac57275d3c090
71 files changed +3918 -221
cmake/external_projects/flashmla.cmake
+76 -11
@@ -18,8 +18,8 @@ if(FLASH_MLA_SRC_DIR)
18 else()
19 FetchContent_Declare(
20 flashmla
21 - GIT_REPOSITORY https://github.com/vllm-project/FlashMLA.git
22 - GIT_TAG a757314c04eedd166e329e846c820eb1bdd702de
21 + GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
22 + GIT_TAG 5f65b85703c7ed75fda01e06495077caad207c3f
23 GIT_PROGRESS TRUE
24 CONFIGURE_COMMAND ""
25 BUILD_COMMAND ""
@@ -33,23 +33,64 @@ message(STATUS "FlashMLA is available at ${flashmla_SOURCE_DIR}")
33 # The FlashMLA kernels only work on hopper and require CUDA 12.3 or later.
34 # Only build FlashMLA kernels if we are building for something compatible with
35 # sm90a
36 -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)
36 +
37 +set(SUPPORT_ARCHS)
38 +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.3)
39 + list(APPEND SUPPORT_ARCHS 9.0a)
40 +endif()
41 +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.8)
42 + list(APPEND SUPPORT_ARCHS 10.0a)
43 +endif()
44 +
45 +
46 +cuda_archs_loose_intersection(FLASH_MLA_ARCHS "${SUPPORT_ARCHS}" "${CUDA_ARCHS}")
47 +if(FLASH_MLA_ARCHS)
48 + set(VLLM_FLASHMLA_GPU_FLAGS ${VLLM_GPU_FLAGS})
49 + list(APPEND VLLM_FLASHMLA_GPU_FLAGS "--expt-relaxed-constexpr" "--expt-extended-lambda" "--use_fast_math")
50 +
51 set(FlashMLA_SOURCES
39 - ${flashmla_SOURCE_DIR}/csrc/flash_api.cpp
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/splitkv_mla.cu
43 - ${flashmla_SOURCE_DIR}/csrc/kernels_fp8/flash_fwd_mla_fp8_sm90.cu)
52 + ${flashmla_SOURCE_DIR}/csrc/torch_api.cpp
53 + ${flashmla_SOURCE_DIR}/csrc/pybind.cpp
54 + ${flashmla_SOURCE_DIR}/csrc/smxx/get_mla_metadata.cu
55 + ${flashmla_SOURCE_DIR}/csrc/smxx/mla_combine.cu
56 + ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/splitkv_mla.cu
57 + ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/splitkv_mla.cu
58 + ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/fwd.cu
59 + ${flashmla_SOURCE_DIR}/csrc/sm100/decode/sparse_fp8/splitkv_mla.cu
60 + ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
61 + ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu
62 + ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd.cu
63 + )
64 +
65 + set(FlashMLA_Extension_SOURCES
66 + ${flashmla_SOURCE_DIR}/csrc/extension/torch_api.cpp
67 + ${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/pybind.cpp
68 + ${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/flash_fwd_mla_fp8_sm90.cu
69 + )
70
71 set(FlashMLA_INCLUDES
72 + ${flashmla_SOURCE_DIR}/csrc
73 + ${flashmla_SOURCE_DIR}/csrc/sm90
74 + ${flashmla_SOURCE_DIR}/csrc/cutlass/include
75 + ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
76 + )
77 +
78 + set(FlashMLA_Extension_INCLUDES
79 + ${flashmla_SOURCE_DIR}/csrc
80 + ${flashmla_SOURCE_DIR}/csrc/sm90
81 + ${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/
82 ${flashmla_SOURCE_DIR}/csrc/cutlass/include
47 - ${flashmla_SOURCE_DIR}/csrc)
83 + ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
84 + )
85
86 set_gencode_flags_for_srcs(
87 SRCS "${FlashMLA_SOURCES}"
88 CUDA_ARCHS "${FLASH_MLA_ARCHS}")
89
90 + set_gencode_flags_for_srcs(
91 + SRCS "${FlashMLA_Extension_SOURCES}"
92 + CUDA_ARCHS "${FLASH_MLA_ARCHS}")
93 +
94 define_gpu_extension_target(
95 _flashmla_C
96 DESTINATION vllm
@@ -60,8 +101,32 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.3 AND FLASH_MLA_ARCHS)
101 INCLUDE_DIRECTORIES ${FlashMLA_INCLUDES}
102 USE_SABI 3
103 WITH_SOABI)
104 +
105 + # Keep Stable ABI for the module, but *not* for CUDA/C++ files.
106 + # This prevents Py_LIMITED_API from affecting nvcc and C++ compiles.
107 + target_compile_options(_flashmla_C PRIVATE
108 + $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
109 + $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>)
110 +
111 + define_gpu_extension_target(
112 + _flashmla_extension_C
113 + DESTINATION vllm
114 + LANGUAGE ${VLLM_GPU_LANG}
115 + SOURCES ${FlashMLA_Extension_SOURCES}
116 + COMPILE_FLAGS ${VLLM_FLASHMLA_GPU_FLAGS}
117 + ARCHITECTURES ${VLLM_GPU_ARCHES}
118 + INCLUDE_DIRECTORIES ${FlashMLA_Extension_INCLUDES}
119 + USE_SABI 3
120 + WITH_SOABI)
121 +
122 + # Keep Stable ABI for the module, but *not* for CUDA/C++ files.
123 + # This prevents Py_LIMITED_API from affecting nvcc and C++ compiles.
124 + target_compile_options(_flashmla_extension_C PRIVATE
125 + $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
126 + $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>)
127 else()
64 - # Create an empty target for setup.py when not targeting sm90a systems
128 + # Create empty targets for setup.py when not targeting sm90a systems
129 add_custom_target(_flashmla_C)
130 + add_custom_target(_flashmla_extension_C)
131 endif()
132
csrc/cache.h
+8
@@ -56,3 +56,11 @@ void cp_gather_cache(
56 torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
57 torch::Tensor const& cu_seq_lens, // [BATCH+1]
58 int64_t batch_size, std::optional<torch::Tensor> seq_starts = std::nullopt);
59 +
60 +// Indexer K quantization and cache function
61 +void indexer_k_quant_and_cache(
62 + torch::Tensor& k, // [num_tokens, head_dim]
63 + torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride]
64 + torch::Tensor& slot_mapping, // [num_tokens]
65 + int64_t quant_block_size, // quantization block size
66 + const std::string& scale_fmt);
csrc/cache_kernels.cu
+252 -6
@@ -16,6 +16,7 @@
16
17 #include <algorithm>
18 #include <cassert>
19 +#include <cfloat> // FLT_MIN
20 #include <map>
21 #include <vector>
22
@@ -396,6 +397,176 @@ __global__ void concat_and_cache_mla_kernel(
397 copy(k_pe, kv_cache, k_pe_stride, block_stride, pe_dim, kv_lora_rank);
398 }
399
400 +template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
401 +__global__ void concat_and_cache_ds_mla_kernel(
402 + const scalar_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank]
403 + const scalar_t* __restrict__ k_pe, // [num_tokens, pe_dim]
404 + cache_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank
405 + // + pe_dim)]
406 + const int64_t* __restrict__ slot_mapping, // [num_tokens]
407 + const int block_stride, //
408 + const int entry_stride, //
409 + const int kv_c_stride, //
410 + const int k_pe_stride, //
411 + const int kv_lora_rank, //
412 + const int pe_dim, //
413 + const int block_size, //
414 + const float* scale //
415 +) {
416 + const int64_t token_idx = blockIdx.x;
417 + const int64_t slot_idx = slot_mapping[token_idx];
418 + // NOTE: slot_idx can be -1 if the token is padded
419 + if (slot_idx < 0) {
420 + return;
421 + }
422 + const int64_t block_idx = slot_idx / block_size;
423 + const int64_t block_offset = slot_idx % block_size;
424 + const int64_t dst_idx_start =
425 + block_idx * block_stride + block_offset * entry_stride;
426 +
427 + // Create 4 tile scales in shared memory
428 + __shared__ float smem[20];
429 + float* shard_abs_max = smem;
430 + float* tile_scales = smem + 16;
431 +
432 + // For the NoPE part, each tile of 128 elements is handled by 4 warps
433 + // (128 threads). There are 4 total tiles, so 16 warps (512 threads).
434 + // The first thread of the first warp in each tile writes the scale
435 + // value for the tile. The RoPE part (last 64 elements) is handled
436 + // by another 2 warps (64 threads).
437 + // So in total, we use 18 warps (576 threads) per block.
438 +
439 + // Cast kv_cache to 16_bit for RoPE values
440 + scalar_t* kv_cache_16bit =
441 + reinterpret_cast<scalar_t*>(&kv_cache[dst_idx_start]);
442 +
443 + // The last 64 threads handle the RoPE part
444 + if (threadIdx.x >= kv_lora_rank) {
445 + const int8_t pe_idx = threadIdx.x - kv_lora_rank;
446 + const int64_t src_idx = token_idx * k_pe_stride + pe_idx;
447 + // RoPE values start after the packed 8-bit NoPE values and the
448 + // 32-bit scales
449 + const int64_t dst_idx = kv_lora_rank / 2 + 8 + pe_idx;
450 + kv_cache_16bit[dst_idx] = k_pe[src_idx];
451 + return;
452 + }
453 +
454 + // Determine the scale for each chunk of NoPE
455 + const int16_t tile_idx = threadIdx.x >> 7;
456 + const int16_t warp_idx = (threadIdx.x & 127) >> 5;
457 + const int16_t lane_idx = threadIdx.x & 31;
458 +
459 + // Load the NoPE element for this thread into registers
460 + const int64_t src_idx = token_idx * kv_c_stride + threadIdx.x;
461 + const scalar_t src_val = kv_c[src_idx];
462 +
463 + // Warp-level reduction to find the max absolute value in the warp
464 + float max_abs = fabsf(src_val);
465 +#pragma unroll
466 + for (int offset = 16; offset > 0; offset /= 2) {
467 +#ifdef USE_ROCM
468 + max_abs = fmaxf(max_abs, __shfl_down_sync(UINT64_MAX, max_abs, offset));
469 +#else
470 + max_abs = fmaxf(max_abs, __shfl_down_sync(0xFFFFFFFF, max_abs, offset));
471 +#endif
472 + }
473 +
474 + // The first lane of each warp in each tile writes the max_abs of this part
475 + // of the tile to shared memory
476 + if (lane_idx == 0) {
477 + shard_abs_max[tile_idx * 4 + warp_idx] = max_abs;
478 + }
479 + __syncthreads();
480 +
481 + // The first lane of the first warp in each tile computes the scale for the
482 + // tile and writes it to shared memory and to kv_cache
483 + if (warp_idx == 0 && lane_idx == 0) {
484 + float4 shard_abs_max_vec =
485 + reinterpret_cast<float4*>(shard_abs_max)[tile_idx];
486 + float tile_scale = fmaxf(fmaxf(shard_abs_max_vec.x, shard_abs_max_vec.y),
487 + fmaxf(shard_abs_max_vec.z, shard_abs_max_vec.w)) /
488 + 448.f;
489 +
490 + // Avoid division by zero in `scaled_convert`
491 + tile_scales[tile_idx] = fmaxf(tile_scale, FLT_MIN);
492 + float* kv_cache_32bit = reinterpret_cast<float*>(&kv_cache[dst_idx_start]);
493 + const uint64_t dst_idx = kv_lora_rank / 4 + tile_idx;
494 + kv_cache_32bit[dst_idx] = tile_scales[tile_idx];
495 + }
496 +
497 + __syncthreads();
498 +
499 + // Now all threads in the block scale and write their element
500 + const float scale_val = tile_scales[tile_idx];
501 + const int64_t dst_idx = dst_idx_start + threadIdx.x;
502 + kv_cache[dst_idx] =
503 + fp8::scaled_convert<uint8_t, scalar_t, Fp8KVCacheDataType::kFp8E4M3>(
504 + src_val, scale_val);
505 +}
506 +
507 +template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
508 +__global__ void indexer_k_quant_and_cache_kernel(
509 + const scalar_t* __restrict__ k, // [num_tokens, head_dim]
510 + cache_t* __restrict__ kv_cache, // [num_blocks, block_size, cache_stride]
511 + const int64_t* __restrict__ slot_mapping, // [num_tokens]
512 + const int head_dim, // dimension of each head
513 + const int quant_block_size, // quantization block size
514 + const int cache_block_size, // cache block size
515 + const int cache_stride, // stride for each token in kv_cache
516 + const bool use_ue8m0 // use ue8m0 scale format
517 +) {
518 + constexpr int VEC_SIZE = 4;
519 + const int64_t token_idx = blockIdx.x;
520 + const int64_t head_dim_idx = (blockIdx.y * blockDim.y * blockDim.x +
521 + threadIdx.y * blockDim.x + threadIdx.x) *
522 + VEC_SIZE;
523 + const int64_t slot_idx = slot_mapping[token_idx];
524 + const int64_t block_idx = slot_idx / cache_block_size;
525 + const int64_t block_offset = slot_idx % cache_block_size;
526 +
527 + // NOTE: slot_idx can be -1 if the token is padded
528 + if (slot_idx < 0 || (head_dim_idx >= head_dim)) {
529 + return;
530 + }
531 +
532 + float2 k_val = (reinterpret_cast<const float2*>(
533 + k))[(token_idx * head_dim + head_dim_idx) / VEC_SIZE];
534 + scalar_t* k_val_ptr = reinterpret_cast<scalar_t*>(&k_val);
535 + float amax = 0.0f;
536 + for (int i = 0; i < VEC_SIZE; i++) {
537 + amax = fmaxf(amax, fabsf(float(k_val_ptr[i])));
538 + }
539 + __syncwarp();
540 +
541 + // Reduced amax
542 + for (int mask = 16; mask > 0; mask /= 2) {
543 +#ifdef USE_ROCM
544 + amax = fmaxf(amax, __shfl_xor_sync(uint64_t(-1), amax, mask));
545 +#else
546 + amax = fmaxf(amax, __shfl_xor_sync(unsigned(-1), amax, mask));
547 +#endif
548 + }
549 + __syncwarp();
550 + float scale = fmaxf(amax, 1e-4) / 448.0f;
551 + if (use_ue8m0) {
552 + scale = exp2f(ceilf(log2f(scale)));
553 + }
554 +
555 + const int64_t dst_offset = block_idx * cache_block_size * cache_stride +
556 + block_offset * head_dim + head_dim_idx;
557 + for (int i = 0; i < VEC_SIZE; i++) {
558 + kv_cache[dst_offset + i] =
559 + fp8::scaled_convert<cache_t, scalar_t, kv_dt>(k_val_ptr[i], scale);
560 + }
561 + if (threadIdx.x == 0) {
562 + const int64_t dst_scale_idx =
563 + block_idx * cache_block_size * cache_stride +
564 + cache_block_size * head_dim +
565 + (block_offset * head_dim + head_dim_idx) * 4 / quant_block_size;
566 + reinterpret_cast<float*>(kv_cache)[dst_scale_idx / 4] = scale;
567 + }
568 +}
569 +
570 } // namespace vllm
571
572 // KV_T is the data type of key and value tensors.
@@ -438,7 +609,7 @@ void reshape_and_cache(
609 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
610
611 DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype,
441 - CALL_RESHAPE_AND_CACHE)
612 + CALL_RESHAPE_AND_CACHE);
613 }
614
615 // KV_T is the data type of key and value tensors.
@@ -509,6 +680,18 @@ void reshape_and_cache_flash(
680 kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \
681 reinterpret_cast<const float*>(scale.data_ptr()));
682
683 +// KV_T is the data type of key and value tensors.
684 +// CACHE_T is the stored data type of kv-cache.
685 +#define CALL_CONCAT_AND_CACHE_DS_MLA(KV_T, CACHE_T, KV_DTYPE) \
686 + vllm::concat_and_cache_ds_mla_kernel<KV_T, CACHE_T, KV_DTYPE> \
687 + <<<grid, block, 0, stream>>>( \
688 + reinterpret_cast<KV_T*>(kv_c.data_ptr()), \
689 + reinterpret_cast<KV_T*>(k_pe.data_ptr()), \
690 + reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
691 + slot_mapping.data_ptr<int64_t>(), block_stride, entry_stride, \
692 + kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \
693 + reinterpret_cast<const float*>(scale.data_ptr()));
694 +
695 void concat_and_cache_mla(
696 torch::Tensor& kv_c, // [num_tokens, kv_lora_rank]
697 torch::Tensor& k_pe, // [num_tokens, pe_dim]
@@ -531,20 +714,44 @@ void concat_and_cache_mla(
714 int pe_dim = k_pe.size(1);
715 int block_size = kv_cache.size(1);
716
534 - TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + pe_dim);
717 + if (kv_cache_dtype == "fp8_ds_mla") {
718 + TORCH_CHECK(kv_lora_rank == 512, "kv_lora_rank must be 512 for fp8_ds_mla");
719 + TORCH_CHECK(pe_dim == 64, "pe_dim must be 64 for fp8_ds_mla");
720 + TORCH_CHECK(kv_cache.size(2) == 656 / kv_cache.itemsize(),
721 + "kv_cache.size(2) must be 656 bytes for fp8_ds_mla");
722 + TORCH_CHECK(kv_c.itemsize() == 2,
723 + "kv_c.itemsize() must be 2 for fp8_ds_mla");
724 + TORCH_CHECK(k_pe.itemsize() == 2,
725 + "k_pe.itemsize() must be 2 for fp8_ds_mla");
726 + } else {
727 + TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + pe_dim);
728 + }
729
730 int kv_c_stride = kv_c.stride(0);
731 int k_pe_stride = k_pe.stride(0);
732 int block_stride = kv_cache.stride(0);
733 int entry_stride = kv_cache.stride(1);
734
541 - dim3 grid(num_tokens);
542 - dim3 block(std::min(kv_lora_rank, 512));
735 const at::cuda::OptionalCUDAGuard device_guard(device_of(kv_c));
736 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
737
546 - DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype,
547 - CALL_CONCAT_AND_CACHE_MLA);
738 + if (kv_cache_dtype == "fp8_ds_mla") {
739 + dim3 grid(num_tokens);
740 + // For the NoPE part, each tile of 128 elements is handled by 4 warps
741 + // (128 threads). There are 4 total tiles, so 16 warps (512 threads).
742 + // The first thread of the first warp in each tile writes the scale
743 + // value for the tile. The RoPE part (last 64 elements) is handled
744 + // by another 2 warps (64 threads).
745 + // So in total, we use 18 warps (576 threads) per block.
746 + dim3 block(576);
747 + DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype,
748 + CALL_CONCAT_AND_CACHE_DS_MLA);
749 + } else {
750 + dim3 grid(num_tokens);
751 + dim3 block(std::min(kv_lora_rank, 512));
752 + DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype,
753 + CALL_CONCAT_AND_CACHE_MLA);
754 + }
755 }
756
757 namespace vllm {
@@ -922,3 +1129,42 @@ void cp_gather_cache(
1129 TORCH_CHECK(false, "Unsupported data type width: ", dtype_bits);
1130 }
1131 }
1132 +
1133 +// Macro to dispatch the kernel based on the data type.
1134 +#define CALL_INDEXER_K_QUANT_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \
1135 + vllm::indexer_k_quant_and_cache_kernel<KV_T, CACHE_T, KV_DTYPE> \
1136 + <<<grid, block, 0, stream>>>( \
1137 + reinterpret_cast<KV_T*>(k.data_ptr()), \
1138 + reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
1139 + slot_mapping.data_ptr<int64_t>(), head_dim, quant_block_size, \
1140 + cache_block_size, cache_stride, use_ue8m0);
1141 +
1142 +void indexer_k_quant_and_cache(
1143 + torch::Tensor& k, // [num_tokens, head_dim]
1144 + torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride]
1145 + torch::Tensor& slot_mapping, // [num_tokens]
1146 + int64_t quant_block_size, // quantization block size
1147 + const std::string& scale_fmt) {
1148 + int num_tokens = k.size(0);
1149 + int head_dim = k.size(1);
1150 + int cache_block_size = kv_cache.size(1);
1151 + int cache_stride = kv_cache.size(2);
1152 + bool use_ue8m0 = scale_fmt == "ue8m0";
1153 +
1154 + TORCH_CHECK(k.device() == kv_cache.device(),
1155 + "k and kv_cache must be on the same device");
1156 + TORCH_CHECK(k.device() == slot_mapping.device(),
1157 + "k and slot_mapping must be on the same device");
1158 + TORCH_CHECK(head_dim % quant_block_size == 0,
1159 + "head_dim must be divisible by quant_block_size");
1160 +
1161 + constexpr int vec_size = 4;
1162 + dim3 grid(num_tokens, (head_dim + quant_block_size * vec_size - 1) /
1163 + (quant_block_size * vec_size));
1164 + dim3 block(32, vec_size);
1165 + const at::cuda::OptionalCUDAGuard device_guard(device_of(k));
1166 + const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1167 +
1168 + DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), "fp8_e4m3",
1169 + CALL_INDEXER_K_QUANT_AND_CACHE);
1170 +}
\ No newline at end of file
csrc/quantization/fp8/nvidia/quant_utils.cuh
+11
@@ -576,6 +576,17 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
576 TORCH_CHECK(false, \
577 "Unsupported input type of kv cache: ", SRC_DTYPE); \
578 } \
579 + } else if (KV_DTYPE == "fp8_ds_mla") { \
580 + if (SRC_DTYPE == at::ScalarType::Float) { \
581 + FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
582 + } else if (SRC_DTYPE == at::ScalarType::Half) { \
583 + FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
584 + } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
585 + FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
586 + } else { \
587 + TORCH_CHECK(false, \
588 + "Unsupported input type of kv cache: ", SRC_DTYPE); \
589 + } \
590 } else { \
591 TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
592 } \
csrc/torch_bindings.cpp
+7
@@ -713,6 +713,13 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
713 "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, "
714 "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()");
715 cache_ops.impl("cp_gather_cache", torch::kCUDA, &cp_gather_cache);
716 +
717 + cache_ops.def(
718 + "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "
719 + "slot_mapping, "
720 + "int quant_block_size, str kv_cache_dtype) -> ()");
721 + cache_ops.impl("indexer_k_quant_and_cache", torch::kCUDA,
722 + &indexer_k_quant_and_cache);
723 }
724
725 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
setup.py
+4
@@ -322,6 +322,8 @@ class precompiled_wheel_utils:
322 "vllm/_C.abi3.so",
323 "vllm/_moe_C.abi3.so",
324 "vllm/_flashmla_C.abi3.so",
325 + "vllm/_flashmla_extension_C.abi3.so",
326 + "vllm/_sparse_flashmla_C.abi3.so",
327 "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
328 "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
329 "vllm/cumem_allocator.abi3.so",
@@ -589,6 +591,8 @@ if _is_cuda():
591 # not targeting a hopper system
592 ext_modules.append(
593 CMakeExtension(name="vllm._flashmla_C", optional=True))
594 + ext_modules.append(
595 + CMakeExtension(name="vllm._flashmla_extension_C", optional=True))
596 ext_modules.append(CMakeExtension(name="vllm.cumem_allocator"))
597
598 if _build_custom_ops():
tests/compile/test_fusion_attn.py
-1
@@ -191,7 +191,6 @@ class AttentionQuantPatternModel(torch.nn.Module):
191 num_kv_heads=self.num_kv_heads,
192 head_size=self.head_size,
193 dtype=self.kv_cache_dtype,
194 - use_mla=False,
194 ),
195 layer_names=[self.attn.layer_name],
196 vllm_config=self.vllm_config,
tests/kernels/attention/test_cache.py
+113
@@ -593,6 +593,119 @@ def test_concat_and_cache_mla(
593 torch.testing.assert_close(kv_cache, ref_kv_cache)
594
595
596 +@pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
597 +@pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
598 +@pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA)
599 +@pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA)
600 +@pytest.mark.parametrize("num_blocks", NUM_BLOCKS_MLA)
601 +@pytest.mark.parametrize("dtype", DTYPES)
602 +@pytest.mark.parametrize("seed", SEEDS)
603 +@pytest.mark.parametrize("device", CUDA_DEVICES)
604 +@torch.inference_mode()
605 +def test_concat_and_cache_ds_mla(
606 + kv_lora_rank: int,
607 + qk_rope_head_dim: int,
608 + num_tokens: int,
609 + block_size: int,
610 + num_blocks: int,
611 + dtype: torch.dtype,
612 + seed: int,
613 + device: str,
614 +) -> None:
615 + if dtype.itemsize != 2:
616 + pytest.skip("ds_mla only supports 16-bit input")
617 + kv_cache_dtype = "fp8_ds_mla"
618 + current_platform.seed_everything(seed)
619 + torch.set_default_device(device)
620 +
621 + total_slots = num_blocks * block_size
622 + slot_mapping_lst = random.sample(range(total_slots), num_tokens)
623 + slot_mapping = torch.tensor(slot_mapping_lst,
624 + dtype=torch.long,
625 + device=device)
626 +
627 + kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=dtype, device=device)
628 + k_pe = torch.randn(num_tokens,
629 + qk_rope_head_dim,
630 + dtype=dtype,
631 + device=device)
632 + entry_size = kv_lora_rank + (4 * 4) + (2 * qk_rope_head_dim)
633 +
634 + scale = torch.tensor(1.0, dtype=torch.float32, device=device)
635 + kv_cache = _create_mla_cache(num_blocks,
636 + block_size,
637 + entry_size,
638 + dtype=torch.uint8,
639 + kv_cache_dtype=kv_cache_dtype,
640 + device=device)
641 +
642 + ref_cache = torch.zeros_like(kv_cache, dtype=kv_cache.dtype)
643 + tile_data = torch.zeros(128, dtype=dtype, device=device)
644 +
645 + for i in range(num_tokens):
646 + slot = slot_mapping[i].item()
647 + block_idx = slot // block_size
648 + block_offset = slot % block_size
649 +
650 + ref_cache_slice = ref_cache[block_idx, block_offset]
651 + ref_cache_16bit = ref_cache_slice.view(dtype)
652 + ref_cache_32bit = ref_cache_slice.view(torch.float32)
653 +
654 + kv_c_data = kv_c[i]
655 + for tile_idx in range(4):
656 + tile_start = tile_idx * 128
657 + tile_end = (tile_idx + 1) * 128
658 + tile_data[:] = kv_c_data[tile_start:tile_end]
659 +
660 + # tile_scale = tile_data.amax().to(torch.float32) / 448.
661 + # NOTE: Using torch's amax() gives different results,
662 + # so this must be manually computed.
663 + tile_data_float = tile_data.to(torch.float32)
664 + manual_max = abs(tile_data_float[0])
665 + for j in range(1, 128):
666 + manual_max = max(manual_max, abs(tile_data_float[j]))
667 + tile_scale = manual_max / 448.
668 +
669 + ref_cache_32bit[kv_lora_rank // 4 + tile_idx] = tile_scale
670 +
671 + ops.convert_fp8(ref_cache_slice[tile_start:tile_end],
672 + tile_data,
673 + tile_scale.item(),
674 + kv_dtype="fp8")
675 +
676 + for j in range(qk_rope_head_dim):
677 + ref_cache_16bit[kv_lora_rank // 2 + 8 + j] = k_pe[i, j]
678 +
679 + opcheck(
680 + torch.ops._C_cache_ops.concat_and_cache_mla,
681 + (kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale),
682 + test_utils=DEFAULT_OPCHECK_TEST_UTILS,
683 + )
684 +
685 + ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping,
686 + kv_cache_dtype, scale)
687 +
688 + for i in range(num_tokens):
689 + slot = slot_mapping[i].item()
690 + block_idx = slot // block_size
691 + block_offset = slot % block_size
692 + kv_cache_slice = kv_cache[block_idx, block_offset]
693 + ref_cache_slice = ref_cache[block_idx, block_offset]
694 +
695 + kv_nope = kv_cache_slice[:kv_lora_rank]
696 + ref_nope = ref_cache_slice[:kv_lora_rank]
697 + kv_scales = kv_cache_slice.view(torch.float32)[kv_lora_rank //
698 + 4:kv_lora_rank // 4 + 4]
699 + ref_scales = ref_cache_slice.view(
700 + torch.float32)[kv_lora_rank // 4:kv_lora_rank // 4 + 4]
701 + kv_rope = kv_cache_slice.view(dtype)[kv_lora_rank // 2 + 8:]
702 + ref_rope = ref_cache_slice.view(dtype)[kv_lora_rank // 2 + 8:]
703 +
704 + torch.testing.assert_close(kv_nope, ref_nope, atol=0.001, rtol=0.1)
705 + torch.testing.assert_close(kv_scales, ref_scales, atol=0.001, rtol=0.1)
706 + torch.testing.assert_close(kv_rope, ref_rope, atol=0.001, rtol=0.1)
707 +
708 +
709 @pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
710 @pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
711 @pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA)
tests/kernels/attention/test_deepgemm_attention.py new
+279
@@ -0,0 +1,279 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import random
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm.platforms import current_platform
9 +from vllm.utils import cdiv, has_deep_gemm
10 +from vllm.utils.deep_gemm import (_ceil_to_ue8m0, calc_diff, fp8_mqa_logits,
11 + fp8_paged_mqa_logits, get_num_sms,
12 + get_paged_mqa_logits_metadata)
13 +
14 +
15 +def kv_cache_cast_to_fp8(x: torch.Tensor) -> torch.Tensor:
16 + # x: (num_blocks, block_size, 1, head_dim)
17 + num_blocks, block_size, num_heads, head_dim = x.shape
18 + assert num_heads == 1
19 + x_amax = x.abs().float().amax(dim=3, keepdim=True).clamp(1e-4)
20 + sf = x_amax / 448.0
21 + x_scaled = (x * (1.0 / sf)).to(torch.float8_e4m3fn)
22 + x_fp8 = torch.empty(
23 + (num_blocks, block_size * (head_dim + 4)),
24 + device=x.device,
25 + dtype=torch.uint8,
26 + )
27 + x_fp8[:, :block_size * head_dim] = x_scaled.view(
28 + num_blocks, block_size * head_dim).view(dtype=torch.uint8)
29 + x_fp8[:,
30 + block_size * head_dim:] = sf.view(num_blocks,
31 + block_size).view(dtype=torch.uint8)
32 + return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4)
33 +
34 +
35 +def per_custom_dims_cast_to_fp8(
36 + x: torch.Tensor, dims: tuple,
37 + use_ue8m0: bool) -> tuple[torch.Tensor, torch.Tensor]:
38 + excluded_dims = tuple([i for i in range(x.dim()) if i not in set(dims)])
39 + x_amax = x.abs().float().amax(dim=excluded_dims, keepdim=True).clamp(1e-4)
40 + sf = x_amax / 448.0
41 + sf = _ceil_to_ue8m0(sf) if use_ue8m0 else sf
42 + x_scaled = (x * (1.0 / sf)).to(torch.float8_e4m3fn)
43 + return x_scaled, sf.squeeze()
44 +
45 +
46 +def _generate_cp_test_data(seq_len: int, seq_len_kv: int):
47 + assert seq_len_kv % seq_len == 0 and seq_len % 2 == 0
48 + chunk_size = seq_len // 2
49 + cp_size = seq_len_kv // seq_len
50 + cp_id = cp_size // 3
51 + ks = torch.zeros(seq_len, dtype=torch.int, device="cuda")
52 + ke = torch.zeros(seq_len, dtype=torch.int, device="cuda")
53 + for i in range(chunk_size):
54 + ke[i] = cp_id * chunk_size + i
55 + ke[i + chunk_size] = (cp_size * 2 - 1 - cp_id) * chunk_size + i
56 + return ks, ke
57 +
58 +
59 +def _ref_fp8_mqa_logits(
60 + q: torch.Tensor,
61 + kv: torch.Tensor,
62 + weights: torch.Tensor,
63 + cu_seqlen_ks: torch.Tensor,
64 + cu_seqlen_ke: torch.Tensor,
65 +):
66 + seq_len_kv = kv.shape[0]
67 +
68 + k = kv
69 + q = q.float()
70 + k = k.float()
71 +
72 + mask_lo = (torch.arange(0, seq_len_kv, device="cuda")[None, :]
73 + >= cu_seqlen_ks[:, None])
74 + mask_hi = (torch.arange(0, seq_len_kv, device="cuda")[None, :]
75 + < cu_seqlen_ke[:, None])
76 + mask = mask_lo & mask_hi
77 +
78 + score = torch.einsum("mhd,and->hmn", q, k)
79 + logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0)
80 + logits = logits.masked_fill(~mask, float("-inf"))
81 +
82 + return logits
83 +
84 +
85 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
86 +@pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available")
87 +@pytest.mark.skipif(not current_platform.has_device_capability(90),
88 + reason="SM90 and SM100 only")
89 +def test_deepgemm_fp8_mqa_logits():
90 + torch.manual_seed(0)
91 + random.seed(0)
92 + num_heads, head_dim = 32, 128
93 + for seq_len in (512, ):
94 + for seq_len_kv in (1024, ):
95 + for disable_cp in (False, True):
96 + q = torch.randn(
97 + seq_len,
98 + num_heads,
99 + head_dim,
100 + device="cuda",
101 + dtype=torch.bfloat16,
102 + )
103 + kv = torch.randn(seq_len_kv,
104 + head_dim,
105 + device="cuda",
106 + dtype=torch.bfloat16)
107 + weights = torch.randn(seq_len,
108 + num_heads,
109 + device="cuda",
110 + dtype=torch.float32)
111 +
112 + if disable_cp:
113 + ks = torch.zeros(seq_len, dtype=torch.int, device="cuda")
114 + ke = torch.arange(seq_len, dtype=torch.int,
115 + device="cuda") + (seq_len_kv - seq_len)
116 + else:
117 + ks, ke = _generate_cp_test_data(seq_len, seq_len_kv)
118 +
119 + q_fp8 = q.to(torch.float8_e4m3fn)
120 + kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0, ), False)
121 + logits = fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke)
122 +
123 + ref_logits = _ref_fp8_mqa_logits(
124 + q=q,
125 + kv=kv,
126 + weights=weights,
127 + cu_seqlen_ks=ks,
128 + cu_seqlen_ke=ke,
129 + )
130 +
131 + ref_neginf_mask = ref_logits == float("-inf")
132 + neginf_mask = logits == float("-inf")
133 + assert torch.equal(neginf_mask, ref_neginf_mask)
134 +
135 + ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0)
136 + logits = logits.masked_fill(neginf_mask, 0)
137 + diff = calc_diff(logits, ref_logits)
138 + assert diff < 1e-3, f"{diff=}"
139 +
140 +
141 +def _ref_fp8_paged_mqa_logits(
142 + q: torch.Tensor,
143 + kv_cache: torch.Tensor,
144 + weights: torch.Tensor,
145 + context_lens: torch.Tensor,
146 + block_tables: torch.Tensor,
147 + max_model_len: int,
148 +):
149 + batch_size, next_n, _, _ = q.size()
150 + _, block_size, _, _ = kv_cache.size()
151 + logits = torch.full(
152 + [batch_size * next_n, max_model_len],
153 + float("-inf"),
154 + device=q.device,
155 + dtype=torch.float32,
156 + )
157 + context_lens_list = context_lens.tolist()
158 + for i in range(batch_size):
159 + context_len = context_lens_list[i]
160 + q_offsets = torch.arange(context_len - next_n,
161 + context_len,
162 + device="cuda")
163 + weight_slice = (weights[i * next_n:(i + 1) * next_n, :].transpose(
164 + 0, 1).contiguous())
165 + for block_rk in range(cdiv(context_len, block_size)):
166 + block_idx = block_tables[i][block_rk]
167 + qx, kx = q[i], kv_cache[block_idx]
168 + k_offsets = torch.arange(
169 + block_rk * block_size,
170 + (block_rk + 1) * block_size,
171 + device="cuda",
172 + )
173 + mask = (k_offsets[None, :] < context_len) & (k_offsets[None, :]
174 + <= q_offsets[:, None])
175 + s = torch.where(
176 + mask[None, :, :],
177 + (qx.transpose(0, 1) @ kx.transpose(0, 1).transpose(1, 2)).to(
178 + logits.dtype),
179 + float("-inf"),
180 + )
181 + s = torch.relu(s) * weight_slice[..., None]
182 + s = s.sum(dim=0)
183 + logits[
184 + i * next_n:(i + 1) * next_n,
185 + block_rk * block_size:(block_rk + 1) * block_size,
186 + ] = torch.where(k_offsets[None, :] <= q_offsets[:, None], s,
187 + float("-inf"))
188 + return logits
189 +
190 +
191 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
192 +@pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available")
193 +@pytest.mark.skipif(not current_platform.has_device_capability(90),
194 + reason="SM90 and SM100 only")
195 +def test_deepgemm_fp8_paged_mqa_logits():
196 + torch.manual_seed(0)
197 + random.seed(0)
198 +
199 + max_model_len = 4096
200 + for batch_size, next_n in [(4, 1), (2, 2)]:
201 + for heads, index_dim in [(32, 128)]:
202 + for avg_kv in (2048, ):
203 + num_blocks, blocksize = max_model_len * 2, 64
204 +
205 + q = torch.randn(
206 + (batch_size, next_n, heads, index_dim),
207 + device="cuda",
208 + dtype=torch.bfloat16,
209 + )
210 + kv_cache = torch.randn(
211 + (num_blocks, blocksize, 1, index_dim),
212 + device="cuda",
213 + dtype=torch.bfloat16,
214 + )
215 + weights = torch.randn(
216 + (batch_size * next_n, heads),
217 + device="cuda",
218 + dtype=torch.float32,
219 + )
220 +
221 + context_lens = (torch.randint(int(0.8 * avg_kv),
222 + int(1.2 * avg_kv),
223 + (batch_size, )).cuda().to(
224 + torch.int32))
225 + max_block_len = ((context_lens.max().item() + blocksize - 1) //
226 + blocksize * blocksize)
227 + block_tables = torch.zeros(
228 + (batch_size, max_block_len),
229 + device="cuda",
230 + dtype=torch.int32,
231 + )
232 +
233 + counter = 0
234 + block_idx_pool = list(range(num_blocks))
235 + random.shuffle(block_idx_pool)
236 + for i in range(batch_size):
237 + ctx_len = int(context_lens[i].item())
238 + for j in range((ctx_len + blocksize - 1) // blocksize):
239 + block_tables[i][j] = block_idx_pool[counter]
240 + counter += 1
241 +
242 + q_fp8 = q.to(torch.float8_e4m3fn)
243 + kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache)
244 +
245 + schedule_metadata = get_paged_mqa_logits_metadata(
246 + context_lens, blocksize, get_num_sms())
247 + logits = fp8_paged_mqa_logits(
248 + q_fp8,
249 + kv_cache_fp8,
250 + weights,
251 + context_lens,
252 + block_tables,
253 + schedule_metadata,
254 + max_model_len,
255 + )
256 +
257 + ref_logits = _ref_fp8_paged_mqa_logits(
258 + q,
259 + kv_cache,
260 + weights,
261 + context_lens,
262 + block_tables,
263 + max_model_len,
264 + )
265 +
266 + positions = (torch.arange(max_model_len,
267 + device="cuda").unsqueeze(0).expand(
268 + batch_size * next_n, -1))
269 + row_indices = (
270 + torch.arange(batch_size * next_n, device="cuda") // next_n)
271 + next_n_offset = (
272 + torch.arange(batch_size * next_n, device="cuda") % next_n)
273 + mask = positions <= (context_lens[row_indices] - next_n +
274 + next_n_offset).unsqueeze(1)
275 +
276 + logits = logits.masked_fill(~mask, 0)
277 + ref_logits = ref_logits.masked_fill(~mask, 0)
278 + diff = calc_diff(logits, ref_logits)
279 + assert diff < 1e-3, f"{diff=}"
tests/kernels/attention/test_flashmla.py
+10 -12
@@ -97,18 +97,16 @@ def test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal,
97 descale_k = None
98
99 def flash_mla():
100 - return flash_mla_with_kvcache(
101 - q,
102 - blocked_k,
103 - block_table,
104 - cache_seqlens,
105 - dv,
106 - tile_scheduler_metadata,
107 - num_splits,
108 - causal=causal,
109 - descale_q=descale_q,
110 - descale_k=descale_k,
111 - )
100 + return flash_mla_with_kvcache(q,
101 + blocked_k,
102 + block_table,
103 + cache_seqlens,
104 + dv,
105 + tile_scheduler_metadata,
106 + num_splits,
107 + causal=causal,
108 + descale_q=descale_q,
109 + descale_k=descale_k)
110
111 def scaled_dot_product_attention(query, key, value, is_causal=False):
112 query = query.float()
tests/kernels/attention/test_flashmla_sparse.py new
+119
@@ -0,0 +1,119 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import pytest
4 +import torch
5 +
6 +
7 +def _cuda_sm90_available() -> bool:
8 + if not torch.cuda.is_available():
9 + return False
10 + major, _ = torch.cuda.get_device_capability()
11 + return major == 9
12 +
13 +
14 +def test_sparse_flashmla_metadata_smoke():
15 + import vllm.attention.ops.flashmla as fm
16 + ok, reason = fm.is_flashmla_supported()
17 + if not ok or not _cuda_sm90_available():
18 + pytest.skip(reason or "SM90 not available")
19 +
20 + device = torch.device("cuda")
21 + batch_size = 1
22 + seqlen_q = 1
23 + num_heads_q = 128
24 + num_heads_k = 1
25 + q_seq_per_hk = seqlen_q * num_heads_q // num_heads_k
26 + topk = 128
27 +
28 + cache_seqlens = torch.zeros(batch_size, dtype=torch.int32, device=device)
29 +
30 + tile_md, num_splits = fm.get_mla_metadata(cache_seqlens,
31 + q_seq_per_hk,
32 + num_heads_k,
33 + num_heads_q=num_heads_q,
34 + topk=topk,
35 + is_fp8_kvcache=True)
36 + assert tile_md.dtype == torch.int32
37 + assert num_splits.dtype == torch.int32
38 +
39 +
40 +def test_sparse_flashmla_decode_smoke():
41 + import vllm.attention.ops.flashmla as fm
42 + ok, reason = fm.is_flashmla_supported()
43 + if not ok or not _cuda_sm90_available():
44 + pytest.skip(reason or "SM90 not available")
45 +
46 + device = torch.device("cuda")
47 + batch_size = 1
48 + seqlen_q = 1
49 + num_heads_q = 1
50 + head_dim_k = 576
51 + head_dim_v = 512
52 + num_heads_k = 1
53 + page_block_size = 64
54 + bytes_per_token = 656
55 + topk = 128
56 +
57 + # Metadata
58 + q_seq_per_hk = seqlen_q * num_heads_q // num_heads_k
59 + # q_heads_per_hk = num_heads_q // num_heads_k
60 + cache_seqlens = torch.zeros(batch_size, dtype=torch.int32, device=device)
61 + tile_md, num_splits = fm.get_mla_metadata(cache_seqlens,
62 + q_seq_per_hk,
63 + num_heads_k,
64 + num_heads_q=num_heads_q,
65 + topk=topk,
66 + is_fp8_kvcache=True)
67 +
68 + # Inputs
69 + q = torch.zeros((batch_size, seqlen_q, num_heads_q, head_dim_k),
70 + dtype=torch.bfloat16,
71 + device=device)
72 + k_cache = torch.zeros((1, page_block_size, num_heads_k, bytes_per_token),
73 + dtype=torch.uint8,
74 + device=device)
75 + indices = torch.zeros((batch_size, seqlen_q, topk),
76 + dtype=torch.int32,
77 + device=device)
78 +
79 + block_table = torch.zeros((batch_size, 128),
80 + dtype=torch.int32,
81 + device=device)
82 + out, lse = fm.flash_mla_with_kvcache(q,
83 + k_cache,
84 + block_table,
85 + cache_seqlens,
86 + head_dim_v,
87 + tile_md,
88 + num_splits,
89 + indices=indices,
90 + is_fp8_kvcache=True)
91 + assert out.shape[0] == batch_size
92 + assert out.shape[-1] == head_dim_v
93 + assert lse.shape[0] == batch_size
94 +
95 +
96 +def test_sparse_flashmla_prefill_smoke():
97 + import vllm.attention.ops.flashmla as fm
98 + ok, reason = fm.is_flashmla_supported()
99 + if not ok or not _cuda_sm90_available():
100 + pytest.skip(reason or "SM90 not available")
101 +
102 + device = torch.device("cuda")
103 + s_q = 1
104 + s_kv = 1
105 + h_q = 64 # kernel expects multiple of 64
106 + h_kv = 1
107 + d_qk = 576
108 + d_v = 512
109 + topk = 128
110 +
111 + q = torch.zeros((s_q, h_q, d_qk), dtype=torch.bfloat16, device=device)
112 + kv = torch.zeros((s_kv, h_kv, d_qk), dtype=torch.bfloat16, device=device)
113 + indices = torch.zeros((s_q, h_kv, topk), dtype=torch.int32, device=device)
114 +
115 + out, max_logits, lse = fm.flash_mla_sparse_prefill(q, kv, indices, 1.0,
116 + d_v)
117 + assert out.shape == (s_q, h_q, d_v)
118 + assert max_logits.shape == (s_q, h_q)
119 + assert lse.shape == (s_q, h_q)
tests/kernels/attention/test_pack_unpack_triton.py new
+245
@@ -0,0 +1,245 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +from torch.testing import assert_close
6 +
7 +from vllm.attention.ops.common import pack_seq_triton, unpack_seq_triton
8 +
9 +
10 +def test_pack_seq_basic_fp8():
11 + """Test basic functionality of pack_seq_triton with fp8 and 3D tensors."""
12 + device = "cuda"
13 + dtype = torch.float8_e4m3fn
14 +
15 + # Test cases with 3D tensors (N, H, D)
16 + test_cases = [
17 + (6, 8, 4, 2, [3, 3]), # (6, 8, 4) -> (2, 3, 8, 4)
18 + (10, 4, 8, 3, [2, 4, 4]), # (10, 4, 8) -> (3, 4, 4, 8)
19 + (20, 16, 32, 4, [5, 5, 5, 5]), # (20, 16, 32) -> (4, 5, 16, 32)
20 + ]
21 +
22 + for N, H, D, B, lengths_list in test_cases:
23 + # Create input tensor with small values for fp8
24 + x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1
25 + x = x.to(dtype=dtype)
26 + lengths = torch.tensor(lengths_list, device=device)
27 +
28 + # Pack the data
29 + packed = pack_seq_triton(x, lengths)
30 +
31 + # Check output shape and properties
32 + expected_shape = (B, max(lengths_list), H, D)
33 + assert packed.shape == expected_shape
34 + assert packed.dtype == dtype
35 + assert packed.device == x.device
36 +
37 + # Check that valid data is preserved (within fp8 precision)
38 + for b in range(B):
39 + start_idx = sum(lengths_list[:b])
40 + seq_len = lengths_list[b]
41 +
42 + expected_data = x[start_idx:start_idx + seq_len].to(torch.float32)
43 + actual_data = packed[b, :seq_len].to(torch.float32)
44 +
45 + assert_close(actual_data, expected_data, rtol=1e-1, atol=1e-2)
46 +
47 +
48 +def test_pack_seq_custom_padding_fp8():
49 + """Test pack_seq_triton with custom padding values for fp8."""
50 + device = "cuda"
51 + dtype = torch.float8_e4m3fn
52 + N, H, D, B = 20, 8, 16, 2
53 + lengths = torch.tensor([10, 10], device=device)
54 +
55 + x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1
56 + x = x.to(dtype=dtype)
57 +
58 + # Test with different padding values
59 + for pad_value in [-100.0, -10.0, 0.0, 10.0, 100.0]:
60 + result = pack_seq_triton(x, lengths, pad_value=pad_value)
61 +
62 + # Check valid data
63 + for b in range(B):
64 + start_idx = b * 10
65 + expected_data = x[start_idx:start_idx + 10].to(torch.float32)
66 + actual_data = result[b, :10].to(torch.float32)
67 + assert_close(actual_data, expected_data, rtol=1e-1, atol=1e-2)
68 +
69 + # Check padding (fp8 has limited range, so check for large values)
70 + padded_data = result[:, 10:].to(torch.float32)
71 + if pad_value < 0:
72 + assert torch.all(padded_data < -50) # Large negative values
73 + elif pad_value > 0:
74 + assert torch.all(padded_data > 50) # Large positive values
75 + else:
76 + assert torch.allclose(padded_data,
77 + torch.zeros_like(padded_data),
78 + atol=1e-2)
79 +
80 +
81 +def test_pack_seq_default_negative_inf_padding_fp8():
82 + """Test that pack_seq_triton uses -inf padding by default for fp8."""
83 + device = "cuda"
84 + dtype = torch.float8_e4m3fn
85 + # B = 2
86 + N, H, D = 20, 8, 16
87 + lengths = torch.tensor([10, 10], device=device)
88 +
89 + x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1
90 + x = x.to(dtype=dtype)
91 + result = pack_seq_triton(x, lengths)
92 +
93 + # Check that padding is large negative values (fp8 representation of -inf)
94 + padded_data = result[:, 10:].to(torch.float32)
95 + assert torch.all(
96 + padded_data < -100) # fp8 -inf is represented as large negative number
97 +
98 +
99 +def test_pack_seq_edge_cases_fp8():
100 + """Test pack_seq_triton with edge cases for fp8."""
101 + device = "cuda"
102 + dtype = torch.float8_e4m3fn
103 +
104 + # Test with single batch element
105 + x = torch.randn(10, 8, 16, dtype=torch.float32, device=device) * 0.1
106 + x = x.to(dtype=dtype)
107 + lengths = torch.tensor([10], device=device)
108 + result = pack_seq_triton(x, lengths)
109 + assert result.shape == (1, 10, 8, 16)
110 +
111 + # Test with very short sequences
112 + x = torch.randn(20, 4, 8, dtype=torch.float32, device=device) * 0.1
113 + x = x.to(dtype=dtype)
114 + lengths = torch.tensor([1, 1, 1], device=device)
115 + result = pack_seq_triton(x, lengths)
116 + assert result.shape == (3, 1, 4, 8)
117 +
118 + # Test with different sequence lengths
119 + x = torch.randn(15, 8, 16, dtype=torch.float32, device=device) * 0.1
120 + x = x.to(dtype=dtype)
121 + lengths = torch.tensor([5, 7, 3], device=device)
122 + result = pack_seq_triton(x, lengths)
123 + assert result.shape == (3, 7, 8, 16)
124 +
125 +
126 +def test_pack_seq_different_block_sizes_fp8():
127 + """Test pack_seq_triton with different block sizes for fp8."""
128 + device = "cuda"
129 + dtype = torch.float8_e4m3fn
130 + N, H, D, B = 100, 16, 32, 4
131 + lengths = torch.tensor([25, 25, 25, 25], device=device)
132 +
133 + x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1
134 + x = x.to(dtype=dtype)
135 +
136 + # Test different block sizes
137 + for block_t, block_d in [(32, 32), (64, 64), (128, 128)]:
138 + result = pack_seq_triton(x, lengths, block_t=block_t, block_d=block_d)
139 +
140 + assert result.shape == (B, 25, H, D)
141 +
142 + # Check that valid data is preserved (within fp8 precision)
143 + for b in range(B):
144 + start_idx = b * 25
145 + expected_data = x[start_idx:start_idx + 25].to(torch.float32)
146 + actual_data = result[b, :25].to(torch.float32)
147 + assert_close(actual_data, expected_data, rtol=1e-1, atol=1e-2)
148 +
149 +
150 +def test_pack_seq_shape_consistency():
151 + """Test that pack_seq_triton maintains shape consistency."""
152 + device = "cuda"
153 + dtype = torch.float8_e4m3fn
154 + N, H, D, B = 20, 8, 16, 2
155 + lengths = torch.tensor([10, 10], device=device)
156 +
157 + x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1
158 + x = x.to(dtype=dtype)
159 +
160 + result = pack_seq_triton(x, lengths)
161 +
162 + # Check shape consistency
163 + assert result.shape[0] == B # Batch dimension
164 + assert result.shape[1] == lengths.max().item() # Max sequence length
165 + assert result.shape[2:] == x.shape[1:] # Feature dimensions preserved
166 +
167 +
168 +def test_pack_unpack_roundtrip_fp8():
169 + """Test that pack -> unpack gives us back the original data for fp8."""
170 + device = "cuda"
171 + dtype = torch.float8_e4m3fn
172 +
173 + # Test cases with 3D tensors
174 + test_cases = [
175 + (6, 8, 4, 2, [3, 3]),
176 + (10, 4, 8, 3, [2, 4, 4]),
177 + (20, 16, 32, 4, [5, 5, 5, 5]),
178 + (15, 8, 16, 3, [7, 5, 3]),
179 + ]
180 +
181 + for N, H, D, B, lengths_list in test_cases:
182 + # Create input tensor with small values for fp8
183 + x = torch.randn(N, H, D, dtype=torch.float32, device=device) * 0.1
184 + x = x.to(dtype=dtype)
185 + lengths = torch.tensor(lengths_list, device=device)
186 +
187 + # Pack the data
188 + packed = pack_seq_triton(x, lengths)
189 +
190 + # Unpack the data
191 + unpacked = unpack_seq_triton(packed, lengths)
192 +
193 + # Check that we get back the original data (within fp8 precision)
194 + assert unpacked.shape == x.shape
195 + x_f32 = x.to(torch.float32)
196 + unpacked_f32 = unpacked.to(torch.float32)
197 + assert_close(x_f32, unpacked_f32, rtol=1e-3, atol=1e-3)
198 +
199 + # Unpack without explicit start locations (computed in kernel)
200 + unpacked_with_loc = unpack_seq_triton(packed, lengths)
201 + assert_close(x_f32,
202 + unpacked_with_loc.to(torch.float32),
203 + rtol=1e-3,
204 + atol=1e-2)
205 +
206 +
207 +def test_unpack_seq_triton_edge_cases_fp8():
208 + """Test unpack function with edge cases for fp8."""
209 + device = "cuda"
210 + dtype = torch.float8_e4m3fn
211 +
212 + # Test with single batch element
213 + x = torch.randn(10, 8, 16, dtype=torch.float32, device=device) * 0.1
214 + x = x.to(dtype=dtype)
215 + lengths = torch.tensor([10], device=device)
216 + packed = pack_seq_triton(x, lengths)
217 + unpacked = unpack_seq_triton(packed, lengths)
218 + assert unpacked.shape == x.shape
219 + assert_close(x.to(torch.float32),
220 + unpacked.to(torch.float32),
221 + rtol=1e-1,
222 + atol=1e-2)
223 +
224 + # Test with very short sequences
225 + x = torch.randn(20, 4, 8, dtype=torch.float32, device=device) * 0.1
226 + x = x.to(dtype=dtype)
227 + lengths = torch.tensor([1, 1, 1], device=device)
228 + packed = pack_seq_triton(x, lengths)
229 + unpacked = unpack_seq_triton(packed, lengths)
230 + # Only compare the first 3 elements that were actually packed
231 + assert_close(x[:3].to(torch.float32),
232 + unpacked.to(torch.float32),
233 + rtol=1e-1,
234 + atol=1e-2)
235 +
236 + x = torch.randn(15, 8, 16, dtype=torch.float32, device=device) * 0.1
237 + x = x.to(dtype=dtype)
238 + lengths = torch.tensor([5, 7, 3], device=device)
239 + packed = pack_seq_triton(x, lengths)
240 + unpacked = unpack_seq_triton(packed, lengths)
241 + assert unpacked.shape == x.shape
242 + assert_close(x.to(torch.float32),
243 + unpacked.to(torch.float32),
244 + rtol=1e-1,
245 + atol=1e-2)
tests/models/registry.py
+1
@@ -207,6 +207,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
207 trust_remote_code=True),
208 "DeepseekV3ForCausalLM": _HfExamplesInfo("deepseek-ai/DeepSeek-V3", # noqa: E501
209 trust_remote_code=True),
210 + "DeepseekV32ForCausalLM": _HfExamplesInfo("deepseek-ai/DeepSeek-V3.2-Exp"),
211 "Ernie4_5ForCausalLM": _HfExamplesInfo("baidu/ERNIE-4.5-0.3B-PT",
212 min_transformers_version="4.54"),
213 "Ernie4_5_MoeForCausalLM": _HfExamplesInfo("baidu/ERNIE-4.5-21B-A3B-PT",
tests/models/test_initialization.py
+6 -3
@@ -8,7 +8,8 @@ import pytest
8
9 from vllm import LLM
10 from vllm.utils import GiB_bytes
11 -from vllm.v1.core.kv_cache_utils import get_kv_cache_configs
11 +from vllm.v1.core.kv_cache_utils import (generate_scheduler_kv_cache_config,
12 + get_kv_cache_configs)
13 from vllm.v1.engine.core import EngineCore as V1EngineCore
14
15 from ..utils import create_new_process_for_each_test
@@ -62,11 +63,13 @@ def can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch,
63 # Avoid calling model.forward()
64 def _initialize_kv_caches_v1(self, vllm_config):
65 kv_cache_specs = self.model_executor.get_kv_cache_specs()
65 - scheduler_kv_cache_config = get_kv_cache_configs(
66 + kv_cache_configs = get_kv_cache_configs(
67 vllm_config,
68 kv_cache_specs,
69 [10 * GiB_bytes],
69 - )[0]
70 + )
71 + scheduler_kv_cache_config = generate_scheduler_kv_cache_config(
72 + kv_cache_configs)
73
74 # gpu_blocks (> 0), cpu_blocks, scheduler_kv_cache_config
75 return 1, 0, scheduler_kv_cache_config
tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_platform.py
+1 -1
@@ -26,5 +26,5 @@ class DummyPlatform(Platform):
26
27 def get_attn_backend_cls(self, backend_name, head_size, dtype,
28 kv_cache_dtype, block_size, use_v1, use_mla,
29 - has_sink):
29 + has_sink, use_sparse):
30 return "vllm_add_dummy_platform.dummy_attention_backend.DummyAttentionBackend" # noqa E501
tests/v1/attention/test_mla_backends.py
+58 -11
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for v1 MLA backends without GPUModelRunner dependency."""
4 +from typing import Optional, Union
5
6 import pytest
7 import torch
@@ -10,6 +11,7 @@ from tests.v1.attention.utils import (BatchSpec, _Backend,
11 create_standard_kv_cache_spec,
12 create_vllm_config,
13 get_attention_backend)
14 +from vllm import _custom_ops as ops
15 from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv
16 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
17 from vllm.v1.kv_cache_interface import FullAttentionSpec
@@ -78,7 +80,9 @@ def create_and_prepopulate_kv_cache(
80 device: torch.device,
81 num_blocks: int,
82 common_attn_metadata: CommonAttentionMetadata,
81 - randomize_blocks: bool = True) -> torch.Tensor:
83 + randomize_blocks: bool = True,
84 + kv_cache_dtype: Optional[str] = None,
85 + scale: Union[float, torch.Tensor] = 1.0) -> torch.Tensor:
86 """Create and prepopulate an MLA KV cache with context data.
87
88 Args:
@@ -93,6 +97,11 @@ def create_and_prepopulate_kv_cache(
97 common_attn_metadata: Common attention metadata
98 randomize_blocks: Whether to randomly permute blocks
99 or use sequential order
100 + kv_cache_dtype: Optional kv cache dtype string. When set to
101 + "fp8_ds_mla" the cache is populated using the
102 + fp8 DeepSeek MLA layout via concat_and_cache_mla.
103 + scale: Scaling factor forwarded to concat_and_cache_mla when the
104 + fp8 cache layout is requested.
105
106 Returns:
107 MLA KV cache tensor
@@ -105,23 +114,61 @@ def create_and_prepopulate_kv_cache(
114 block_table = common_attn_metadata.block_table_tensor
115 slot_mapping = common_attn_metadata.slot_mapping
116
108 - # Create MLA KV cache: (num_blocks, block_size, head_size)
109 - kv_cache = torch.empty(num_blocks,
110 - block_size,
111 - head_size,
112 - dtype=dtype,
113 - device=device)
114 - kv_cache_flat = kv_cache.view(-1, head_size)
117 + use_fp8_ds_mla = kv_cache_dtype == "fp8_ds_mla"
118 +
119 + if use_fp8_ds_mla:
120 + if not kv_c_contexts:
121 + raise ValueError("kv_c_contexts cannot be empty when using"
122 + " fp8_ds_mla cache dtype")
123 + kv_lora_rank = kv_c_contexts[0].shape[-1]
124 + rope_dim = k_pe_contexts[0].shape[-1]
125 + entry_size = kv_lora_rank + 4 * 4 + 2 * rope_dim
126 + kv_cache = torch.zeros(num_blocks,
127 + block_size,
128 + entry_size,
129 + dtype=torch.uint8,
130 + device=device)
131 + scale_tensor = (scale
132 + if isinstance(scale, torch.Tensor) else torch.tensor(
133 + scale, dtype=torch.float32, device=device))
134 + scale_tensor = scale_tensor.to(device=device, dtype=torch.float32)
135 + else:
136 + # Create MLA KV cache: (num_blocks, block_size, head_size)
137 + kv_cache = torch.empty(num_blocks,
138 + block_size,
139 + head_size,
140 + dtype=dtype,
141 + device=device)
142 + kv_cache_flat = kv_cache.view(-1, head_size)
143
144 # Populate the cache with the context tokens
145 # Start from block_id=1 since block_id=0 is considered the null block
146 start_block_idx = 1
147 for i in range(batch_size):
148 kv_c_context, k_pe_context = kv_c_contexts[i], k_pe_contexts[i]
121 - kv_context = torch.cat([kv_c_context, k_pe_context.squeeze(1)], dim=-1)
149 + context_len = kv_c_context.shape[0]
150 + if context_len == 0:
151 + start_block_idx += cdiv(int(seq_lens[i]), block_size)
152 + continue
153 +
154 start = start_block_idx * block_size
123 - end = start + kv_context.shape[0]
124 - kv_cache_flat[start:end, ...] = kv_context
155 +
156 + if use_fp8_ds_mla:
157 + slots = torch.arange(context_len, device=device,
158 + dtype=torch.long) + start
159 + ops.concat_and_cache_mla(
160 + kv_c_context,
161 + k_pe_context.squeeze(1),
162 + kv_cache,
163 + slots,
164 + kv_cache_dtype="fp8_ds_mla",
165 + scale=scale_tensor,
166 + )
167 + else:
168 + kv_context = torch.cat(
169 + [kv_c_context, k_pe_context.squeeze(1)], dim=-1)
170 + end = start + kv_context.shape[0]
171 + kv_cache_flat[start:end, ...] = kv_context
172
173 # Stay block aligned and allocate enough blocks for the new tokens
174 start_block_idx += cdiv(int(seq_lens[i]), block_size)
tests/v1/attention/test_sparse_mla_backends.py new
+426
@@ -0,0 +1,426 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for the FlashMLA sparse backend utilities."""
4 +
5 +import math
6 +from types import MethodType, SimpleNamespace
7 +
8 +import numpy as np
9 +import pytest
10 +import torch
11 +
12 +from tests.v1.attention.test_mla_backends import (
13 + BATCH_SPECS, BatchSpec, MockAttentionLayer,
14 + create_and_prepopulate_kv_cache)
15 +from tests.v1.attention.utils import (create_common_attn_metadata,
16 + create_standard_kv_cache_spec,
17 + create_vllm_config)
18 +from vllm import _custom_ops as ops
19 +from vllm.attention.ops import flashmla
20 +from vllm.model_executor.layers.linear import ColumnParallelLinear
21 +from vllm.utils import cdiv
22 +from vllm.v1.attention.backends.mla.flashmla_sparse import (
23 + FlashMLASparseBackend, FlashMLASparseDecodeAndContextMetadata,
24 + FlashMLASparseImpl, FlashMLASparseMetadata)
25 +
26 +SPARSE_BACKEND_BATCH_SPECS = {
27 + name: BATCH_SPECS[name]
28 + for name in [
29 + "mixed_small",
30 + "mixed_medium",
31 + "small_prefill",
32 + "medium_prefill",
33 + "single_prefill",
34 + ]
35 +}
36 +
37 +SPARSE_BACKEND_BATCH_SPECS["large_q_prefill"] = BatchSpec(seq_lens=[1024] * 2,
38 + query_lens=[256] * 2)
39 +SPARSE_BACKEND_BATCH_SPECS["large_q_pure_prefill"] = BatchSpec(
40 + seq_lens=[256] * 2, query_lens=[256] * 2)
41 +
42 +
43 +def _dequantize_fp8_ds_mla_entry(
44 + cache_slice: torch.Tensor, kv_lora_rank: int, rope_dim: int,
45 + dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
46 + """Dequantize a single fp8_ds_mla cache entry back to latent + rope."""
47 +
48 + # The first kv_lora_rank bytes store FP8 latent values with one scale per
49 + # 128 element tile written as float32 right after the latent payload.
50 + scales = cache_slice.view(torch.float32)[kv_lora_rank //
51 + 4:kv_lora_rank // 4 + 4]
52 + latent = torch.empty(kv_lora_rank,
53 + dtype=torch.float16,
54 + device=cache_slice.device)
55 + for tile_idx in range(4):
56 + tile_start = tile_idx * 128
57 + tile_end = tile_start + 128
58 + ops.convert_fp8(latent[tile_start:tile_end],
59 + cache_slice[tile_start:tile_end],
60 + float(scales[tile_idx].item()),
61 + kv_dtype="fp8")
62 + latent = latent.to(dtype)
63 +
64 + rope_offset = kv_lora_rank // 2 + 8
65 + rope_vals = cache_slice.view(dtype)[rope_offset:rope_offset + rope_dim]
66 + return latent, rope_vals.clone()
67 +
68 +
69 +def _quantize_dequantize_fp8_ds_mla(
70 + kv_c: torch.Tensor, k_pe: torch.Tensor, block_size: int,
71 + scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
72 + """Round-trip kv_c/k_pe though the fp8_ds_mla cache layout."""
73 +
74 + if kv_c.numel() == 0:
75 + return kv_c.clone(), k_pe.clone()
76 +
77 + kv_lora_rank = kv_c.shape[-1]
78 + rope_dim = k_pe.shape[-1]
79 + num_tokens = kv_c.shape[0]
80 + num_blocks = max(1, math.ceil(num_tokens / block_size))
81 + entry_size = kv_lora_rank + 4 * 4 + 2 * rope_dim
82 +
83 + tmp_cache = torch.zeros(num_blocks,
84 + block_size,
85 + entry_size,
86 + dtype=torch.uint8,
87 + device=kv_c.device)
88 + slot_mapping = torch.arange(num_tokens,
89 + dtype=torch.long,
90 + device=kv_c.device)
91 +
92 + ops.concat_and_cache_mla(kv_c,
93 + k_pe,
94 + tmp_cache,
95 + slot_mapping,
96 + kv_cache_dtype="fp8_ds_mla",
97 + scale=scale)
98 +
99 + dequant_kv_c = torch.empty_like(kv_c)
100 + dequant_k_pe = torch.empty_like(k_pe)
101 +
102 + for token_idx in range(num_tokens):
103 + slot = slot_mapping[token_idx].item()
104 + block_idx = slot // block_size
105 + block_offset = slot % block_size
106 + cache_slice = tmp_cache[block_idx, block_offset]
107 + latent, rope_vals = _dequantize_fp8_ds_mla_entry(
108 + cache_slice, kv_lora_rank, rope_dim, kv_c.dtype)
109 + dequant_kv_c[token_idx] = latent
110 + dequant_k_pe[token_idx] = rope_vals
111 +
112 + return dequant_kv_c, dequant_k_pe
113 +
114 +
115 +def test_sparse_backend_metadata_registration():
116 + backend = FlashMLASparseBackend
117 +
118 + assert backend.get_name() == "FLASHMLA_SPARSE_VLLM_V1"
119 + assert backend.get_metadata_cls() is FlashMLASparseMetadata
120 + assert backend.get_impl_cls() is FlashMLASparseImpl
121 +
122 + dtype_list = backend.get_supported_dtypes()
123 + assert torch.bfloat16 in dtype_list
124 +
125 + shape = backend.get_kv_cache_shape(num_blocks=2,
126 + block_size=64,
127 + num_kv_heads=1,
128 + head_size=576)
129 + assert shape == (2, 64, 576)
130 +
131 +
132 +def test_sparse_decode_metadata_filters_prefill_indices():
133 + prefill_context_lengths = torch.tensor([4, 2], dtype=torch.int32)
134 + metadata = FlashMLASparseDecodeAndContextMetadata(
135 + scheduler_metadata=torch.tensor([[0]], dtype=torch.int32),
136 + num_splits=torch.tensor([1, 1], dtype=torch.int32),
137 + cache_lens=torch.tensor([10, 12], dtype=torch.int32),
138 + prefill_context_lengths=prefill_context_lengths,
139 + )
140 +
141 + indices = torch.tensor([[0, 3, 5], [1, 2, 4]], dtype=torch.int32)
142 +
143 + context_indices, new_token_indices = metadata.filter_prefill_indices(
144 + indices)
145 +
146 + expected_context = torch.tensor([[-1, -1, 5], [-1, -1, 4]],
147 + dtype=torch.int32)
148 + expected_new_tokens = torch.tensor([[-1, -1, 1], [-1, 0, 2]],
149 + dtype=torch.int32)
150 +
151 + assert torch.equal(context_indices, expected_context)
152 + assert torch.equal(new_token_indices, expected_new_tokens)
153 +
154 +
155 +def test_sparse_impl_zero_fills_when_metadata_missing():
156 + impl = FlashMLASparseImpl.__new__(FlashMLASparseImpl)
157 + dummy_layer = object()
158 + q = torch.zeros((2, 1, 3))
159 + k_c = torch.zeros((2, 3))
160 + k_pe = torch.zeros((2, 1, 1))
161 + kv_cache = torch.zeros((1, 1, 1))
162 + output = torch.ones((2, 4))
163 +
164 + result = FlashMLASparseImpl.forward(impl,
165 + dummy_layer,
166 + q,
167 + k_c,
168 + k_pe,
169 + kv_cache,
170 + attn_metadata=None,
171 + output=output)
172 +
173 + assert result is output
174 + assert torch.all(result == 0)
175 +
176 +
177 +@pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys()))
178 +@pytest.mark.parametrize("kv_cache_dtype", ["fp8_ds_mla", "auto"])
179 +def test_sparse_backend_decode_correctness(dist_init, batch_name,
180 + kv_cache_dtype):
181 + if not torch.cuda.is_available():
182 + pytest.skip("CUDA is required for sparse MLA decode test")
183 +
184 + device = torch.device("cuda")
185 + dtype = torch.bfloat16
186 +
187 + batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name]
188 +
189 + # Model hyper-parameters (kept intentionally small for the unit test)
190 + num_heads = 128
191 + kv_lora_rank = 512
192 + qk_nope_head_dim = 128
193 + qk_rope_head_dim = 64
194 + v_head_dim = 128
195 + head_size = kv_lora_rank + qk_rope_head_dim
196 + topk_tokens = 2048
197 +
198 + max_seqlen = max(batch_spec.seq_lens)
199 + total_cache_tokens = sum(batch_spec.seq_lens)
200 + block_size = 64
201 +
202 + vllm_config = create_vllm_config(
203 + model_name="deepseek-ai/DeepSeek-V2-Lite-Chat",
204 + max_model_len=max_seqlen,
205 + num_gpu_blocks=max(2048,
206 + cdiv(total_cache_tokens, block_size) + 1),
207 + block_size=block_size)
208 + model_config = vllm_config.model_config
209 + model_config.hf_config = SimpleNamespace(
210 + attn_module_list_cfg=[{
211 + "topk_tokens": topk_tokens
212 + }])
213 + model_config.hf_text_config = SimpleNamespace(
214 + q_lora_rank=None,
215 + kv_lora_rank=kv_lora_rank,
216 + qk_nope_head_dim=qk_nope_head_dim,
217 + qk_rope_head_dim=qk_rope_head_dim,
218 + v_head_dim=v_head_dim,
219 + model_type="deepseek_v2",
220 + )
221 + model_config.dtype = dtype
222 + model_config.get_num_attention_heads = MethodType(
223 + lambda self, parallel_config: num_heads, model_config)
224 + model_config.get_num_kv_heads = MethodType(lambda self, parallel_config: 1,
225 + model_config)
226 + model_config.get_head_size = MethodType(lambda self: head_size,
227 + model_config)
228 + model_config.get_sliding_window = MethodType(lambda self: None,
229 + model_config)
230 +
231 + kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
232 +
233 + torch.manual_seed(0)
234 +
235 + scale = 1.0 / math.sqrt(head_size)
236 +
237 + # Shared MLA projection weights to keep reference and backend in sync
238 + W_UK = torch.randn(kv_lora_rank,
239 + num_heads,
240 + qk_nope_head_dim,
241 + dtype=dtype,
242 + device=device)
243 + W_UV = torch.randn(kv_lora_rank,
244 + num_heads,
245 + v_head_dim,
246 + dtype=dtype,
247 + device=device)
248 +
249 + # Build synthetic decode-only workload
250 + seq_lens = batch_spec.seq_lens
251 + query_lens = batch_spec.query_lens
252 +
253 + all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []
254 + kv_c_contexts, k_pe_contexts = [], []
255 + reference_outputs = []
256 +
257 + kv_cache_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
258 +
259 + for i in range(batch_spec.batch_size):
260 + s_len = seq_lens[i]
261 + q_len = query_lens[i]
262 + ctx_len = s_len - q_len
263 +
264 + q_c = torch.rand(q_len,
265 + num_heads,
266 + qk_nope_head_dim + qk_rope_head_dim,
267 + dtype=dtype,
268 + device=device)
269 + kv_c_full = torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device)
270 + k_pe_full = torch.rand(s_len,
271 + 1,
272 + qk_rope_head_dim,
273 + dtype=dtype,
274 + device=device)
275 +
276 + kv_c_full, k_pe_full = _quantize_dequantize_fp8_ds_mla(
277 + kv_c_full,
278 + k_pe_full.squeeze(1),
279 + block_size=vllm_config.cache_config.block_size,
280 + scale=kv_cache_scale,
281 + )
282 +
283 + q_nope, q_pe = q_c.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)
284 + ql_nope = torch.einsum("qnh,lnh->qnl", q_nope, W_UK)
285 + q_mqa = torch.cat([ql_nope, q_pe], dim=-1)
286 +
287 + k_mqa = torch.cat([kv_c_full, k_pe_full], dim=-1)
288 + k_mqa = k_mqa.unsqueeze(1).expand(-1, num_heads, -1)
289 + v_mqa = kv_c_full.unsqueeze(1).expand(-1, num_heads, -1)
290 +
291 + attn_mask = torch.ones(q_len, s_len, dtype=torch.bool, device=device)
292 + causal_mask = torch.tril(torch.ones(q_len, q_len, device=device))
293 + attn_mask[:, ctx_len:] = causal_mask
294 +
295 + q_sdpa_in = q_mqa.unsqueeze(0).transpose(1, 2)
296 + k_sdpa_in = k_mqa.unsqueeze(0).transpose(1, 2)
297 + v_sdpa_in = v_mqa.unsqueeze(0).transpose(1, 2)
298 +
299 + sdpa_out = torch.nn.functional.scaled_dot_product_attention(
300 + q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=scale)
301 + sdpa_out = sdpa_out.transpose(1, 2).squeeze(0)
302 +
303 + sdpa_out = torch.einsum("qnl,lnv->qnv", sdpa_out, W_UV)
304 + reference_outputs.append(sdpa_out.flatten(start_dim=-2))
305 +
306 + all_q_vllm.append(q_c)
307 + all_kv_c_vllm.append(kv_c_full[ctx_len:])
308 + all_k_pe_vllm.append(k_pe_full[ctx_len:])
309 + kv_c_contexts.append(kv_c_full[:ctx_len + 1])
310 + k_pe_contexts.append(k_pe_full[:ctx_len + 1])
311 +
312 + query_vllm = torch.cat(all_q_vllm, dim=0)
313 + kv_c_vllm = torch.cat(all_kv_c_vllm, dim=0)
314 + k_pe_vllm = torch.cat(all_k_pe_vllm, dim=0)
315 + sdpa_reference = torch.cat(reference_outputs, dim=0)
316 +
317 + vllm_config.cache_config.cache_dtype = kv_cache_dtype
318 +
319 + common_attn_metadata = create_common_attn_metadata(
320 + batch_spec,
321 + vllm_config.cache_config.block_size,
322 + device,
323 + arange_block_indices=True)
324 +
325 + kv_cache = create_and_prepopulate_kv_cache(
326 + kv_c_contexts=kv_c_contexts,
327 + k_pe_contexts=k_pe_contexts,
328 + block_size=vllm_config.cache_config.block_size,
329 + head_size=head_size,
330 + dtype=dtype,
331 + device=device,
332 + num_blocks=vllm_config.cache_config.num_gpu_blocks,
333 + common_attn_metadata=common_attn_metadata,
334 + randomize_blocks=False,
335 + kv_cache_dtype=vllm_config.cache_config.cache_dtype,
336 + scale=kv_cache_scale,
337 + )
338 +
339 + builder_cls = FlashMLASparseBackend.get_builder_cls()
340 + builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device)
341 + metadata = builder.build(common_prefix_len=0,
342 + common_attn_metadata=common_attn_metadata)
343 +
344 + starts = np.asarray(common_attn_metadata.query_start_loc_cpu,
345 + dtype=np.int32)
346 + seg_lengths = np.diff(starts)
347 + positions = np.arange(starts[-1], dtype=np.int32) - np.repeat(
348 + starts[:-1], seg_lengths)
349 + seq_lengths = np.asarray(common_attn_metadata.seq_lens_cpu, dtype=np.int32)
350 + prefix_lengths = seq_lengths - seg_lengths
351 + positions += np.repeat(prefix_lengths, seg_lengths)
352 +
353 + pos_gpu = torch.as_tensor(positions, device=device, dtype=torch.int32)
354 + topk = metadata.topk_tokens
355 + debug_indices = torch.arange(topk, device=device,
356 + dtype=torch.int32).unsqueeze(0)
357 + token_positions = pos_gpu.unsqueeze(1)
358 + causal_mask = (debug_indices <= token_positions)
359 + debug_indices = torch.where(causal_mask, debug_indices,
360 + torch.full_like(debug_indices, -1))
361 +
362 + # FlashMLASparseImpl now reads top-k indices from the indexer-provided
363 + # buffer, so emulate that contract with a simple namespace mock.
364 + debug_indices = debug_indices.expand(metadata.num_actual_tokens,
365 + -1).clone()
366 + mock_indexer = SimpleNamespace(topk_indices_buffer=debug_indices)
367 +
368 + ok, reason = flashmla.is_flashmla_supported()
369 + if not ok:
370 + pytest.skip(reason)
371 +
372 + kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1)
373 + kv_b_proj_weight = kv_b_proj_weight.view(
374 + kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim))
375 +
376 + mock_kv_b_proj = ColumnParallelLinear(input_size=kv_lora_rank,
377 + output_size=num_heads *
378 + (qk_nope_head_dim + v_head_dim),
379 + bias=False).to(device=device,
380 + dtype=dtype)
381 + mock_kv_b_proj.weight = torch.nn.Parameter(kv_b_proj_weight.T.contiguous())
382 +
383 + impl_cls = FlashMLASparseBackend.get_impl_cls()
384 + impl = impl_cls(num_heads=num_heads,
385 + head_size=head_size,
386 + scale=scale,
387 + num_kv_heads=1,
388 + alibi_slopes=None,
389 + sliding_window=None,
390 + kv_cache_dtype=vllm_config.cache_config.cache_dtype,
391 + logits_soft_cap=None,
392 + attn_type="decoder",
393 + kv_sharing_target_layer_name=None,
394 + q_lora_rank=None,
395 + kv_lora_rank=kv_lora_rank,
396 + qk_nope_head_dim=qk_nope_head_dim,
397 + qk_rope_head_dim=qk_rope_head_dim,
398 + qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,
399 + v_head_dim=v_head_dim,
400 + kv_b_proj=mock_kv_b_proj,
401 + indexer=mock_indexer)
402 +
403 + impl.process_weights_after_loading(dtype)
404 +
405 + layer = MockAttentionLayer(device)
406 + out_buffer = torch.empty(metadata.num_actual_tokens,
407 + num_heads * v_head_dim,
408 + dtype=dtype,
409 + device=device)
410 +
411 + backend_output = impl.forward(layer,
412 + query_vllm,
413 + kv_c_vllm,
414 + k_pe_vllm,
415 + kv_cache,
416 + metadata,
417 + output=out_buffer)
418 +
419 + assert backend_output.shape == sdpa_reference.shape
420 + assert backend_output.dtype == sdpa_reference.dtype
421 + assert torch.isfinite(backend_output).all()
422 +
423 + torch.testing.assert_close(backend_output,
424 + sdpa_reference,
425 + rtol=0.5,
426 + atol=0.5)
tests/v1/attention/utils.py
-1
@@ -168,7 +168,6 @@ def create_standard_kv_cache_spec(
168 vllm_config.parallel_config),
169 head_size=vllm_config.model_config.get_head_size(),
170 dtype=vllm_config.model_config.dtype,
171 - use_mla=vllm_config.model_config.use_mla,
171 sliding_window=vllm_config.model_config.get_sliding_window(),
172 )
173
tests/v1/core/test_kv_cache_utils.py
+47 -9
@@ -24,7 +24,8 @@ from vllm.v1.core.kv_cache_utils import (
24 make_block_hash_with_group_id)
25 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
26 KVCacheGroupSpec, KVCacheSpec,
27 - KVCacheTensor, SlidingWindowSpec,
27 + KVCacheTensor, MLAAttentionSpec,
28 + SlidingWindowSpec,
29 UniformTypeKVCacheSpecs)
30 from vllm.v1.metrics.stats import PrefixCacheStats
31 from vllm.v1.request import Request
@@ -77,13 +78,11 @@ def new_kv_cache_spec(block_size=16,
78 num_kv_heads=2,
79 head_size=64,
80 dtype=torch.float32,
80 - use_mla=False,
81 sliding_window=None):
82 return FullAttentionSpec(block_size=block_size,
83 num_kv_heads=num_kv_heads,
84 head_size=head_size,
85 dtype=dtype,
86 - use_mla=use_mla,
86 sliding_window=sliding_window)
87
88
@@ -91,13 +90,11 @@ def new_sliding_window_spec(block_size=16,
90 num_kv_heads=2,
91 head_size=64,
92 dtype=torch.float32,
94 - use_mla=False,
93 sliding_window=1):
94 return SlidingWindowSpec(block_size=block_size,
95 num_kv_heads=num_kv_heads,
96 head_size=head_size,
97 dtype=dtype,
100 - use_mla=use_mla,
98 sliding_window=sliding_window)
99
100
@@ -894,7 +891,6 @@ def test_merge_kv_cache_spec():
891 num_kv_heads=full_spec.num_kv_heads,
892 head_size=full_spec.head_size,
893 dtype=full_spec.dtype,
897 - use_mla=full_spec.use_mla,
894 sliding_window=1,
895 ),
896 ]
@@ -991,7 +987,6 @@ def test_estimate_max_model_len(model_id, max_model_len,
987 num_kv_heads=32,
988 head_size=128,
989 dtype=torch.float16,
994 - use_mla=False,
990 )
991 # Estimate the maximum model length, 16384 model_len need 8GB
992 estimated_max_len = estimate_max_model_len(vllm_config, kv_cache_spec,
@@ -1022,7 +1017,6 @@ def test_get_max_concurrency_for_kv_cache_config():
1017 num_kv_heads=32,
1018 head_size=128,
1019 dtype=torch.float16,
1025 - use_mla=False,
1020 )
1021
1022 sliding_window_spec = SlidingWindowSpec(
@@ -1030,7 +1024,6 @@ def test_get_max_concurrency_for_kv_cache_config():
1024 num_kv_heads=32,
1025 head_size=128,
1026 dtype=torch.float16,
1033 - use_mla=False,
1027 sliding_window=1024,
1028 )
1029
@@ -1412,3 +1405,48 @@ def test_generate_scheduler_kv_cache_config():
1405 KVCacheGroupSpec(['layer_1', 'layer_2'], new_kv_cache_spec())
1406 ],
1407 )
1408 +
1409 +
1410 +def new_mla_spec(cache_dtype_str=None):
1411 + return MLAAttentionSpec(block_size=16,
1412 + num_kv_heads=16,
1413 + head_size=64,
1414 + dtype=torch.float32,
1415 + cache_dtype_str=cache_dtype_str)
1416 +
1417 +
1418 +def test_merge_mla_spec():
1419 + kv_cache_specs = [
1420 + new_mla_spec(),
1421 + new_mla_spec(),
1422 + ]
1423 + mla_spec = kv_cache_specs[0].merge(kv_cache_specs)
1424 + assert mla_spec == new_mla_spec()
1425 +
1426 + kv_cache_specs = [
1427 + new_mla_spec(cache_dtype_str="fp8_ds_mla"),
1428 + new_mla_spec(cache_dtype_str="fp8_ds_mla"),
1429 + ]
1430 + mla_spec = kv_cache_specs[0].merge(kv_cache_specs)
1431 + assert mla_spec == new_mla_spec(cache_dtype_str="fp8_ds_mla")
1432 +
1433 + kv_cache_specs = [
1434 + new_mla_spec(cache_dtype_str="fp8_ds_mla"),
1435 + new_mla_spec(cache_dtype_str=None),
1436 + ]
1437 + with pytest.raises(AssertionError):
1438 + kv_cache_specs[0].merge(kv_cache_specs)
1439 +
1440 + kv_cache_specs = [
1441 + new_kv_cache_spec(),
1442 + new_mla_spec(),
1443 + ]
1444 + with pytest.raises(AssertionError):
1445 + kv_cache_specs[0].merge(kv_cache_specs)
1446 +
1447 + kv_cache_specs = [
1448 + new_mla_spec(cache_dtype_str="fp8_ds_mla"),
1449 + new_kv_cache_spec(),
1450 + ]
1451 + with pytest.raises(AssertionError):
1452 + kv_cache_specs[0].merge(kv_cache_specs)
tests/v1/core/test_prefix_caching.py
+2 -5
@@ -76,7 +76,7 @@ def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig:
76 kv_cache_groups=[
77 KVCacheGroupSpec(
78 ["layer"],
79 - FullAttentionSpec(block_size, 1, 1, torch.float32, False),
79 + FullAttentionSpec(block_size, 1, 1, torch.float32),
80 )
81 ],
82 )
@@ -90,7 +90,7 @@ def make_kv_cache_config_hybrid_model(block_size: int,
90 kv_cache_groups=[
91 KVCacheGroupSpec(
92 ["layer1"],
93 - FullAttentionSpec(block_size, 1, 1, torch.float32, False),
93 + FullAttentionSpec(block_size, 1, 1, torch.float32),
94 ),
95 KVCacheGroupSpec(
96 ["layer2"],
@@ -98,7 +98,6 @@ def make_kv_cache_config_hybrid_model(block_size: int,
98 1,
99 1,
100 torch.float32,
101 - False,
101 sliding_window=2 * block_size),
102 ),
103 KVCacheGroupSpec(
@@ -107,7 +106,6 @@ def make_kv_cache_config_hybrid_model(block_size: int,
106 1,
107 1,
108 torch.float32,
110 - False,
109 sliding_window=2 * block_size),
110 ),
111 ],
@@ -1338,7 +1336,6 @@ def test_eagle_with_sliding_window():
1336 head_size=1,
1337 dtype=torch.float32,
1338 sliding_window=block_size,
1341 - use_mla=False,
1339 )
1340 manager = KVCacheManager(
1341 KVCacheConfig(
tests/v1/core/test_single_type_kv_cache_manager.py
-6
@@ -35,7 +35,6 @@ def test_chunked_local_attention_possible_cached_prefix():
35 head_size=1,
36 dtype=torch.float32,
37 attention_chunk_size=4,
38 - use_mla=False,
38 )
39
40 block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
@@ -100,7 +99,6 @@ def test_sliding_window_possible_cached_prefix():
99 head_size=1,
100 dtype=torch.float32,
101 sliding_window=4,
103 - use_mla=False,
102 )
103
104 block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
@@ -165,7 +163,6 @@ def test_chunked_local_attention_remove_skipped_blocks():
163 head_size=1,
164 dtype=torch.float32,
165 attention_chunk_size=4,
168 - use_mla=False,
166 )
167
168 block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True)
@@ -217,7 +214,6 @@ def test_sliding_window_remove_skipped_blocks():
214 head_size=1,
215 dtype=torch.float32,
216 sliding_window=4,
220 - use_mla=False,
217 )
218
219 block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True)
@@ -285,7 +281,6 @@ def test_get_num_blocks_to_allocate():
281 head_size=1,
282 dtype=torch.float32,
283 sliding_window=4, # Placeholder value, not related to test result
288 - use_mla=False,
284 )
285
286 block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
@@ -308,7 +303,6 @@ def test_chunked_local_attention_get_num_blocks_to_allocate():
303 head_size=1,
304 dtype=torch.float32,
305 attention_chunk_size=4, # Placeholder value, not related to test result
311 - use_mla=False,
306 )
307
308 block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
tests/v1/engine/test_engine_core_client.py
+1 -2
@@ -836,8 +836,7 @@ def test_engine_core_proc_instantiation_cuda_empty(
836 mock_spec = FullAttentionSpec(block_size=16,
837 num_kv_heads=1,
838 head_size=64,
839 - dtype=torch.float16,
840 - use_mla=False)
839 + dtype=torch.float16)
840
841 mock_executor.get_kv_cache_specs.return_value = [{
842 "default": mock_spec
tests/v1/worker/test_gpu_model_runner.py
-1
@@ -39,7 +39,6 @@ def initialize_kv_cache(runner: GPUModelRunner):
39 runner.parallel_config),
40 head_size=runner.model_config.get_head_size(),
41 dtype=runner.kv_cache_dtype,
42 - use_mla=False,
42 )
43 tensor_size = attn_spec.page_size_bytes * NUM_BLOCKS
44 kv_cache_config = KVCacheConfig(
vllm/_custom_ops.py
+9
@@ -1678,6 +1678,15 @@ def cp_gather_cache(src_cache: torch.Tensor,
1678 cu_seq_lens, batch_size, seq_starts)
1679
1680
1681 +def indexer_k_quant_and_cache(k: torch.Tensor, kv_cache: torch.Tensor,
1682 + slot_mapping: torch.Tensor,
1683 + quant_block_size: int,
1684 + kv_cache_dtype: str) -> None:
1685 + torch.ops._C_cache_ops.indexer_k_quant_and_cache(k, kv_cache, slot_mapping,
1686 + quant_block_size,
1687 + kv_cache_dtype)
1688 +
1689 +
1690 def get_device_attribute(attribute: int, device: int) -> int:
1691 return torch.ops._C_cuda_utils.get_device_attribute(attribute, device)
1692
vllm/attention/backends/abstract.py
+1
@@ -70,6 +70,7 @@ class AttentionBackend(ABC):
70 block_size: int,
71 num_kv_heads: int,
72 head_size: int,
73 + cache_dtype_str: str = "auto",
74 ) -> Tuple[int, ...]:
75 raise NotImplementedError
76
vllm/attention/layer.py
+4 -1
@@ -94,6 +94,7 @@ class Attention(nn.Module, AttentionLayerBase):
94 logits_soft_cap: Optional[float] = None,
95 per_layer_sliding_window: Optional[int] = None,
96 use_mla: bool = False,
97 + use_sparse: bool = False,
98 prefix: str = "",
99 attn_type: str = AttentionType.DECODER,
100 kv_sharing_target_layer_name: Optional[str] = None,
@@ -154,6 +155,7 @@ class Attention(nn.Module, AttentionLayerBase):
155 self._o_scale_float: Optional[float] = None
156
157 self.use_mla = use_mla
158 + self.use_sparse = use_sparse
159 self.num_heads = num_heads
160 self.head_size = head_size
161 self.num_kv_heads = num_kv_heads
@@ -186,7 +188,8 @@ class Attention(nn.Module, AttentionLayerBase):
188 kv_cache_dtype,
189 block_size,
190 use_mla=use_mla,
189 - has_sink=self.has_sink)
191 + has_sink=self.has_sink,
192 + use_sparse=use_sparse)
193 else:
194 self.attn_backend = attn_backend
195
vllm/attention/ops/common.py
+205
@@ -138,3 +138,208 @@ def cp_lse_ag_out_rs(cp_attn_out: torch.Tensor,
138 out, _ = correct_attn_out(cp_attn_out, lses, cp_group.rank_in_group, ctx)
139 out = cp_group.reduce_scatter(out, dim=1)
140 return out
141 +
142 +
143 +@triton.jit
144 +def _pack_seq_kernel(
145 + x_ptr, # [N, D]
146 + out_ptr, # [B, Lmax, D]
147 + lengths_ptr, # *i32, [B]
148 + N: tl.constexpr,
149 + D: tl.constexpr,
150 + Lmax: tl.constexpr,
151 + PAD_VALUE: tl.constexpr,
152 + BLOCK_T: tl.constexpr, # timesteps per program
153 + BLOCK_D: tl.constexpr # features per program
154 +):
155 + pid_b = tl.program_id(0) # batch id
156 + pid_t = tl.program_id(1) # block over time dimension
157 + pid_d = tl.program_id(2) # block over feature dimension
158 + off_t = pid_t * BLOCK_T + tl.arange(0, BLOCK_T) # [BLOCK_T]
159 + off_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) # [BLOCK_D]
160 +
161 + # Compute start index and sequence length from cumulative lengths
162 + in_start = 0
163 + for i in range(pid_b):
164 + in_start += tl.load(lengths_ptr + i)
165 + seq_len = tl.load(lengths_ptr + pid_b)
166 +
167 + # valid time positions for this block
168 + t_mask = off_t < Lmax
169 +
170 + # compute input row indices for valid (b, t)
171 + in_row = in_start + off_t
172 + valid_row = (off_t < seq_len) & t_mask
173 +
174 + # Pointers
175 + # x_ptr: row-major [N, D]
176 + x_row_ptr = x_ptr + in_row[:, None] * D + off_d[None, :]
177 +
178 + # out_ptr: row-major [B, Lmax, D]
179 + out_row_ptr = out_ptr + (pid_b * Lmax + off_t)[:,
180 + None] * D + off_d[None, :]
181 +
182 + # Initialize with PAD (cast will occur as needed based on out_ptr dtype)
183 + d_mask = off_d[None, :] < D
184 + pad_vals = tl.full([BLOCK_T, BLOCK_D], PAD_VALUE, tl.float32)
185 + tl.store(out_row_ptr, pad_vals, mask=t_mask[:, None] & d_mask)
186 +
187 + # Load & write only where within seq_len
188 + x_vals = tl.load(x_row_ptr, mask=valid_row[:, None] & d_mask)
189 + tl.store(out_row_ptr, x_vals, mask=valid_row[:, None] & d_mask)
190 +
191 +
192 +def pack_seq_triton(x: torch.Tensor,
193 + lengths: torch.Tensor,
194 + pad_value: float = -float('inf'),
195 + block_t: int = 64,
196 + block_d: int = 64) -> torch.Tensor:
197 + """
198 + Pack sequences of different lengths into a batched tensor.
199 +
200 + Args:
201 + x: [N, ...] - input tensor where N is total number of tokens
202 + lengths: [B] - sequence lengths for each batch
203 + pad_value: value to use for padding
204 + block_t: block size for time dimension
205 + block_d: block size for feature dimension
206 +
207 + Returns:
208 + packed: [B, Lmax, ...] - packed tensor
209 + """
210 +
211 + # Handle multi-dimensional input by reshaping to (N, -1)
212 + original_shape = x.shape
213 + if len(original_shape) > 2:
214 + N = original_shape[0]
215 + x_reshaped = x.reshape(N, -1)
216 + D = x_reshaped.shape[1]
217 + else:
218 + N, D = x.shape
219 + x_reshaped = x
220 +
221 + B = lengths.numel()
222 + Lmax = int(lengths.max().item())
223 +
224 + # Starts are computed inside the kernel from lengths
225 +
226 + out = torch.empty((B, Lmax, D), device=x.device, dtype=x.dtype)
227 +
228 + grid = (B, triton.cdiv(Lmax, block_t), triton.cdiv(D, block_d))
229 + _pack_seq_kernel[grid](x_reshaped,
230 + out,
231 + lengths.int(),
232 + N,
233 + D,
234 + Lmax,
235 + PAD_VALUE=float(pad_value),
236 + BLOCK_T=block_t,
237 + BLOCK_D=block_d,
238 + num_warps=4,
239 + num_stages=2)
240 +
241 + # Reshape output back to original dimensions (except first dimension)
242 + if len(original_shape) > 2:
243 + output_shape = (B, Lmax) + original_shape[1:]
244 + out = out.reshape(output_shape)
245 +
246 + return out
247 +
248 +
249 +@triton.jit
250 +def _unpack_seq_triton_kernel(
251 + packed_ptr, # [B, Lmax, D]
252 + out_ptr, # [N, D]
253 + lengths_ptr, # *i32, [B]
254 + B: tl.constexpr,
255 + Lmax: tl.constexpr,
256 + D: tl.constexpr,
257 + BLOCK_T: tl.constexpr, # timesteps per program
258 + BLOCK_D: tl.constexpr # features per program
259 +):
260 + pid_b = tl.program_id(0) # batch id
261 + pid_t = tl.program_id(1) # block over time dimension
262 + pid_d = tl.program_id(2) # block over feature dimension
263 + off_t = pid_t * BLOCK_T + tl.arange(0, BLOCK_T) # [BLOCK_T]
264 + off_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) # [BLOCK_D]
265 +
266 + # bounds: compute start from cumulative lengths
267 + in_start = 0
268 + for i in range(pid_b):
269 + in_start += tl.load(lengths_ptr + i)
270 + seq_len = tl.load(lengths_ptr + pid_b)
271 +
272 + # valid time positions for this block
273 + t_mask = off_t < Lmax
274 + valid_row = (off_t < seq_len) & t_mask
275 +
276 + # compute output row indices for valid (b, t)
277 + out_row = in_start + off_t
278 +
279 + # Pointers
280 + # packed_ptr: row-major [B, Lmax, D]
281 + packed_row_ptr = packed_ptr + (pid_b * Lmax +
282 + off_t)[:, None] * D + off_d[None, :]
283 +
284 + # out_ptr: row-major [N, D]
285 + out_row_ptr = out_ptr + out_row[:, None] * D + off_d[None, :]
286 +
287 + # Load from packed tensor and store to output
288 + d_mask = off_d[None, :] < D
289 + packed_vals = tl.load(packed_row_ptr, mask=valid_row[:, None] & d_mask)
290 + tl.store(out_row_ptr, packed_vals, mask=valid_row[:, None] & d_mask)
291 +
292 +
293 +def unpack_seq_triton(packed_tensor: torch.Tensor,
294 + lengths: torch.Tensor,
295 + block_t: int = 64,
296 + block_d: int = 64) -> torch.Tensor:
297 + """
298 + Unpack a packed decode query tensor back to the original format.
299 + Efficient Triton implementation.
300 +
301 + Args:
302 + packed_tensor: [B, Lmax, ...] - packed tensor from pack_seq_triton
303 + lengths: [B] - sequence lengths for each batch
304 + block_t: block size for time dimension
305 + block_d: block size for feature dimension
306 +
307 + Returns:
308 + unpacked_tensor: [N, ...] where N = sum(lengths)
309 + """
310 +
311 + # Handle multi-dimensional input by reshaping to (B, Lmax, -1)
312 + original_shape = packed_tensor.shape
313 + if len(original_shape) > 3:
314 + B, Lmax = original_shape[:2]
315 + packed_reshaped = packed_tensor.reshape(B, Lmax, -1)
316 + D = packed_reshaped.shape[2]
317 + else:
318 + B, Lmax, D = packed_tensor.shape
319 + packed_reshaped = packed_tensor
320 +
321 + # Calculate total number of elements
322 + N = int(lengths.sum().item())
323 +
324 + out = torch.empty((N, D),
325 + device=packed_tensor.device,
326 + dtype=packed_tensor.dtype)
327 +
328 + grid = (B, triton.cdiv(Lmax, block_t), triton.cdiv(D, block_d))
329 + _unpack_seq_triton_kernel[grid](packed_reshaped,
330 + out,
331 + lengths.int(),
332 + B,
333 + Lmax,
334 + D,
335 + BLOCK_T=block_t,
336 + BLOCK_D=block_d,
337 + num_warps=4,
338 + num_stages=2)
339 +
340 + # Reshape output back to original dimensions (except first dimension)
341 + if len(original_shape) > 3:
342 + output_shape = (N, ) + original_shape[2:]
343 + out = out.reshape(output_shape)
344 +
345 + return out
vllm/attention/ops/flashmla.py
+117 -48
@@ -19,6 +19,15 @@ if current_platform.is_cuda():
19 else:
20 _flashmla_C_AVAILABLE = False
21
22 +if current_platform.is_cuda():
23 + try:
24 + import vllm._flashmla_extension_C # noqa: F401
25 + _flashmla_extension_C_AVAILABLE = True
26 + except ImportError:
27 + _flashmla_extension_C_AVAILABLE = False
28 +else:
29 + _flashmla_extension_C_AVAILABLE = False
30 +
31
32 def is_flashmla_supported() -> Tuple[bool, Optional[str]]:
33 """
@@ -37,24 +46,34 @@ def is_flashmla_supported() -> Tuple[bool, Optional[str]]:
46
47
48 def get_mla_metadata(
40 - cache_seqlens: torch.Tensor,
41 - num_heads_per_head_k: int,
42 - num_heads_k: int,
43 -) -> Tuple[torch.Tensor, torch.Tensor]:
49 + cache_seqlens: torch.Tensor,
50 + num_q_tokens_per_head_k: int,
51 + num_heads_k: int,
52 + num_heads_q: Optional[int] = None,
53 + is_fp8_kvcache: bool = False,
54 + topk: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]:
55 """
56 Arguments:
46 - cache_seqlens: (batch_size), dtype torch.int32.
47 - num_heads_per_head_k: Equals to seq_len_q * num_heads_q // num_heads_k.
48 - num_heads_k: num_heads_k.
49 -
50 - Return:
51 - tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize),
52 - dtype torch.int32.
53 - num_splits: (batch_size + 1), dtype torch.int32.
57 + - cache_seqlens: (batch_size), dtype torch.int32.
58 + - num_q_tokens_per_head_k:
59 + Equals to num_q_tokens_per_q_seq * num_heads_q // num_heads_k.
60 + - num_heads_k: The number of k heads.
61 + - num_heads_q:
62 + The number of q heads.
63 + This argument is optional when sparse attention is not enabled
64 + - is_fp8_kvcache: Whether the k_cache and v_cache are in fp8 format.
65 + - topk: If not None, sparse attention will be enabled,
66 + and only tokens in the `indices` array
67 + passed to `flash_mla_with_kvcache_sm90` will be attended to.
68 +
69 + Returns:
70 + - tile_scheduler_metadata:
71 + (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32.
72 + - num_splits: (batch_size + 1), dtype torch.int32.
73 """
55 - return torch.ops._flashmla_C.get_mla_metadata(cache_seqlens,
56 - num_heads_per_head_k,
57 - num_heads_k)
74 + return torch.ops._flashmla_C.get_mla_decoding_metadata(
75 + cache_seqlens, num_q_tokens_per_head_k, num_heads_k, num_heads_q,
76 + is_fp8_kvcache, topk)
77
78
79 def flash_mla_with_kvcache(
@@ -69,45 +88,95 @@ def flash_mla_with_kvcache(
88 causal: bool = False,
89 descale_q: Optional[torch.Tensor] = None,
90 descale_k: Optional[torch.Tensor] = None,
91 + is_fp8_kvcache: bool = False,
92 + indices: Optional[torch.Tensor] = None,
93 ) -> Tuple[torch.Tensor, torch.Tensor]:
94 """
95 Arguments:
75 - q: (batch_size, seq_len_q, num_heads_q, head_dim).
76 - k_cache: (num_blocks, page_block_size, num_heads_k, head_dim).
77 - block_table: (batch_size, max_num_blocks_per_seq), torch.int32.
78 - cache_seqlens: (batch_size), torch.int32.
79 - head_dim_v: Head_dim of v.
80 - tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize),
81 - torch.int32, return by get_mla_metadata.
82 - num_splits: (batch_size + 1), torch.int32, return by get_mla_metadata.
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.
88 -
89 - Return:
90 - out: (batch_size, seq_len_q, num_heads_q, head_dim_v).
91 - softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32.
96 + - q: (batch_size, seq_len_q, num_heads_q, head_dim).
97 + - k_cache: (num_blocks, page_block_size, num_heads_k, head_dim).
98 + - block_table: (batch_size, max_num_blocks_per_seq), torch.int32.
99 + - cache_seqlens: (batch_size), torch.int32.
100 + - head_dim_v: Head dimension of v.
101 + - tile_scheduler_metadata:
102 + (num_sm_parts, TileSchedulerMetaDataSize), torch.int32,
103 + returned by get_mla_metadata.
104 + - num_splits:
105 + (batch_size + 1), torch.int32, returned by get_mla_metadata.
106 + - softmax_scale: float.
107 + The scale of QK^T before applying softmax.
108 + Default to 1 / sqrt(head_dim).
109 + - causal: bool. Whether to apply causal attention mask.
110 + - descale_q: (batch_size),
111 + torch.float32. Descaling factors for Q, used for fp8 quantization.
112 + - descale_k: (batch_size),
113 + torch.float32. Descaling factors for K, used for fp8 quantization.
114 + - is_fp8_kvcache: bool.
115 + Whether the k_cache and v_cache are in fp8 format.
116 + For the format of FP8 KV cache, please refer to README.md
117 + - indices: (batch_size, seq_len_q, topk), torch.int32.
118 + If not None, sparse attention will be enabled,
119 + and only tokens in the `indices` array will be attended to.
120 + Invalid indices should be set to -1 or numbers >= total_seq_len_kv.
121 + For details about how to set up `indices`, please refer to README.md.
122 +
123 + Returns:
124 + - out: (batch_size, seq_len_q, num_heads_q, head_dim_v).
125 + - softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32.
126 """
127 if softmax_scale is None:
128 softmax_scale = q.shape[-1]**(-0.5)
95 - out, softmax_lse = torch.ops._flashmla_C.fwd_kvcache_mla(
96 - q,
97 - k_cache,
98 - head_dim_v,
99 - cache_seqlens,
100 - block_table,
101 - softmax_scale,
102 - causal,
103 - tile_scheduler_metadata,
104 - num_splits,
105 - descale_q,
106 - descale_k,
107 - )
108 -
109 - # Note(hc): need revisit when we support DCP with decode query_len > 1.
110 - return out.squeeze(1), softmax_lse.squeeze(-1)
129 + if indices is not None:
130 + # NOTE (zyongye): sparse attention is also causal
131 + # since it only attend to the tokens before
132 + # but here `causal` should not be specified
133 + assert not causal, \
134 + "causal must be `false` if sparse attention is enabled."
135 + assert (descale_q is None) == (
136 + descale_k is None
137 + ), "descale_q and descale_k should be both None or both not None"
138 +
139 + if (descale_q is not None) and (descale_k is not None):
140 + out, softmax_lse = torch.ops._flashmla_extension_C.fwd_kvcache_mla_fp8(
141 + q, k_cache, head_dim_v, cache_seqlens, block_table, softmax_scale,
142 + causal, tile_scheduler_metadata, num_splits, descale_q, descale_k)
143 + else:
144 + out, softmax_lse = torch.ops._flashmla_C.fwd_kvcache_mla(
145 + q, k_cache, head_dim_v, cache_seqlens, block_table, softmax_scale,
146 + causal, tile_scheduler_metadata, num_splits, is_fp8_kvcache,
147 + indices)
148 + return out, softmax_lse
149 +
150 +
151 +def flash_mla_sparse_prefill(
152 + q: torch.Tensor,
153 + kv: torch.Tensor,
154 + indices: torch.Tensor,
155 + sm_scale: float,
156 + d_v: int = 512,
157 +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
158 + """
159 + Sparse attention prefill kernel
160 +
161 + Args:
162 + - q: [s_q, h_q, d_qk], bfloat16
163 + - kv: [s_kv, h_kv, d_qk], bfloat16
164 + - indices: [s_q, h_kv, topk], int32.
165 + Invalid indices should be set to -1 or numbers >= s_kv
166 + - sm_scale: float
167 + - d_v: The dimension of value vectors. Can only be 512
168 +
169 + Returns:
170 + - (output, max_logits, lse)
171 + About the definition of output,
172 + max_logits and lse, please refer to README.md
173 + - output: [s_q, h_q, d_v], bfloat16
174 + - max_logits: [s_q, h_q], float
175 + - lse: [s_q, h_q], float, 2-based log-sum-exp
176 + """
177 + results = torch.ops._flashmla_C.sparse_prefill_fwd(q, kv, indices,
178 + sm_scale, d_v)
179 + return results
180
181
182 #
vllm/attention/ops/paged_attn.py
+1
@@ -50,6 +50,7 @@ class PagedAttention:
50 block_size: int,
51 num_kv_heads: int,
52 head_size: int,
53 + cache_dtype_str: str = "auto",
54 ) -> Tuple[int, ...]:
55 return (2, num_blocks, block_size * num_kv_heads * head_size)
56
vllm/attention/selector.py
+4 -1
@@ -144,6 +144,7 @@ def get_attn_backend(
144 block_size: int,
145 use_mla: bool = False,
146 has_sink: bool = False,
147 + use_sparse: bool = False,
148 ) -> type[AttentionBackend]:
149 """Selects which attention backend to use and lazily imports it."""
150 # Accessing envs.* behind an @lru_cache decorator can cause the wrong
@@ -158,6 +159,7 @@ def get_attn_backend(
159 use_v1=envs.VLLM_USE_V1,
160 use_mla=use_mla,
161 has_sink=has_sink,
162 + use_sparse=use_sparse,
163 )
164
165
@@ -170,6 +172,7 @@ def _cached_get_attn_backend(
172 use_v1: bool = False,
173 use_mla: bool = False,
174 has_sink: bool = False,
175 + use_sparse: bool = False,
176 ) -> type[AttentionBackend]:
177
178 # Check whether a particular choice of backend was
@@ -203,7 +206,7 @@ def _cached_get_attn_backend(
206 # get device-specific attn_backend
207 attention_cls = current_platform.get_attn_backend_cls(
208 selected_backend, head_size, dtype, kv_cache_dtype, block_size, use_v1,
206 - use_mla, has_sink)
209 + use_mla, has_sink, use_sparse)
210 if not attention_cls:
211 raise ValueError(
212 f"Invalid attention backend for {current_platform.device_name}")
vllm/config/cache.py
+13 -7
@@ -22,7 +22,8 @@ else:
22 logger = init_logger(__name__)
23
24 BlockSize = Literal[1, 8, 16, 32, 64, 128]
25 -CacheDType = Literal["auto", "fp8", "fp8_e4m3", "fp8_e5m2", "fp8_inc"]
25 +CacheDType = Literal["auto", "bfloat16", "fp8", "fp8_e4m3", "fp8_e5m2",
26 + "fp8_inc"]
27 MambaDType = Literal["auto", "float32"]
28 PrefixCachingHashAlgo = Literal["sha256", "sha256_cbor"]
29
@@ -52,7 +53,11 @@ class CacheConfig:
53 cache_dtype: CacheDType = "auto"
54 """Data type for kv cache storage. If "auto", will use model data type.
55 CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. ROCm (AMD GPU) supports
55 - fp8 (=fp8_e4m3). Intel Gaudi (HPU) supports fp8 (using fp8_inc)."""
56 + fp8 (=fp8_e4m3). Intel Gaudi (HPU) supports fp8 (using fp8_inc).
57 + Some models (namely DeepSeekV3.2) default to fp8, set to bfloat16 to use
58 + bfloat16 instead, this is an invalid option for models that do not default
59 + to fp8.
60 + """
61 is_attention_free: bool = False
62 """Whether the model is attention-free. This is primarily set in
63 `ModelConfig` and that value should be manually duplicated here."""
@@ -171,11 +176,12 @@ class CacheConfig:
176 if self.cache_dtype == "auto":
177 pass
178 elif self.cache_dtype in get_args(CacheDType):
174 - logger.info(
175 - "Using fp8 data type to store kv cache. It reduces the GPU "
176 - "memory footprint and boosts the performance. "
177 - "Meanwhile, it may cause accuracy drop without a proper "
178 - "scaling factor.")
179 + if self.cache_dtype.startswith("fp8"):
180 + logger.info(
181 + "Using fp8 data type to store kv cache. It reduces the GPU "
182 + "memory footprint and boosts the performance. "
183 + "Meanwhile, it may cause accuracy drop without a proper "
184 + "scaling factor.")
185 else:
186 raise ValueError(f"Unknown kv cache dtype: {self.cache_dtype}")
187
vllm/config/compilation.py
+1
@@ -374,6 +374,7 @@ class CompilationConfig:
374 "vllm.linear_attention",
375 "vllm.plamo2_mamba_mixer",
376 "vllm.gdn_attention",
377 + "vllm.sparse_attn_indexer",
378 ]
379
380 def compute_hash(self) -> str:
vllm/config/model.py
+2 -2
@@ -1082,14 +1082,14 @@ class ModelConfig:
1082 if not hasattr(self.hf_text_config, "model_type"):
1083 return False
1084 elif self.hf_text_config.model_type in \
1085 - ('deepseek_v2', 'deepseek_v3', 'deepseek_mtp',
1085 + ('deepseek_v2', 'deepseek_v3', 'deepseek_v32', 'deepseek_mtp',
1086 'kimi_k2', 'longcat_flash'):
1087 return self.hf_text_config.kv_lora_rank is not None
1088 elif self.hf_text_config.model_type == 'eagle':
1089 # if the model is an EAGLE module, check for the
1090 # underlying architecture
1091 return self.hf_text_config.model.model_type in \
1092 - ('deepseek_v2', 'deepseek_v3') \
1092 + ('deepseek_v2', 'deepseek_v3', 'deepseek_v32') \
1093 and self.hf_text_config.kv_lora_rank is not None
1094 return False
1095
vllm/config/speculative.py
+1 -1
@@ -145,7 +145,7 @@ class SpeculativeConfig:
145
146 @staticmethod
147 def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
148 - if hf_config.model_type == "deepseek_v3":
148 + if hf_config.model_type in ("deepseek_v3", "deepseek_v32"):
149 hf_config.model_type = "deepseek_mtp"
150 if hf_config.model_type == "deepseek_mtp":
151 n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
vllm/model_executor/layers/layernorm.py
+18
@@ -5,6 +5,7 @@ from typing import Optional, Union
5
6 import torch
7 import torch.nn as nn
8 +import torch.nn.functional as F
9
10 import vllm.envs as envs
11 from vllm.model_executor.custom_op import CustomOp
@@ -375,3 +376,20 @@ class PolyNorm(CustomOp):
376 x: torch.Tensor,
377 ) -> torch.Tensor:
378 return poly_norm(x, self.weight, self.bias, self.variance_epsilon)
379 +
380 +
381 +class LayerNorm(nn.Module):
382 + """
383 + Layer Normalization.
384 + """
385 +
386 + def __init__(self, dim: int, eps: float = 1e-6):
387 + super().__init__()
388 + self.dim = dim
389 + self.eps = eps
390 + self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))
391 + self.bias = nn.Parameter(torch.zeros(dim, dtype=torch.float32))
392 +
393 + def forward(self, x: torch.Tensor):
394 + return F.layer_norm(x.float(), (self.dim, ), self.weight, self.bias,
395 + self.eps).type_as(x)
vllm/model_executor/layers/mla.py
+16
@@ -24,6 +24,9 @@ class MLAModules:
24 q_a_layernorm: Optional[torch.nn.Module]
25 q_b_proj: Optional[torch.nn.Module]
26 q_proj: Optional[torch.nn.Module]
27 + indexer: Optional[torch.nn.Module]
28 + is_sparse: bool
29 + topk_indices_buffer: Optional[torch.Tensor]
30
31
32 @CustomOp.register("multi_head_latent_attention")
@@ -76,6 +79,13 @@ class MultiHeadLatentAttention(CustomOp):
79 self.kv_b_proj = mla_modules.kv_b_proj
80 self.rotary_emb = mla_modules.rotary_emb
81 self.o_proj = mla_modules.o_proj
82 + self.indexer = mla_modules.indexer
83 + self.is_sparse = mla_modules.is_sparse
84 +
85 + if self.indexer is not None:
86 + assert hasattr(self.indexer, "topk_tokens")
87 + self.topk_tokens = self.indexer.topk_tokens
88 + self.topk_indices_buffer = mla_modules.topk_indices_buffer
89
90 # In the MLA backend, kv_cache includes both k_c and
91 # pe (i.e. decoupled position embeddings). In particular,
@@ -92,6 +102,7 @@ class MultiHeadLatentAttention(CustomOp):
102 quant_config=quant_config,
103 prefix=f"{prefix}.attn",
104 use_mla=True,
105 + use_sparse=mla_modules.is_sparse,
106 # MLA Args
107 q_lora_rank=self.q_lora_rank,
108 kv_lora_rank=self.kv_lora_rank,
@@ -100,6 +111,7 @@ class MultiHeadLatentAttention(CustomOp):
111 qk_head_dim=self.qk_head_dim,
112 v_head_dim=self.v_head_dim,
113 kv_b_proj=self.kv_b_proj,
114 + indexer=self.indexer,
115 )
116
117 self.prefix = prefix
@@ -145,6 +157,10 @@ class MultiHeadLatentAttention(CustomOp):
157 q[..., self.qk_nope_head_dim:], k_pe = self.rotary_emb(
158 positions, q[..., self.qk_nope_head_dim:], k_pe)
159
160 + if self.indexer and self.is_sparse:
161 + _topk_indices = self.indexer(hidden_states, q_c, positions,
162 + self.rotary_emb)
163 +
164 attn_out = self.mla_attn(
165 q,
166 kv_c_normed,
vllm/model_executor/models/config.py
+27 -2
@@ -346,8 +346,7 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
346 block_size=1,
347 num_kv_heads=model_config.get_num_kv_heads(parallel_config),
348 head_size=model_config.get_head_size(),
349 - dtype=kv_cache_dtype,
350 - use_mla=model_config.use_mla).page_size_bytes
349 + dtype=kv_cache_dtype).page_size_bytes
350
351 model_cls, _ = ModelRegistry.resolve_model_cls(
352 model_config.architecture,
@@ -401,6 +400,31 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
400 "exactly equal.", mamba_padding_pct)
401
402
403 +class DeepseekV3ForCausalLM(VerifyAndUpdateConfig):
404 +
405 + @classmethod
406 + def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
407 + """
408 + Updated fp8 cache to custom "fp8_ds_mla" format for DeepSeekV32
409 + """
410 + hf_config = vllm_config.model_config.hf_config
411 +
412 + is_v32 = hasattr(hf_config, "index_topk")
413 +
414 + if is_v32:
415 + # For DeepSeekV3.2, we use a custom fp8 format as default (i.e.
416 + # "auto")
417 + cache_config = vllm_config.cache_config
418 + if cache_config.cache_dtype == "auto" or \
419 + cache_config.cache_dtype.startswith("fp8"):
420 + cache_config.cache_dtype = "fp8_ds_mla"
421 + logger.info(
422 + "Using custom fp8 kv-cache format for DeepSeekV3.2")
423 + if cache_config.cache_dtype == "bfloat16":
424 + cache_config.cache_dtype = "auto"
425 + logger.info("Using bfloat16 kv-cache for DeepSeekV3.2")
426 +
427 +
428 MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
429 "GteModel": SnowflakeGteNewModelConfig,
430 "GteNewModel": GteNewModelConfig,
@@ -417,4 +441,5 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
441 "MambaForCausalLM": MambaModelConfig,
442 "Mamba2ForCausalLM": MambaModelConfig,
443 "FalconMambaForCausalLM": MambaModelConfig,
444 + "DeepseekV3ForCausalLM": DeepseekV3ForCausalLM,
445 }
vllm/model_executor/models/deepseek_mtp.py
+13 -1
@@ -53,8 +53,20 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
53 self.eh_proj = nn.Linear(config.hidden_size * 2,
54 config.hidden_size,
55 bias=False)
56 +
57 + self.is_v32 = hasattr(config, "index_topk")
58 + if self.is_v32:
59 + topk_tokens = config.index_topk
60 + topk_indices_buffer = torch.empty(
61 + vllm_config.scheduler_config.max_num_batched_tokens,
62 + topk_tokens,
63 + dtype=torch.int32,
64 + device="cuda")
65 + else:
66 + topk_indices_buffer = None
67 self.shared_head = SharedHead(config=config, quant_config=quant_config)
57 - self.mtp_block = DeepseekV2DecoderLayer(vllm_config, prefix)
68 + self.mtp_block = DeepseekV2DecoderLayer(vllm_config, prefix,
69 + topk_indices_buffer)
70
71 def forward(
72 self,
vllm/model_executor/models/deepseek_v2.py
+445 -4
@@ -33,15 +33,21 @@ from torch import nn
33 from transformers import DeepseekV2Config, DeepseekV3Config
34
35 from vllm.attention import Attention
36 +from vllm.attention.backends.abstract import AttentionBackend
37 +from vllm.attention.ops.common import pack_seq_triton, unpack_seq_triton
38 from vllm.compilation.decorators import support_torch_compile
37 -from vllm.config import CacheConfig, ParallelConfig, VllmConfig
39 +from vllm.config import (CacheConfig, ParallelConfig, VllmConfig,
40 + get_current_vllm_config)
41 from vllm.distributed import (get_ep_group, get_pp_group,
42 get_tensor_model_parallel_rank,
43 get_tensor_model_parallel_world_size,
44 tensor_model_parallel_all_gather)
45 +from vllm.forward_context import get_forward_context
46 +from vllm.logger import init_logger
47 from vllm.model_executor.layers.activation import SiluAndMul
48 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
49 from vllm.model_executor.layers.fused_moe import FusedMoE
44 -from vllm.model_executor.layers.layernorm import RMSNorm
50 +from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
51 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
52 MergedColumnParallelLinear,
53 ReplicatedLinear,
@@ -49,6 +55,8 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
55 from vllm.model_executor.layers.logits_processor import LogitsProcessor
56 from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttention
57 from vllm.model_executor.layers.quantization import QuantizationConfig
58 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
59 + per_token_group_quant_fp8)
60 from vllm.model_executor.layers.rotary_embedding import get_rope
61 from vllm.model_executor.layers.shared_fused_moe import SharedFusedMoE
62 from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -56,13 +64,26 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
64 from vllm.model_executor.model_loader.weight_utils import (
65 default_weight_loader, maybe_remap_kv_scale_name)
66 from vllm.model_executor.models.utils import sequence_parallel_chunk
67 +from vllm.platforms import current_platform
68 from vllm.sequence import IntermediateTensors
69 +from vllm.utils import cdiv, direct_register_custom_op
70 +from vllm.utils.deep_gemm import fp8_mqa_logits, fp8_paged_mqa_logits
71 +from vllm.v1.attention.backends.mla.indexer import (DeepseekV32IndexerBackend,
72 + DeepseekV32IndexerMetadata)
73 +from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
74
75 from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
76 from .utils import (PPMissingLayer, is_pp_missing_parameter,
77 make_empty_intermediate_tensors_factory, make_layers,
78 maybe_prefix)
79
80 +if current_platform.is_cuda_alike():
81 + from vllm import _custom_ops as ops
82 +elif current_platform.is_xpu():
83 + from vllm._ipex_ops import ipex_ops as ops
84 +
85 +logger = init_logger(__name__)
86 +
87
88 class DeepseekV2MLP(nn.Module):
89
@@ -276,6 +297,7 @@ class DeepseekV2Attention(nn.Module):
297
298 def __init__(
299 self,
300 + vllm_config: VllmConfig,
301 config: Union[DeepseekV2Config, DeepseekV3Config],
302 hidden_size: int,
303 num_heads: int,
@@ -289,6 +311,7 @@ class DeepseekV2Attention(nn.Module):
311 max_position_embeddings: int = 8192,
312 cache_config: Optional[CacheConfig] = None,
313 quant_config: Optional[QuantizationConfig] = None,
314 + topk_indices_buffer: Optional[torch.Tensor] = None,
315 prefix: str = "",
316 ) -> None:
317 super().__init__()
@@ -306,6 +329,8 @@ class DeepseekV2Attention(nn.Module):
329 self.scaling = self.qk_head_dim**-0.5
330 self.rope_theta = rope_theta
331 self.max_position_embeddings = max_position_embeddings
332 + assert topk_indices_buffer is None, "topk_indices_buffer is not \
333 + supported for DeepseekV2Attention"
334
335 if self.q_lora_rank is not None:
336 self.q_a_proj = ReplicatedLinear(self.hidden_size,
@@ -418,6 +443,391 @@ class DeepseekV2Attention(nn.Module):
443 return output
444
445
446 +class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):
447 +
448 + def __init__(self, head_dim: int, dtype: torch.dtype, prefix: str,
449 + cache_config: CacheConfig):
450 + super().__init__()
451 + self.kv_cache = [torch.tensor([])]
452 + self.head_dim = head_dim
453 + self.prefix = prefix
454 + self.cache_config = cache_config
455 + self.dtype = dtype
456 + compilation_config = get_current_vllm_config().compilation_config
457 + if prefix in compilation_config.static_forward_context:
458 + raise ValueError(f"Duplicate layer name: {prefix}")
459 + compilation_config.static_forward_context[prefix] = self
460 +
461 + def get_kv_cache_spec(self) -> KVCacheSpec:
462 + return MLAAttentionSpec( # Only has one vector instead of K + V
463 + block_size=self.cache_config.block_size,
464 + num_kv_heads=1,
465 + head_size=self.head_dim,
466 + dtype=self.dtype,
467 + )
468 +
469 + def forward(self):
470 + ...
471 +
472 + def get_attn_backend(self) -> AttentionBackend:
473 + return DeepseekV32IndexerBackend
474 +
475 +
476 +@torch.inference_mode()
477 +def cp_gather_indexer_k_quant_cache(
478 + kv_cache, # [num_blocks, block_size, head_dim + 1]
479 + dst_value, # [cu_seq_lens[-1], head_dim]
480 + dst_scale, # [cu_seq_lens[-1], 4]
481 + block_table, # [batch_size, num_blocks]
482 + cu_seq_lens, # [batch_size + 1, ]
483 + batch_size,
484 +):
485 + num_blocks, block_size, _ = kv_cache.shape
486 + head_dim = dst_value.shape[-1]
487 + kv_cache = kv_cache.view(num_blocks, -1)
488 +
489 + expected_value = []
490 + expected_scale = []
491 + for b in range(batch_size):
492 + s = cu_seq_lens[b + 1] - cu_seq_lens[b]
493 + if s == 0:
494 + continue
495 + tot = cdiv(s, block_size)
496 + blocks = block_table[b, :tot]
497 +
498 + value = []
499 + scale = []
500 + full_block = torch.arange(tot - 1,
501 + device=kv_cache.device,
502 + dtype=torch.int32)
503 + non_remaining_value = kv_cache[blocks[full_block], :block_size *
504 + head_dim].view(-1, head_dim)
505 + non_remaining_scale = kv_cache[blocks[full_block],
506 + block_size * head_dim:].view(-1, 4)
507 +
508 + remaining = s - (tot - 1) * block_size
509 +
510 + value = torch.cat([
511 + non_remaining_value,
512 + kv_cache[blocks[-1], :remaining * head_dim].view(-1, head_dim)
513 + ],
514 + dim=0)
515 + scale = torch.cat([
516 + non_remaining_scale,
517 + kv_cache[blocks[-1], block_size * head_dim:block_size * head_dim +
518 + remaining * 4].view(-1, 4)
519 + ],
520 + dim=0)
521 +
522 + expected_value.append(value)
523 + expected_scale.append(scale)
524 +
525 + gather_value = torch.cat(expected_value, dim=0).view(-1, head_dim)
526 + gather_scale = torch.cat(expected_scale, dim=0).view(-1, 4)
527 + gather_value = gather_value.view(torch.float8_e4m3fn)
528 + gather_scale = gather_scale.view(torch.float32)
529 + dst_value.copy_(gather_value)
530 + dst_scale.copy_(gather_scale)
531 +
532 +
533 +def sparse_attn_indexer(
534 + hidden_states: torch.Tensor,
535 + k_cache_prefix: str,
536 + kv_cache: torch.Tensor,
537 + q_fp8: torch.Tensor,
538 + k: torch.Tensor,
539 + weights: torch.Tensor,
540 + quant_block_size: int,
541 + scale_fmt: Optional[str],
542 + topk_tokens: int,
543 + head_dim: int,
544 + max_model_len: int,
545 + total_seq_lens: int,
546 + topk_indices_buffer: Optional[torch.Tensor],
547 +) -> torch.Tensor:
548 +
549 + # careful! this will be None in dummy run
550 + attn_metadata = get_forward_context().attn_metadata
551 + # assert isinstance(attn_metadata, dict)
552 + if not isinstance(attn_metadata, dict):
553 + return sparse_attn_indexer_fake(
554 + hidden_states,
555 + k_cache_prefix,
556 + kv_cache,
557 + q_fp8,
558 + k,
559 + weights,
560 + quant_block_size,
561 + scale_fmt,
562 + topk_tokens,
563 + head_dim,
564 + max_model_len,
565 + total_seq_lens,
566 + topk_indices_buffer,
567 + )
568 + attn_metadata = attn_metadata[k_cache_prefix]
569 + assert isinstance(attn_metadata, DeepseekV32IndexerMetadata)
570 + slot_mapping = attn_metadata.slot_mapping
571 + has_decode = attn_metadata.num_decodes > 0
572 + has_prefill = attn_metadata.num_prefills > 0
573 + num_decode_tokens = attn_metadata.num_decode_tokens
574 +
575 + ops.indexer_k_quant_and_cache(
576 + k,
577 + kv_cache,
578 + slot_mapping,
579 + quant_block_size,
580 + scale_fmt,
581 + )
582 +
583 + topk_indices_buffer[:hidden_states.shape[0]] = -1
584 + if has_prefill:
585 + prefill_metadata = attn_metadata.prefill
586 + num_prefills = attn_metadata.num_prefills
587 + k_fp8 = torch.empty([prefill_metadata.total_seq_lens, head_dim],
588 + device=k.device,
589 + dtype=torch.float8_e4m3fn)
590 + k_scale = torch.empty([prefill_metadata.total_seq_lens, 1],
591 + device=k.device,
592 + dtype=torch.float32)
593 + cp_gather_indexer_k_quant_cache(
594 + kv_cache,
595 + k_fp8,
596 + k_scale,
597 + prefill_metadata.block_table,
598 + prefill_metadata.cu_seq_lens,
599 + num_prefills,
600 + )
601 + cu_seqlen_ks = prefill_metadata.cu_seqlen_ks
602 + cu_seqlen_ke = prefill_metadata.cu_seqlen_ke
603 + num_tokens = attn_metadata.num_actual_tokens
604 + logits = fp8_mqa_logits(
605 + q_fp8[num_decode_tokens:num_tokens],
606 + (k_fp8, k_scale),
607 + weights[num_decode_tokens:num_tokens],
608 + cu_seqlen_ks,
609 + cu_seqlen_ke,
610 + )
611 + topk_indices = logits.topk(min(topk_tokens, logits.shape[-1]),
612 + dim=-1)[1]
613 + topk_indices -= cu_seqlen_ks[:, None]
614 + mask_lo = topk_indices >= 0
615 + mask_hi = topk_indices - (cu_seqlen_ke - cu_seqlen_ks)[:, None] < 0
616 + mask = torch.full_like(topk_indices,
617 + False,
618 + dtype=torch.bool,
619 + device=topk_indices.device)
620 + mask = mask_lo & mask_hi
621 + topk_indices = topk_indices.masked_fill(~mask, -1)
622 + topk_indices_buffer[num_decode_tokens:num_tokens, :topk_indices.
623 + shape[-1]] = topk_indices.to(dtype=torch.int32)
624 +
625 + if has_decode:
626 + decode_metadata = attn_metadata.decode
627 + # kv_cache size requirement [num_block, block_size, n_head, head_dim],
628 + # we only have [num_block, block_size, head_dim],
629 + kv_cache = kv_cache.unsqueeze(-2)
630 + decode_lens = decode_metadata.decode_lens
631 + if decode_metadata.requires_padding:
632 + # pad in edge case where we have short chunked prefill length <
633 + # decode_threshold since we unstrictly split
634 + # prefill and decode by decode_threshold
635 + # (currently set to 1 + speculative tokens)
636 + padded_q_fp8_decode_tokens = pack_seq_triton(
637 + q_fp8[:num_decode_tokens], decode_lens)
638 + else:
639 + padded_q_fp8_decode_tokens = q_fp8[:num_decode_tokens].reshape(
640 + decode_lens.shape[0], -1, *q_fp8.shape[1:])
641 + # TODO: move and optimize below logic with triton kernels
642 + batch_size = padded_q_fp8_decode_tokens.shape[0]
643 + next_n = padded_q_fp8_decode_tokens.shape[1]
644 + assert batch_size == decode_metadata.seq_lens.shape[0]
645 + num_padded_tokens = batch_size * next_n
646 + logits = fp8_paged_mqa_logits(
647 + padded_q_fp8_decode_tokens,
648 + kv_cache,
649 + weights[:num_padded_tokens],
650 + decode_metadata.seq_lens,
651 + decode_metadata.block_table,
652 + decode_metadata.schedule_metadata,
653 + max_model_len=max_model_len,
654 + )
655 + # padded query len
656 + current_device = padded_q_fp8_decode_tokens.device
657 + padded_num_tokens = batch_size * next_n
658 + positions = torch.arange(max_model_len,
659 + device=current_device).unsqueeze(0).expand(
660 + batch_size * next_n, -1)
661 + row_indices = torch.arange(padded_num_tokens,
662 + device=current_device) // next_n
663 + next_n_offset = torch.arange(
664 + padded_num_tokens,
665 + device=padded_q_fp8_decode_tokens.device) % next_n
666 + index_end_pos = (decode_metadata.seq_lens[row_indices] - next_n +
667 + next_n_offset).unsqueeze(1)
668 + # index_end_pos: [B * N, 1]
669 + mask = positions <= index_end_pos
670 + # mask: [B * N, L]
671 + logits = logits.masked_fill(~mask, float('-inf'))
672 + topk_indices = logits.topk(topk_tokens,
673 + dim=-1)[1].to(torch.int32) # [B * N, K]
674 + # ensure we don't set indices for the top k
675 + # that is out of range(masked already)
676 + # this will happen if context length is shorter than K
677 + topk_indices[topk_indices > index_end_pos] = -1
678 + if decode_metadata.requires_padding:
679 + # if padded, we need to unpack
680 + # the topk indices removing padded tokens
681 + topk_indices = unpack_seq_triton(
682 + topk_indices.reshape(batch_size, -1, topk_indices.shape[-1]),
683 + decode_lens)
684 + topk_indices_buffer[:num_decode_tokens, :topk_indices.
685 + shape[-1]] = topk_indices.to(dtype=torch.int32)
686 +
687 + return topk_indices_buffer
688 +
689 +
690 +def sparse_attn_indexer_fake(
691 + hidden_states: torch.Tensor,
692 + k_cache_prefix: str,
693 + kv_cache: torch.Tensor,
694 + q_fp8: torch.Tensor,
695 + k: torch.Tensor,
696 + weights: torch.Tensor,
697 + quant_block_size: int,
698 + scale_fmt: Optional[str],
699 + topk_tokens: int,
700 + head_dim: int,
701 + max_model_len: int,
702 + total_seq_lens: int,
703 + topk_indices_buffer: Optional[torch.Tensor],
704 +) -> torch.Tensor:
705 + # profile run
706 + # NOTE(Chen): create the max possible flattened_kv. So that
707 + # profile_run can get correct memory usage.
708 + _flattened_kv = torch.empty([total_seq_lens, head_dim + 4],
709 + device=k.device,
710 + dtype=torch.uint8)
711 + _k_fp8 = _flattened_kv[..., :head_dim].view(
712 + torch.float8_e4m3fn).contiguous()
713 + _k_scale = _flattened_kv[..., head_dim:].view(torch.float32).contiguous()
714 + return topk_indices_buffer
715 +
716 +
717 +direct_register_custom_op(
718 + op_name="sparse_attn_indexer",
719 + op_func=sparse_attn_indexer,
720 + mutates_args=["topk_indices_buffer"],
721 + fake_impl=sparse_attn_indexer_fake,
722 + dispatch_key=current_platform.dispatch_key,
723 +)
724 +
725 +
726 +class Indexer(nn.Module):
727 +
728 + def __init__(self,
729 + vllm_config: VllmConfig,
730 + config: Union[DeepseekV2Config, DeepseekV3Config],
731 + hidden_size: int,
732 + q_lora_rank: int,
733 + quant_config: Optional[QuantizationConfig],
734 + cache_config: Optional[CacheConfig],
735 + topk_indices_buffer: Optional[torch.Tensor],
736 + prefix: str = ""):
737 + super().__init__()
738 + self.vllm_config = vllm_config
739 + self.config = config
740 + # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"]
741 + self.topk_tokens = config.index_topk
742 + self.n_head = config.index_n_heads # 64
743 + self.head_dim = config.index_head_dim # 128
744 + self.rope_dim = config.qk_rope_head_dim # 64
745 + self.q_lora_rank = q_lora_rank # 1536
746 + # no tensor parallel, just replicated
747 + self.wq_b = ReplicatedLinear(self.q_lora_rank,
748 + self.head_dim * self.n_head,
749 + bias=False,
750 + quant_config=quant_config,
751 + prefix=f"{prefix}.wq_b")
752 + self.wk = ReplicatedLinear(hidden_size,
753 + self.head_dim,
754 + bias=False,
755 + quant_config=quant_config,
756 + prefix=f"{prefix}.wk")
757 + self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
758 + self.weights_proj = ReplicatedLinear(hidden_size,
759 + self.n_head,
760 + quant_config=None,
761 + prefix=f"{prefix}.weights_proj")
762 + self.softmax_scale = self.head_dim**-0.5
763 +
764 + self.scale_fmt = "ue8m0"
765 + self.quant_block_size = 128 # TODO: get from config
766 + self.topk_indices_buffer = topk_indices_buffer
767 +
768 + # NOTE: (zyongye) we use fp8 naive cache,
769 + # where we store value in fp8 and scale in fp32
770 + # per self.quant_block_size element
771 + self.k_cache = DeepseekV32IndexerCache(
772 + head_dim=self.head_dim +
773 + self.head_dim // self.quant_block_size * 4,
774 + dtype=torch.uint8,
775 + prefix=f"{prefix}.k_cache",
776 + cache_config=cache_config)
777 + self.max_model_len = vllm_config.model_config.max_model_len
778 + self.prefix = prefix
779 + from vllm.v1.attention.backends.mla.indexer import (
780 + get_max_prefill_buffer_size)
781 + self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config)
782 +
783 + def forward(self, hidden_states: torch.Tensor, qr: torch.Tensor, positions,
784 + rotary_emb) -> torch.Tensor:
785 + q, _ = self.wq_b(qr)
786 + q = q.view(-1, self.n_head, self.head_dim)
787 + q_pe, q_nope = torch.split(
788 + q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1)
789 +
790 + k, _ = self.wk(hidden_states)
791 + k = self.k_norm(k)
792 + k_pe, k_nope = torch.split(
793 + k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1)
794 +
795 + q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
796 + q = torch.cat([q_pe, q_nope], dim=-1)
797 + k = torch.cat([k_pe.squeeze(1), k_nope], dim=-1)
798 +
799 + # we only quant q here since k quant is fused with cache insertion
800 + q = q.view(-1, self.head_dim)
801 + q_fp8, q_scale = per_token_group_quant_fp8(q,
802 + self.quant_block_size,
803 + column_major_scales=False,
804 + use_ue8m0=self.scale_fmt
805 + is not None)
806 + q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim)
807 + q_scale = q_scale.view(-1, self.n_head, 1)
808 +
809 + weights, _ = self.weights_proj(hidden_states)
810 + weights = weights.unsqueeze(
811 + -1) * q_scale * self.softmax_scale * self.n_head**-0.5
812 + weights = weights.squeeze(-1)
813 +
814 + return torch.ops.vllm.sparse_attn_indexer(
815 + hidden_states,
816 + self.k_cache.prefix,
817 + self.k_cache.kv_cache[0],
818 + q_fp8,
819 + k,
820 + weights,
821 + self.quant_block_size,
822 + self.scale_fmt,
823 + self.topk_tokens,
824 + self.head_dim,
825 + self.max_model_len,
826 + self.max_total_seq_len,
827 + self.topk_indices_buffer,
828 + )
829 +
830 +
831 class DeepseekV2MLAAttention(nn.Module):
832 """
833 Main reference: DeepseekV2 paper, and FlashInfer Implementation
@@ -429,6 +839,7 @@ class DeepseekV2MLAAttention(nn.Module):
839
840 def __init__(
841 self,
842 + vllm_config: VllmConfig,
843 config: Union[DeepseekV2Config, DeepseekV3Config],
844 hidden_size: int,
845 num_heads: int,
@@ -443,6 +854,7 @@ class DeepseekV2MLAAttention(nn.Module):
854 cache_config: Optional[CacheConfig] = None,
855 quant_config: Optional[QuantizationConfig] = None,
856 prefix: str = "",
857 + topk_indices_buffer: Optional[torch.Tensor] = None,
858 ) -> None:
859 super().__init__()
860 self.hidden_size = hidden_size
@@ -523,6 +935,15 @@ class DeepseekV2MLAAttention(nn.Module):
935 mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
936 self.scaling = self.scaling * mscale * mscale
937
938 + self.is_v32 = hasattr(config, "index_topk")
939 +
940 + if self.is_v32:
941 + self.indexer = Indexer(vllm_config, config, hidden_size,
942 + q_lora_rank, quant_config, cache_config,
943 + topk_indices_buffer, f"{prefix}.indexer")
944 + else:
945 + self.indexer = None
946 +
947 mla_modules = MLAModules(
948 kv_a_layernorm=self.kv_a_layernorm,
949 kv_b_proj=self.kv_b_proj,
@@ -536,7 +957,11 @@ class DeepseekV2MLAAttention(nn.Module):
957 if self.q_lora_rank is not None else None,
958 q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,
959 q_proj=self.q_proj if self.q_lora_rank is None else None,
960 + indexer=self.indexer,
961 + is_sparse=self.is_v32,
962 + topk_indices_buffer=topk_indices_buffer,
963 )
964 +
965 self.mla_attn = MultiHeadLatentAttention(
966 self.hidden_size,
967 self.num_local_heads,
@@ -562,7 +987,10 @@ class DeepseekV2MLAAttention(nn.Module):
987
988 class DeepseekV2DecoderLayer(nn.Module):
989
565 - def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
990 + def __init__(self,
991 + vllm_config: VllmConfig,
992 + prefix: str,
993 + topk_indices_buffer: Optional[torch.Tensor] = None) -> None:
994 super().__init__()
995
996 config = vllm_config.model_config.hf_config
@@ -585,6 +1013,7 @@ class DeepseekV2DecoderLayer(nn.Module):
1013 else:
1014 attn_cls = DeepseekV2Attention
1015 self.self_attn = attn_cls(
1016 + vllm_config=vllm_config,
1017 config=config,
1018 hidden_size=self.hidden_size,
1019 num_heads=config.num_attention_heads,
@@ -600,6 +1029,7 @@ class DeepseekV2DecoderLayer(nn.Module):
1029 cache_config=cache_config,
1030 quant_config=quant_config,
1031 prefix=f"{prefix}.self_attn",
1032 + topk_indices_buffer=topk_indices_buffer,
1033 )
1034
1035 if (config.n_routed_experts is not None
@@ -683,6 +1113,16 @@ class DeepseekV2Model(nn.Module):
1113 self.config = config
1114
1115 self.vocab_size = config.vocab_size
1116 + self.is_v32 = hasattr(config, "index_topk")
1117 + if self.is_v32:
1118 + topk_tokens = config.index_topk
1119 + topk_indices_buffer = torch.empty(
1120 + vllm_config.scheduler_config.max_num_batched_tokens,
1121 + topk_tokens,
1122 + dtype=torch.int32,
1123 + device="cuda")
1124 + else:
1125 + topk_indices_buffer = None
1126
1127 if get_pp_group().is_first_rank:
1128 self.embed_tokens = VocabParallelEmbedding(
@@ -695,7 +1135,8 @@ class DeepseekV2Model(nn.Module):
1135
1136 self.start_layer, self.end_layer, self.layers = make_layers(
1137 config.num_hidden_layers,
698 - lambda prefix: DeepseekV2DecoderLayer(vllm_config, prefix),
1138 + lambda prefix: DeepseekV2DecoderLayer(vllm_config, prefix,
1139 + topk_indices_buffer),
1140 prefix=f"{prefix}.layers")
1141
1142 if get_pp_group().is_last_rank:
vllm/model_executor/models/longcat_flash.py
+3
@@ -308,6 +308,7 @@ class FlashDecoderLayer(nn.Module):
308
309 def __init__(
310 self,
311 + vllm_config: VllmConfig,
312 config: FlashConfig,
313 cache_config: Optional[CacheConfig] = None,
314 quant_config: Optional[QuantizationConfig] = None,
@@ -329,6 +330,7 @@ class FlashDecoderLayer(nn.Module):
330 # Dual attention structure
331 self.self_attn = nn.ModuleList([
332 DeepseekV2MLAAttention(
333 + vllm_config=vllm_config,
334 config=config,
335 hidden_size=self.hidden_size,
336 num_heads=config.num_attention_heads,
@@ -454,6 +456,7 @@ class FlashModel(nn.Module):
456 self.start_layer, self.end_layer, self.layers = make_layers(
457 config.num_hidden_layers,
458 lambda prefix: FlashDecoderLayer(
459 + vllm_config,
460 config,
461 cache_config=cache_config,
462 quant_config=quant_config,
vllm/model_executor/models/registry.py
+1
@@ -70,6 +70,7 @@ _TEXT_GENERATION_MODELS = {
70 "DeepseekForCausalLM": ("deepseek", "DeepseekForCausalLM"),
71 "DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"),
72 "DeepseekV3ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),
73 + "DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),
74 "Dots1ForCausalLM": ("dots1", "Dots1ForCausalLM"),
75 "Ernie4_5ForCausalLM": ("ernie45", "Ernie4_5ForCausalLM"),
76 "Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"),
vllm/platforms/cpu.py
+4 -1
@@ -93,11 +93,14 @@ class CpuPlatform(Platform):
93 def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,
94 dtype: torch.dtype, kv_cache_dtype: Optional[str],
95 block_size: int, use_v1: bool, use_mla: bool,
96 - has_sink: bool) -> str:
96 + has_sink: bool, use_sparse: bool) -> str:
97 if selected_backend and selected_backend != _Backend.TORCH_SDPA:
98 logger.info("Cannot use %s backend on CPU.", selected_backend)
99 if use_mla:
100 raise NotImplementedError("MLA is not supported on CPU.")
101 + if use_sparse:
102 + raise NotImplementedError(
103 + "Sparse Attention is not supported on CPU.")
104 logger.info("Using Torch SDPA backend.")
105 if not use_v1:
106 raise ValueError("CPU backend only supports V1.")
vllm/platforms/cuda.py
+14 -1
@@ -119,6 +119,8 @@ class CudaPlatformBase(Platform):
119 # TODO(lucas): handle this more gracefully
120 # Note: model_config may be None during testing
121 if model_config is not None and model_config.use_mla:
122 + use_sparse = hasattr(vllm_config.model_config.hf_config,
123 + "index_topk")
124 # If `VLLM_ATTENTION_BACKEND` is not set and we are using MLA,
125 # then we default to FlashMLA backend for non-blackwell GPUs,
126 # else we default to CutlassMLA. For each case, we force the
@@ -165,6 +167,12 @@ class CudaPlatformBase(Platform):
167 "Forcing kv cache block size to 64 for FlashInferMLA "
168 "backend.")
169
170 + # TODO(Chen): remove this hacky code
171 + if use_sparse and cache_config.block_size != 64:
172 + cache_config.block_size = 64
173 + logger.info(
174 + "Forcing kv cache block size to 64 for FlashMLASparse "
175 + "backend.")
176 # lazy import to avoid circular import
177 from vllm.config import CUDAGraphMode
178
@@ -221,7 +229,7 @@ class CudaPlatformBase(Platform):
229 @classmethod
230 def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
231 kv_cache_dtype, block_size, use_v1, use_mla,
224 - has_sink) -> str:
232 + has_sink, use_sparse) -> str:
233 if use_mla:
234 if not use_v1:
235 raise RuntimeError(
@@ -231,6 +239,11 @@ class CudaPlatformBase(Platform):
239 from vllm.attention.ops.flashmla import is_flashmla_supported
240 from vllm.attention.utils.fa_utils import flash_attn_supports_mla
241
242 + if use_sparse:
243 + logger.info_once("Using Sparse MLA backend on V1 engine.")
244 + return ("vllm.v1.attention.backends.mla.flashmla_sparse."
245 + "FlashMLASparseBackend")
246 +
247 use_cutlassmla = selected_backend == _Backend.CUTLASS_MLA or (
248 selected_backend is None and cls.is_device_capability(100)
249 and block_size == 128)
vllm/platforms/interface.py
+1 -1
@@ -194,7 +194,7 @@ class Platform:
194 def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,
195 dtype: torch.dtype, kv_cache_dtype: Optional[str],
196 block_size: int, use_v1: bool, use_mla: bool,
197 - has_sink: bool) -> str:
197 + has_sink: bool, use_sparse: bool) -> str:
198 """Get the attention backend class of a device."""
199 return ""
200
vllm/platforms/rocm.py
+4 -1
@@ -195,7 +195,10 @@ class RocmPlatform(Platform):
195 @classmethod
196 def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
197 kv_cache_dtype, block_size, use_v1, use_mla,
198 - has_sink) -> str:
198 + has_sink, use_sparse) -> str:
199 + if use_sparse:
200 + raise NotImplementedError(
201 + "Sparse Attention is not supported on ROCm.")
202 if use_mla:
203 if not use_v1:
204 raise RuntimeError(
vllm/platforms/tpu.py
+4 -1
@@ -49,7 +49,10 @@ class TpuPlatform(Platform):
49 def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,
50 dtype: torch.dtype, kv_cache_dtype: Optional[str],
51 block_size: int, use_v1: bool, use_mla: bool,
52 - has_sink) -> str:
52 + has_sink, use_sparse) -> str:
53 + if use_sparse:
54 + raise NotImplementedError(
55 + "Sparse Attention is not supported on TPU.")
56 if selected_backend != _Backend.PALLAS:
57 logger.info("Cannot use %s backend on TPU.", selected_backend)
58
vllm/platforms/xpu.py
+4 -1
@@ -36,7 +36,10 @@ class XPUPlatform(Platform):
36 def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,
37 dtype: torch.dtype, kv_cache_dtype: Optional[str],
38 block_size: int, use_v1: bool, use_mla: bool,
39 - has_sink: bool) -> str:
39 + has_sink: bool, use_sparse) -> str:
40 + if use_sparse:
41 + raise NotImplementedError(
42 + "Sparse Attention is not supported on XPU.")
43 use_v1 = envs.VLLM_USE_V1
44 if not use_v1:
45 raise ValueError("XPU backend only supports V1.")
vllm/transformers_utils/config.py
+2
@@ -66,6 +66,8 @@ class LazyConfigDict(dict):
66 _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
67 chatglm="ChatGLMConfig",
68 deepseek_vl_v2="DeepseekVLV2Config",
69 + deepseek_v3="DeepseekV3Config",
70 + deepseek_v32="DeepseekV3Config",
71 kimi_vl="KimiVLConfig",
72 Llama_Nemotron_Nano_VL="Nemotron_Nano_VL_Config",
73 RefinedWeb="RWConfig", # For tiiuae/falcon-40b(-instruct)
vllm/transformers_utils/configs/__init__.py
+2
@@ -8,6 +8,7 @@ Model configs may be defined in this directory for the following reasons:
8 """
9
10 from vllm.transformers_utils.configs.chatglm import ChatGLMConfig
11 +from vllm.transformers_utils.configs.deepseek_v3 import DeepseekV3Config
12 from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config
13 from vllm.transformers_utils.configs.dotsocr import DotsOCRConfig
14 from vllm.transformers_utils.configs.eagle import EAGLEConfig
@@ -37,6 +38,7 @@ from vllm.transformers_utils.configs.ultravox import UltravoxConfig
38 __all__ = [
39 "ChatGLMConfig",
40 "DeepseekVLV2Config",
41 + "DeepseekV3Config",
42 "DotsOCRConfig",
43 "EAGLEConfig",
44 "RWConfig",
vllm/transformers_utils/configs/deepseek_v3.py new
+101
@@ -0,0 +1,101 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from transformers.configuration_utils import PretrainedConfig
4 +from transformers.utils import logging
5 +
6 +logger = logging.get_logger(__name__)
7 +
8 +
9 +class DeepseekV3Config(PretrainedConfig):
10 +
11 + model_type = "deepseek_v3"
12 + keys_to_ignore_at_inference = ["past_key_values"]
13 +
14 + def __init__(
15 + self,
16 + vocab_size=129280,
17 + hidden_size=7168,
18 + intermediate_size=18432,
19 + moe_intermediate_size=2048,
20 + num_hidden_layers=61,
21 + num_nextn_predict_layers=1,
22 + num_attention_heads=128,
23 + num_key_value_heads=128,
24 + n_shared_experts=1,
25 + n_routed_experts=256,
26 + ep_size=1,
27 + routed_scaling_factor=2.5,
28 + kv_lora_rank=512,
29 + q_lora_rank=1536,
30 + qk_rope_head_dim=64,
31 + v_head_dim=128,
32 + qk_nope_head_dim=128,
33 + topk_method='noaux_tc',
34 + n_group=8,
35 + topk_group=4,
36 + num_experts_per_tok=8,
37 + moe_layer_freq=1,
38 + first_k_dense_replace=3,
39 + norm_topk_prob=True,
40 + scoring_func='sigmoid',
41 + hidden_act="silu",
42 + max_position_embeddings=4096,
43 + initializer_range=0.02,
44 + rms_norm_eps=1e-6,
45 + use_cache=True,
46 + pad_token_id=None,
47 + bos_token_id=0,
48 + eos_token_id=1,
49 + tie_word_embeddings=False,
50 + rope_theta=10000.0,
51 + rope_scaling=None,
52 + attention_bias=False,
53 + attention_dropout=0.0,
54 + **kwargs,
55 + ):
56 + self.vocab_size = vocab_size
57 + self.max_position_embeddings = max_position_embeddings
58 + self.hidden_size = hidden_size
59 + self.intermediate_size = intermediate_size
60 + self.moe_intermediate_size = moe_intermediate_size
61 + self.num_hidden_layers = num_hidden_layers
62 + self.num_nextn_predict_layers = num_nextn_predict_layers
63 + self.num_attention_heads = num_attention_heads
64 + self.n_shared_experts = n_shared_experts
65 + self.n_routed_experts = n_routed_experts
66 + self.ep_size = ep_size
67 + self.routed_scaling_factor = routed_scaling_factor
68 + self.kv_lora_rank = kv_lora_rank
69 + self.q_lora_rank = q_lora_rank
70 + self.qk_rope_head_dim = qk_rope_head_dim
71 + self.v_head_dim = v_head_dim
72 + self.qk_nope_head_dim = qk_nope_head_dim
73 + self.topk_method = topk_method
74 + self.n_group = n_group
75 + self.topk_group = topk_group
76 + self.num_experts_per_tok = num_experts_per_tok
77 + self.moe_layer_freq = moe_layer_freq
78 + self.first_k_dense_replace = first_k_dense_replace
79 + self.norm_topk_prob = norm_topk_prob
80 + self.scoring_func = scoring_func
81 + # for backward compatibility
82 + if num_key_value_heads is None:
83 + num_key_value_heads = num_attention_heads
84 +
85 + self.num_key_value_heads = num_key_value_heads
86 + self.hidden_act = hidden_act
87 + self.initializer_range = initializer_range
88 + self.rms_norm_eps = rms_norm_eps
89 + self.use_cache = use_cache
90 + self.rope_theta = rope_theta
91 + self.rope_scaling = rope_scaling
92 + self.attention_bias = attention_bias
93 + self.attention_dropout = attention_dropout
94 +
95 + super().__init__(
96 + pad_token_id=pad_token_id,
97 + bos_token_id=bos_token_id,
98 + eos_token_id=eos_token_id,
99 + tie_word_embeddings=tie_word_embeddings,
100 + **kwargs,
101 + )
vllm/utils/__init__.py
+7
@@ -130,6 +130,7 @@ STR_DTYPE_TO_TORCH_DTYPE = {
130 "fp8_e5m2": torch.uint8,
131 "int8": torch.int8,
132 "fp8_inc": torch.float8_e4m3fn,
133 + "fp8_ds_mla": torch.uint8,
134 }
135
136 TORCH_DTYPE_TO_NUMPY_DTYPE = {
@@ -3433,6 +3434,12 @@ def has_triton_kernels() -> bool:
3434 return _has_module("triton_kernels")
3435
3436
3437 +def has_tilelang() -> bool:
3438 + """Whether the optional `tilelang` package is available."""
3439 +
3440 + return _has_module("tilelang")
3441 +
3442 +
3443 def set_process_title(name: str,
3444 suffix: str = "",
3445 prefix: str = envs.VLLM_PROCESS_NAME_PREFIX) -> None:
vllm/utils/deep_gemm.py
+120 -4
@@ -70,17 +70,25 @@ def _missing(*_: Any, **__: Any) -> NoReturn:
70 _fp8_gemm_nt_impl: Callable[..., Any] | None = None
71 _grouped_impl: Callable[..., Any] | None = None
72 _grouped_masked_impl: Callable[..., Any] | None = None
73 +_fp8_mqa_logits_impl: Callable[..., Any] | None = None
74 +_fp8_paged_mqa_logits_impl: Callable[..., Any] | None = None
75 +_get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None
76 _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None
77
78
79 def _lazy_init() -> None:
80 """Import deep_gemm and resolve symbols on first use."""
78 - global _fp8_gemm_nt_impl, _grouped_impl, _grouped_masked_impl,\
79 - _get_mn_major_tma_aligned_tensor_impl
81 + global _fp8_gemm_nt_impl, _grouped_impl, _grouped_masked_impl
82 + global _fp8_mqa_logits_impl, _fp8_paged_mqa_logits_impl
83 + global _get_paged_mqa_logits_metadata_impl
84 + global _get_mn_major_tma_aligned_tensor_impl
85
86 # fast path
87 if (_fp8_gemm_nt_impl is not None or _grouped_impl is not None
83 - or _grouped_masked_impl is not None):
88 + or _grouped_masked_impl is not None
89 + or _fp8_mqa_logits_impl is not None
90 + or _fp8_paged_mqa_logits_impl is not None
91 + or _get_paged_mqa_logits_metadata_impl is not None):
92 return
93
94 if not has_deep_gemm():
@@ -97,10 +105,20 @@ def _lazy_init() -> None:
105 _fp8_gemm_nt_impl = getattr(_dg, "fp8_gemm_nt", None)
106 _grouped_impl = getattr(_dg, "m_grouped_fp8_gemm_nt_contiguous", None)
107 _grouped_masked_impl = getattr(_dg, "fp8_m_grouped_gemm_nt_masked", None)
108 + _fp8_mqa_logits_impl = getattr(_dg, "fp8_mqa_logits", None)
109 + _fp8_paged_mqa_logits_impl = getattr(_dg, "fp8_paged_mqa_logits", None)
110 + _get_paged_mqa_logits_metadata_impl = getattr(
111 + _dg, "get_paged_mqa_logits_metadata", None)
112 _get_mn_major_tma_aligned_tensor_impl = getattr(
113 _dg, "get_mn_major_tma_aligned_tensor", None)
114
115
116 +def get_num_sms() -> int:
117 + _lazy_init()
118 + _dg = importlib.import_module("deep_gemm")
119 + return int(_dg.get_num_sms())
120 +
121 +
122 def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor:
123 """Wrapper for DeepGEMM's get_mn_major_tma_aligned_tensor"""
124 _lazy_init()
@@ -135,6 +153,100 @@ def fp8_m_grouped_gemm_nt_masked(*args, **kwargs):
153 *args, disable_ue8m0_cast=not is_deep_gemm_e8m0_used(), **kwargs)
154
155
156 +def fp8_mqa_logits(
157 + q: torch.Tensor,
158 + kv: tuple[torch.Tensor, torch.Tensor],
159 + weights: torch.Tensor,
160 + cu_seqlen_ks: torch.Tensor,
161 + cu_seqlen_ke: torch.Tensor,
162 +) -> torch.Tensor:
163 + """Compute FP8 MQA logits for a single sequence without KV paging.
164 +
165 + Args:
166 + q: Query tensor of shape [M, H, D]. Casted to
167 + `torch.float8_e4m3fn` by caller.
168 + kv: Tuple `(k_fp8, k_scales)` where `k_fp8` has shape [N, D] with
169 + dtype `torch.float8_e4m3fn` and `k_scales` has shape [N] (or
170 + [N, 1]) with dtype `torch.float32`.
171 + weights: weights of shape [M, H], dtype `torch.float32`.
172 + cu_seqlen_ks: Start indices (inclusive) for valid K per query position,
173 + shape [M], dtype int32.
174 + cu_seqlen_ke: End indices (exclusive) for valid K per query position,
175 + shape [M], dtype int32.
176 +
177 + Returns:
178 + Logits tensor of shape [M, N], dtype `torch.float32`.
179 + """
180 + _lazy_init()
181 + if _fp8_mqa_logits_impl is None:
182 + return _missing()
183 + return _fp8_mqa_logits_impl(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke)
184 +
185 +
186 +def get_paged_mqa_logits_metadata(context_lens: torch.Tensor, block_size: int,
187 + num_sms: int) -> torch.Tensor:
188 + """Build scheduling metadata for paged MQA logits.
189 +
190 + Args:
191 + context_lens: Tensor of shape [B], dtype int32; effective context length
192 + per batch element.
193 + block_size: KV-cache block size in tokens (e.g., 64).
194 + num_sms: Number of SMs available. 132 for Hopper
195 +
196 + Returns:
197 + Backend-specific tensor consumed by `fp8_paged_mqa_logits` to
198 + schedule work across SMs.
199 + """
200 + _lazy_init()
201 + if _get_paged_mqa_logits_metadata_impl is None:
202 + return _missing()
203 + return _get_paged_mqa_logits_metadata_impl(context_lens, block_size,
204 + num_sms)
205 +
206 +
207 +def fp8_paged_mqa_logits(
208 + q_fp8: torch.Tensor,
209 + kv_cache_fp8: torch.Tensor,
210 + weights: torch.Tensor,
211 + context_lens: torch.Tensor,
212 + block_tables: torch.Tensor,
213 + schedule_metadata: torch.Tensor,
214 + max_model_len: int,
215 +) -> torch.Tensor:
216 + """Compute FP8 MQA logits using paged KV-cache.
217 +
218 + Args:
219 + q_fp8: Query tensor of shape [B, next_n, H, D]. Casted to
220 + `torch.float8_e4m3fn` by caller.
221 + kv_cache_fp8: Paged KV-cache in packed FP8+scale layout with shape
222 + [num_blocks, block_size, 1, D+4], dtype `torch.uint8`. The last
223 + 4 bytes per (block,pos) store the `float` dequant scale.
224 + weights: Tensor of shape [B * next_n, H], dtype `torch.float32`.
225 + context_lens: Tensor of shape [B], dtype int32; effective context length
226 + for each batch element.
227 + block_tables: Tensor of shape [B, max_blocks], dtype int32; maps logical
228 + block indices to physical blocks in the paged cache.
229 + schedule_metadata: Returned by `get_paged_mqa_logits_metadata`;
230 + used to distribute work across SMs.
231 + max_model_len: Maximum sequence length used to size the logits output.
232 +
233 + Returns:
234 + Logits tensor of shape [B * next_n, max_model_len], dtype
235 + `torch.float32`.
236 + """
237 + _lazy_init()
238 + if _fp8_paged_mqa_logits_impl is None:
239 + return _missing()
240 + return _fp8_paged_mqa_logits_impl(q_fp8,
241 + kv_cache_fp8,
242 + weights,
243 + context_lens,
244 + block_tables,
245 + schedule_metadata,
246 + max_model_len,
247 + clean_logits=True)
248 +
249 +
250 def _ceil_to_ue8m0(x: torch.Tensor):
251 return torch.pow(2.0, torch.ceil(torch.log2(x.abs())))
252
@@ -195,9 +307,13 @@ __all__ = [
307 "fp8_gemm_nt",
308 "m_grouped_fp8_gemm_nt_contiguous",
309 "fp8_m_grouped_gemm_nt_masked",
310 + "fp8_mqa_logits",
311 + "fp8_paged_mqa_logits",
312 + "get_paged_mqa_logits_metadata",
313 "per_block_cast_to_fp8",
314 "is_deep_gemm_e8m0_used",
315 "is_deep_gemm_supported",
316 + "get_num_sms",
317 "should_use_deepgemm_for_fp8_linear",
318 "get_col_major_tma_aligned_tensor",
203 -]
\ No newline at end of file
319 +]
vllm/v1/attention/backends/cpu_attn.py
+1
@@ -74,6 +74,7 @@ class TorchSDPABackend(AttentionBackend):
74 block_size: int,
75 num_kv_heads: int,
76 head_size: int,
77 + cache_dtype_str: str = "auto",
78 ) -> tuple[int, ...]:
79 return _get_paged_attn_impl().get_kv_cache_shape(
80 num_blocks, block_size, num_kv_heads, head_size)
vllm/v1/attention/backends/flash_attn.py
+1
@@ -80,6 +80,7 @@ class FlashAttentionBackend(AttentionBackend):
80 block_size: int,
81 num_kv_heads: int,
82 head_size: int,
83 + cache_dtype_str: str = "auto",
84 ) -> tuple[int, ...]:
85 if block_size % 16 != 0:
86 raise ValueError("Block size must be a multiple of 16.")
vllm/v1/attention/backends/flashinfer.py
+1
@@ -187,6 +187,7 @@ class FlashInferBackend(AttentionBackend):
187 block_size: int,
188 num_kv_heads: int,
189 head_size: int,
190 + cache_dtype_str: str = "auto",
191 ) -> tuple[int, ...]:
192 return (num_blocks, 2, block_size, num_kv_heads, head_size)
193
vllm/v1/attention/backends/flex_attention.py
+1
@@ -90,6 +90,7 @@ class FlexAttentionBackend(AttentionBackend):
90 block_size: int,
91 num_kv_heads: int,
92 head_size: int,
93 + cache_dtype_str: str = "auto",
94 ) -> tuple[int, ...]:
95 return (2, num_blocks, block_size, num_kv_heads, head_size)
96
vllm/v1/attention/backends/mla/common.py
+139 -31
@@ -286,6 +286,7 @@ class MLACommonBackend(AttentionBackend):
286 block_size: int,
287 num_kv_heads: int, # assumed to be 1 for MLA
288 head_size: int,
289 + cache_dtype_str: str = "auto",
290 ) -> tuple[int, ...]:
291 return (num_blocks, block_size, head_size)
292
@@ -407,6 +408,7 @@ class MLACommonMetadata(Generic[D]):
408
409
410 M = TypeVar("M", bound=MLACommonMetadata)
411 +A = TypeVar("A")
412
413
414 def use_flashinfer_prefill() -> bool:
@@ -930,7 +932,9 @@ def reorg_kvcache(
932 return reorganized_kv_c_normed, reorganized_k_pe
933
934
933 -class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
935 +# TODO(Lucas): rename MLACommonBaseImpl -> MLACommonImpl,
936 +# and MLACommonImpl -> MLACommonDenseImpl or somthing like that
937 +class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
938 """
939 NOTE: Please read the comment at the top of the file before trying to
940 understand this class
@@ -956,6 +960,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
960 qk_head_dim: int,
961 v_head_dim: int,
962 kv_b_proj: ColumnParallelLinear,
963 + indexer=None,
964 q_pad_num_heads: Optional[int] = None,
965 ) -> None:
966 if kv_sharing_target_layer_name is not None:
@@ -974,8 +979,140 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
979 self.qk_head_dim = qk_head_dim
980 self.v_head_dim = v_head_dim
981 self.kv_b_proj = kv_b_proj
982 + self.indexer = indexer
983 self.q_pad_num_heads = q_pad_num_heads
984
985 + def process_weights_after_loading(self, act_dtype: torch.dtype):
986 +
987 + def get_layer_weight(layer):
988 + WEIGHT_NAMES = ("weight", "qweight", "weight_packed")
989 + for attr in WEIGHT_NAMES:
990 + if hasattr(layer, attr):
991 + return getattr(layer, attr)
992 + raise AttributeError(
993 + f"Layer '{layer}' has no recognized weight attribute:"
994 + f" {WEIGHT_NAMES}.")
995 +
996 + def get_and_maybe_dequant_weights(layer: LinearBase):
997 + if not isinstance(layer.quant_method, UnquantizedLinearMethod):
998 + # NOTE: This should only be used offline, since it's O(N^3)
999 + eye = torch.eye(layer.input_size_per_partition,
1000 + dtype=act_dtype,
1001 + device=get_layer_weight(layer).device)
1002 + dequant_weights = layer.quant_method.apply(layer,
1003 + eye,
1004 + bias=None)
1005 + del eye
1006 + # standardize to (output, input)
1007 + return dequant_weights.T
1008 + return layer.weight
1009 +
1010 + # we currently do not have quantized bmm's which are needed for
1011 + # `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform
1012 + # the bmm's in 16-bit, the extra memory overhead of this is fairly low
1013 + kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T
1014 + assert kv_b_proj_weight.shape == (
1015 + self.kv_lora_rank,
1016 + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim)), (
1017 + f"{kv_b_proj_weight.shape=}, "
1018 + f"{self.kv_lora_rank=}, "
1019 + f"{self.num_heads=}, "
1020 + f"{self.qk_nope_head_dim=}, "
1021 + f"{self.v_head_dim=}")
1022 + kv_b_proj_weight = kv_b_proj_weight.view(
1023 + self.kv_lora_rank,
1024 + self.num_heads,
1025 + self.qk_nope_head_dim + self.v_head_dim,
1026 + )
1027 +
1028 + W_UK, W_UV = kv_b_proj_weight.split(
1029 + [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
1030 +
1031 + if is_rocm_aiter_fp8bmm_enabled():
1032 + W_K = W_UK.transpose(0, 1) # 16 512 128
1033 + W_V = W_UV.permute(1, 2, 0) # 16 128 512
1034 + self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant(
1035 + W_K, dtype=current_platform.fp8_dtype())
1036 + self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant(
1037 + W_V, dtype=current_platform.fp8_dtype())
1038 +
1039 + # The kernel operates on non-padded inputs. Hence, pre-compiling
1040 + # triton kernel to avoid runtime compilation for unseen batch sizes
1041 + # Pre-compile for batch sizes 1 to 1024 to cover most use-cases.
1042 + # On DS-R1, this step adds roughly 50s to the model loading time.
1043 + max_batch_size = 1024 # [ToDo] Find the optimal upper limit
1044 + pre_compilation_list = list(range(1, max_batch_size + 1))
1045 + if is_global_first_rank():
1046 + pre_compilation_list = tqdm(
1047 + pre_compilation_list,
1048 + desc="[Aiter Triton] Pre-compiling fp8 BMM kernel",
1049 + total=max_batch_size,
1050 + )
1051 +
1052 + for m in pre_compilation_list:
1053 + x = torch.empty((self.W_K.shape[0], m, self.W_K.shape[2]),
1054 + dtype=torch.bfloat16,
1055 + device=self.W_K.device)
1056 + aiter_triton_fp8_bmm(x,
1057 + self.W_K,
1058 + self.W_K_scale,
1059 + group_size=128,
1060 + transpose_bm=True)
1061 +
1062 + x = torch.empty((self.W_V.shape[0], m, self.W_V.shape[2]),
1063 + dtype=torch.bfloat16,
1064 + device=self.W_V.device)
1065 + aiter_triton_fp8_bmm(x,
1066 + self.W_V,
1067 + self.W_V_scale,
1068 + group_size=128,
1069 + transpose_bm=True)
1070 + else:
1071 + # Convert from (L, N, V) to (N, L, V)
1072 + self.W_UV = W_UV.transpose(0, 1)
1073 + # Convert from (L, N, P) to (N, P, L)
1074 + self.W_UK_T = W_UK.permute(1, 2, 0)
1075 +
1076 + def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
1077 + # Convert from (B, N, L) to (N, B, L)
1078 + x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
1079 + if is_rocm_aiter_fp8bmm_enabled():
1080 + # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
1081 + x = aiter_triton_fp8_bmm(x,
1082 + self.W_V,
1083 + self.W_V_scale,
1084 + group_size=128,
1085 + transpose_bm=True)
1086 + # Convert from (B, N, V) to (B, N * V)
1087 + x = x.reshape(-1, self.num_heads * self.v_head_dim)
1088 + # Copy result
1089 + out.copy_(x)
1090 + else:
1091 + # Convert from (B, N * V) to (N, B, V)
1092 + out = out.view(-1, self.num_heads, self.v_head_dim).transpose(0, 1)
1093 +
1094 + # Multiply (N, B, L) x (N, L, V) -> (N, B, V)
1095 + torch.bmm(x, self.W_UV, out=out) # Reuse "out" to make it "hot"
1096 +
1097 + # Convert from (N, B, V) to (B, N * V)
1098 + out_new = out.transpose(0, 1).reshape(
1099 + -1, self.num_heads * self.v_head_dim)
1100 +
1101 + # Adjust output buffer shape back to the original (B, N * V)
1102 + N, B, V = out.shape
1103 + out.resize_((B, N * V))
1104 + out.copy_(out_new) # Copy result
1105 +
1106 +
1107 +class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1108 + """
1109 + NOTE: Please read the comment at the top of the file before trying to
1110 + understand this class
1111 + """
1112 +
1113 + def __init__(self, *args, **kwargs) -> None:
1114 + super().__init__(*args, **kwargs)
1115 +
1116 if use_flashinfer_prefill():
1117 logger.debug_once("Using FlashInfer prefill for MLA")
1118 self._run_prefill_context_chunk = self._run_prefill_context_chunk_fi
@@ -1154,36 +1291,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1291 True, #Indicates actual_seq_lens are on GPU or CPU.
1292 )
1293
1157 - def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
1158 - # Convert from (B, N, L) to (N, B, L)
1159 - x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
1160 - if is_rocm_aiter_fp8bmm_enabled():
1161 - # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
1162 - x = aiter_triton_fp8_bmm(x,
1163 - self.W_V,
1164 - self.W_V_scale,
1165 - group_size=128,
1166 - transpose_bm=True)
1167 - # Convert from (B, N, V) to (B, N * V)
1168 - x = x.reshape(-1, self.num_heads * self.v_head_dim)
1169 - # Copy result
1170 - out.copy_(x)
1171 - else:
1172 - # Convert from (B, N * V) to (N, B, V)
1173 - out = out.view(-1, self.num_heads, self.v_head_dim).transpose(0, 1)
1174 -
1175 - # Multiply (N, B, L) x (N, L, V) -> (N, B, V)
1176 - torch.bmm(x, self.W_UV, out=out) # Reuse "out" to make it "hot"
1177 -
1178 - # Convert from (N, B, V) to (B, N * V)
1179 - out_new = out.transpose(0, 1).reshape(
1180 - -1, self.num_heads * self.v_head_dim)
1181 -
1182 - # Adjust output buffer shape back to the original (B, N * V)
1183 - N, B, V = out.shape
1184 - out.resize_((B, N * V))
1185 - out.copy_(out_new) # Copy result
1186 -
1294 def process_weights_after_loading(self, act_dtype: torch.dtype):
1295
1296 def get_layer_weight(layer):
@@ -1455,6 +1562,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1562 attn_metadata: MLACommonMetadata,
1563 k_scale: torch.Tensor,
1564 ) -> torch.Tensor:
1565 + # TODO (zyongye): Prefill function here
1566 assert attn_metadata.prefill is not None
1567 assert self.dcp_world_size is not None
1568
vllm/v1/attention/backends/mla/flashmla.py
+1
@@ -177,6 +177,7 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
177 attn_metadata: FlashMLAMetadata,
178 layer: AttentionLayer,
179 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
180 + # TODO: (zyongye) decode function for mla here
181 assert kv_c_and_k_pe_cache.numel() > 0
182 assert attn_metadata.decode is not None
183
vllm/v1/attention/backends/mla/flashmla_sparse.py new
+544
@@ -0,0 +1,544 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import math
4 +from dataclasses import dataclass
5 +from typing import TYPE_CHECKING, ClassVar, Optional
6 +
7 +import numpy as np
8 +import torch
9 +
10 +from vllm import _custom_ops as ops
11 +from vllm.attention.backends.abstract import (AttentionBackend, AttentionLayer,
12 + AttentionMetadata)
13 +from vllm.attention.backends.utils import get_mla_dims
14 +from vllm.attention.ops.flashmla import (flash_mla_sparse_prefill,
15 + flash_mla_with_kvcache,
16 + get_mla_metadata)
17 +from vllm.config import VllmConfig
18 +from vllm.logger import init_logger
19 +from vllm.platforms import current_platform
20 +from vllm.triton_utils import tl, triton
21 +from vllm.utils import cdiv
22 +from vllm.v1.attention.backends.mla.common import MLACommonBaseImpl
23 +from vllm.v1.attention.backends.utils import (AttentionCGSupport,
24 + AttentionMetadataBuilder,
25 + CommonAttentionMetadata)
26 +from vllm.v1.kv_cache_interface import AttentionSpec
27 +
28 +if TYPE_CHECKING:
29 + from vllm.model_executor.models.deepseek_v2 import Indexer
30 +
31 +logger = init_logger(__name__)
32 +"""
33 +NOTE: FlashMLA Sparse uses an fp8 cache with the following format
34 +
35 +In the "FP8 with scale" format, each token's KV cache is 656 Bytes,
36 +structured as:
37 +- **First 512 bytes:** The "quantized NoPE" part, containing 512
38 + `float8_e4m3` values.
39 +- **Next 16 bytes:** Scale factors, containing 4 `float32` values.
40 + The first `float32` is the scale for the first 128 `float8_e4m3` values,
41 + the second for the next 128, and so on.
42 +- **Last 128 bytes:** The "RoPE" part, containing 64 `bfloat16` values. This
43 + part is not quantized for accuracy.
44 +"""
45 +
46 +
47 +def _lse2_to_lse(lse_base2: torch.Tensor) -> torch.Tensor:
48 + # Convert base-2 LSE to natural-log LSE
49 + # Keep FP32 for numerical stability during the merge.
50 + return (lse_base2.to(torch.float32) * math.log(2.0))
51 +
52 +
53 +class FlashMLASparseBackend(AttentionBackend):
54 +
55 + accept_output_buffer: bool = True
56 +
57 + @staticmethod
58 + def get_name() -> str:
59 + return "FLASHMLA_SPARSE_VLLM_V1"
60 +
61 + @staticmethod
62 + def get_metadata_cls() -> type[AttentionMetadata]:
63 + return FlashMLASparseMetadata
64 +
65 + @staticmethod
66 + def get_builder_cls() -> type["FlashMLASparseMetadataBuilder"]:
67 + return FlashMLASparseMetadataBuilder
68 +
69 + @staticmethod
70 + def get_impl_cls() -> type["FlashMLASparseImpl"]:
71 + return FlashMLASparseImpl
72 +
73 + @staticmethod
74 + def get_kv_cache_shape(
75 + num_blocks: int,
76 + block_size: int,
77 + num_kv_heads: int, # assumed to be 1 for MLA
78 + head_size: int,
79 + cache_dtype_str: str = "auto",
80 + ) -> tuple[int, ...]:
81 + if cache_dtype_str == "fp8_ds_mla":
82 + # custom storage fromat is 656 bytes
83 + # see FlashMLA readme.md for details
84 + return (num_blocks, block_size, 656)
85 + else:
86 + return (num_blocks, block_size, head_size)
87 +
88 + @classmethod
89 + def get_supported_dtypes(cls) -> list[torch.dtype]:
90 + return [torch.bfloat16]
91 +
92 + @classmethod
93 + def get_supported_head_sizes(cls) -> list[int]:
94 + return [576]
95 +
96 +
97 +@dataclass
98 +class MLASparsePrefillMetadata:
99 + # NOTE(Chen): not call it "FlashMLASparsePrefillMetadata" because
100 + # the kernel is not from flashmla
101 + block_table: torch.Tensor
102 + has_context: bool = False
103 + context_lens: Optional[torch.Tensor] = None
104 +
105 +
106 +@dataclass
107 +class FlashMLASparseDecodeAndContextMetadata:
108 + scheduler_metadata: torch.Tensor = None
109 + num_splits: torch.Tensor = None
110 + cache_lens: torch.Tensor = None
111 + prefill_context_lengths: Optional[torch.Tensor] = None
112 + prefill_new_k_start_locs: Optional[torch.Tensor] = None
113 + dummy_block_table: torch.Tensor = None
114 +
115 + def filter_prefill_indices(
116 + self, indices: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
117 + assert self.prefill_context_lengths is not None
118 + prefill_context_lengths = self.prefill_context_lengths.unsqueeze(-1)
119 + context_indices = torch.where(indices < prefill_context_lengths,
120 + indices, -1)
121 + new_token_indices = torch.where(indices >= prefill_context_lengths,
122 + indices - prefill_context_lengths, -1)
123 + return context_indices, new_token_indices
124 +
125 +
126 +@dataclass
127 +class FlashMLASparseMetadata:
128 + num_reqs: int
129 + max_query_len: int
130 + max_seq_len: int
131 +
132 + num_actual_tokens: int # Number of tokens excluding padding.
133 + query_start_loc: torch.Tensor
134 + slot_mapping: torch.Tensor
135 +
136 + block_table: torch.Tensor
137 + req_id_per_token: torch.Tensor
138 + block_size: int = 64
139 + topk_tokens: int = 2048
140 +
141 + @dataclass
142 + class FP8KernelMetadata:
143 + scheduler_metadata: Optional[torch.Tensor]
144 + num_splits: torch.Tensor
145 + dummy_block_table: torch.Tensor
146 + cache_lens: torch.Tensor
147 +
148 + fp8_extra_metadata: Optional[FP8KernelMetadata] = None
149 +
150 +
151 +@triton.jit
152 +def _convert_req_index_to_global_index_kernel(
153 + req_id_ptr, # int32 [num_tokens]
154 + block_table_ptr, # int32 [num_requests, max_num_blocks_per_req]
155 + token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]
156 + out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]
157 + # shapes (compile-time where possible)
158 + max_num_blocks_per_req: tl.constexpr,
159 + BLOCK_SIZE: tl.constexpr,
160 + BLOCK_N: tl.constexpr, # tile width along columns
161 + # strides (in elements)
162 + bt_stride0,
163 + bt_stride1,
164 + ti_stride0,
165 + ti_stride1,
166 + out_stride0,
167 + out_stride1,
168 +):
169 + # program_id(0) -> token_id (row)
170 + # program_id(1) -> tile index along columns
171 + token_id = tl.program_id(0)
172 + tile_id = tl.program_id(1)
173 +
174 + # Each program covers BLOCK_N consecutive columns
175 + indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N)
176 +
177 + # Load request id for this token (no mask: grid is exact)
178 + req = tl.load(req_id_ptr + token_id)
179 +
180 + # Load token indices for this tile
181 + ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1
182 + tok = tl.load(ti_ptr) # int32
183 +
184 + # Only token == -1 should propagate as -1
185 + is_invalid_tok = tok < 0
186 +
187 + # Compute block id and in-block offset
188 + block_id = tok // BLOCK_SIZE
189 + inblock_off = tok % BLOCK_SIZE
190 +
191 + # Guard block_table access
192 + valid_block = block_id < max_num_blocks_per_req
193 + bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1
194 + base = tl.load(bt_ptr, mask=valid_block, other=0)
195 +
196 + # If token == -1 OR block_id OOB, output -1; else base * BLOCK_SIZE + offset
197 + out_val = tl.where(is_invalid_tok | (~valid_block), -1,
198 + base * BLOCK_SIZE + inblock_off)
199 +
200 + # Store results
201 + out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1
202 + tl.store(out_ptr_ij, out_val)
203 +
204 +
205 +def triton_convert_req_index_to_global_index(
206 + req_id: torch.Tensor, # int32 [num_tokens]
207 + block_table: torch.
208 + Tensor, # int32 [num_requests, max_num_blocks_per_req]
209 + token_indices: torch.Tensor, # int32 [num_tokens, NUM_TOPK_TOKENS]
210 + BLOCK_SIZE: int = 64,
211 + NUM_TOPK_TOKENS: int = 2048,
212 + BLOCK_N: int = 128, # tile width along columns
213 +):
214 + """
215 + out[token_id, indice_id] =
216 + block_table[req_id[token_id],
217 + token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE
218 + + token_indices[token_id, indice_id] % BLOCK_SIZE
219 +
220 + Only when token_indices[token_id, indice_id] == -1 do we output -1.
221 + For safety, we also output -1 if the derived block_id would be
222 + out-of-bounds.
223 + """
224 + assert req_id.dtype == torch.int32
225 + assert block_table.dtype == torch.int32
226 + assert token_indices.dtype == torch.int32
227 + assert token_indices.shape[1] == NUM_TOPK_TOKENS
228 + assert NUM_TOPK_TOKENS % BLOCK_N == 0, \
229 + f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible by" \
230 + f"BLOCK_N ({BLOCK_N})"
231 +
232 + num_tokens = req_id.shape[0]
233 + num_requests, max_num_blocks_per_req = block_table.shape
234 + tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N
235 +
236 + # Ensure contiguous tensors on the same device
237 + req_id_c = req_id.contiguous()
238 + block_table_c = block_table.contiguous()
239 + token_indices_c = token_indices.contiguous()
240 + out = torch.empty_like(token_indices_c)
241 +
242 + # Strides in elements
243 + bt_stride0, bt_stride1 = block_table_c.stride()
244 + ti_stride0, ti_stride1 = token_indices_c.stride()
245 + out_stride0, out_stride1 = out.stride()
246 +
247 + # Exact 2D grid: tokens × column tiles
248 + grid = (num_tokens, tiles_per_row)
249 +
250 + _convert_req_index_to_global_index_kernel[grid](
251 + req_id_c,
252 + block_table_c,
253 + token_indices_c,
254 + out,
255 + # shapes / constexprs
256 + max_num_blocks_per_req,
257 + BLOCK_SIZE,
258 + BLOCK_N,
259 + # strides
260 + bt_stride0,
261 + bt_stride1,
262 + ti_stride0,
263 + ti_stride1,
264 + out_stride0,
265 + out_stride1,
266 + )
267 + return out
268 +
269 +
270 +@dataclass
271 +class FlashMLASparseMetadataBuilder(
272 + AttentionMetadataBuilder[FlashMLASparseMetadata]):
273 + cudagraph_support: ClassVar[AttentionCGSupport] = \
274 + AttentionCGSupport.UNIFORM_BATCH
275 +
276 + def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
277 + vllm_config: VllmConfig, device: torch.device):
278 +
279 + cache_config = vllm_config.cache_config
280 + self.kv_cache_spec = kv_cache_spec
281 + self.model_config = vllm_config.model_config
282 + parallel_config = vllm_config.parallel_config
283 + self.device = device
284 +
285 + props = torch.cuda.get_device_properties(device)
286 + sm_count = props.multi_processor_count
287 +
288 + self.num_heads = self.model_config.get_num_attention_heads(
289 + parallel_config)
290 + self.mla_dims = get_mla_dims(self.model_config)
291 + self.topk_tokens = vllm_config.model_config.hf_config.index_topk
292 + self.use_fp8_kv_cache = cache_config.cache_dtype == "fp8_ds_mla"
293 + self.topk_tokens_tensor = torch.tensor([self.topk_tokens],
294 + device=device,
295 + dtype=torch.int32)
296 + self.max_model_len_tensor = torch.tensor(
297 + [self.model_config.max_model_len],
298 + device=device,
299 + dtype=torch.int32)
300 + # this is ignored by `flash_mla_with_kvcache` if indices not None
301 + self.dummy_block_table = torch.empty((1, 1),
302 + dtype=torch.int32,
303 + device=self.device)
304 +
305 + # Equation taken from FlashMLA/csrc/pybind.cpp
306 + h_q, h_k = self.num_heads, 1
307 + s_q = 1 # inversely proportional to s_q, so s_q = 1 is the largest
308 + max_num_sm_parts = int(
309 + max((sm_count // 2) / h_k // (cdiv(h_q // h_k, 2 * 64) * s_q), 1))
310 + if current_platform.is_device_capability(100):
311 + max_num_sm_parts *= 2
312 + self.tile_scheduler_metadata_buffer = torch.empty(
313 + # TileSchedulerMetaDataSize = 8
314 + # see: FlashMLA/csrc/params.h
315 + (max_num_sm_parts, 8),
316 + dtype=torch.int32,
317 + device=device)
318 + self.num_splits_buffer = torch.empty(
319 + # We pack all the tokens into one batch for sparse attention.
320 + # Otherwise, we can exceed the sm of `get_mla_metadata`.
321 + (
322 + 2, ),
323 + dtype=torch.int32,
324 + device=device)
325 + self.req_id_per_token_buffer = torch.empty(
326 + (vllm_config.scheduler_config.max_num_batched_tokens, ),
327 + dtype=torch.int32,
328 + device=device)
329 +
330 + def build(self,
331 + common_prefix_len: int,
332 + common_attn_metadata: CommonAttentionMetadata,
333 + fast_build: bool = False) -> FlashMLASparseMetadata:
334 +
335 + num_tokens = common_attn_metadata.num_actual_tokens
336 + starts = np.asarray(common_attn_metadata.query_start_loc_cpu,
337 + dtype=np.int32)
338 + seg_lengths = np.diff(starts)
339 + req_id_per_token = np.repeat(
340 + np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths)
341 + # Zero-fill for cudagraphs
342 + self.req_id_per_token_buffer.fill_(0)
343 + self.req_id_per_token_buffer[:req_id_per_token.shape[0]]\
344 + .copy_(torch.from_numpy(req_id_per_token), non_blocking=True)
345 + req_id_per_token = self.req_id_per_token_buffer[:num_tokens]
346 +
347 + fp8_extra_metadata = None
348 + if self.use_fp8_kv_cache:
349 + tile_scheduler_metadata, num_splits = get_mla_metadata(
350 + cache_seqlens=self.topk_tokens_tensor,
351 + num_q_tokens_per_head_k=num_tokens * self.num_heads,
352 + topk=self.topk_tokens,
353 + num_heads_q=self.num_heads,
354 + num_heads_k=1,
355 + is_fp8_kvcache=True,
356 + )
357 +
358 + num_sm_parts = tile_scheduler_metadata.size(0)
359 + # Copy to persistent buffer for full-CG support
360 + tile_scheduler_metadata_buffer = \
361 + self.tile_scheduler_metadata_buffer[:num_sm_parts]
362 + tile_scheduler_metadata_buffer.copy_(tile_scheduler_metadata)
363 + self.num_splits_buffer.copy_(num_splits)
364 +
365 + fp8_extra_metadata = FlashMLASparseMetadata.FP8KernelMetadata(
366 + scheduler_metadata=tile_scheduler_metadata_buffer,
367 + num_splits=self.num_splits_buffer,
368 + # cache_lens and block_table are basically unused in sparse case
369 + # but the decode kernel will treat -1 and indices >= cache_lens
370 + # as invalid so we make sure cache_lens is large enough to not
371 + # accidentally mark indices invalid, we will use -1 exclusively
372 + # to mark invalid indices
373 + cache_lens=self.max_model_len_tensor,
374 + dummy_block_table=self.dummy_block_table)
375 +
376 + metadata = FlashMLASparseMetadata(
377 + num_reqs=common_attn_metadata.num_reqs,
378 + max_query_len=common_attn_metadata.max_query_len,
379 + max_seq_len=common_attn_metadata.max_seq_len,
380 + num_actual_tokens=common_attn_metadata.num_actual_tokens,
381 + query_start_loc=common_attn_metadata.query_start_loc,
382 + slot_mapping=common_attn_metadata.slot_mapping,
383 + block_table=common_attn_metadata.block_table_tensor,
384 + req_id_per_token=req_id_per_token,
385 + block_size=self.kv_cache_spec.block_size,
386 + topk_tokens=self.topk_tokens,
387 + fp8_extra_metadata=fp8_extra_metadata,
388 + )
389 + return metadata
390 +
391 +
392 +class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):
393 +
394 + def __init__(
395 + self,
396 + num_heads: int,
397 + head_size: int,
398 + scale: float,
399 + num_kv_heads: int,
400 + alibi_slopes: Optional[list[float]],
401 + sliding_window: Optional[int],
402 + kv_cache_dtype: str,
403 + logits_soft_cap: Optional[float],
404 + attn_type: str,
405 + kv_sharing_target_layer_name: Optional[str],
406 + # MLA Specific Arguments
407 + topk_indice_buffer: Optional[torch.Tensor] = None,
408 + indexer: Optional["Indexer"] = None,
409 + **mla_args) -> None:
410 + super().__init__(num_heads, head_size, scale, num_kv_heads,
411 + alibi_slopes, sliding_window, kv_cache_dtype,
412 + logits_soft_cap, attn_type,
413 + kv_sharing_target_layer_name, **mla_args)
414 + self.softmax_scale = scale
415 + assert indexer is not None
416 + self.topk_indices_buffer = indexer.topk_indices_buffer
417 + self.padding = 128 if current_platform.is_device_capability(
418 + 100) else 64
419 +
420 + def _forward_bf16_kv(
421 + self, q: torch.Tensor, kv_c_and_k_pe_cache: torch.Tensor,
422 + topk_indices: torch.Tensor,
423 + attn_metadata: FlashMLASparseMetadata) -> torch.Tensor:
424 + num_tokens = q.shape[0]
425 + kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(
426 + -1, 1, kv_c_and_k_pe_cache.shape[-1])
427 +
428 + # NOTE(Chen): kernel requires num_local_head to be a multiple of
429 + # 64 on hopper and 128 on blackwell
430 + if self.num_heads % self.padding != 0:
431 + assert self.padding % self.num_heads == 0
432 + logger.warning_once(f"padding num_heads to {self.padding} \
433 + due to sparse attn kernel requirement")
434 + q_padded = q.new_empty((q.shape[0], self.padding, q.shape[2]))
435 + q_padded[:, :self.num_heads, :] = q
436 + q = q_padded
437 +
438 + topk_indices = topk_indices.view(num_tokens, 1, -1)
439 + output = flash_mla_sparse_prefill(q, kv_c_and_k_pe_cache, topk_indices,
440 + self.softmax_scale)[0]
441 + output = output[:, :self.num_heads, :]
442 + return output
443 +
444 + def _forward_fp8_kv(self, q: torch.Tensor,
445 + kv_c_and_k_pe_cache: torch.Tensor,
446 + topk_indices: torch.Tensor,
447 + attn_metadata: FlashMLASparseMetadata) -> torch.Tensor:
448 +
449 + assert attn_metadata.fp8_extra_metadata is not None
450 + extra_metadata = attn_metadata.fp8_extra_metadata
451 +
452 + _attn_out, _ = flash_mla_with_kvcache(
453 + q=q.unsqueeze(0), # unsqueeze to add batch_dim
454 + k_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(-2),
455 + block_table=extra_metadata.dummy_block_table,
456 + head_dim_v=512,
457 + cache_seqlens=extra_metadata.cache_lens,
458 + tile_scheduler_metadata=extra_metadata.scheduler_metadata,
459 + num_splits=extra_metadata.num_splits,
460 + is_fp8_kvcache=True,
461 + indices=topk_indices.unsqueeze(0), # unsqueeze to add batch_dim
462 + softmax_scale=self.softmax_scale,
463 + )
464 +
465 + return _attn_out
466 +
467 + def forward(
468 + self,
469 + layer: AttentionLayer,
470 + q: torch.Tensor,
471 + k_c_normed: torch.Tensor, # key in unified attn
472 + k_pe: torch.Tensor, # value in unified attn
473 + kv_cache: torch.Tensor,
474 + attn_metadata: FlashMLASparseMetadata,
475 + output: Optional[torch.Tensor] = None,
476 + output_scale: Optional[torch.Tensor] = None,
477 + output_block_scale: Optional[torch.Tensor] = None,
478 + ) -> torch.Tensor:
479 + # NOTE(lucas): for the sparse FlashMLA kernels the kernels want to use
480 + # MQA 576/512 approach for both prefill and decode
481 +
482 + assert output is not None, "Output tensor must be provided."
483 +
484 + if output_scale is not None or output_block_scale is not None:
485 + raise NotImplementedError(
486 + "fused output quantization is not yet supported"
487 + " for MLACommonImpl")
488 +
489 + if attn_metadata is None:
490 + # The zero fill is required when used with DP + EP
491 + # to ensure all ranks within a DP group compute the
492 + # same expert outputs.
493 + return output.fill_(0)
494 +
495 + num_actual_toks = attn_metadata.num_actual_tokens
496 +
497 + # Inputs and outputs may be padded for CUDA graphs
498 +
499 + q = q[:num_actual_toks, ...]
500 + k_c_normed = k_c_normed[:num_actual_toks, ...]
501 + k_pe = k_pe[:num_actual_toks, ...]
502 +
503 + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim],
504 + dim=-1)
505 + # Convert from (B, N, P) to (N, B, P)
506 + q_nope = q_nope.transpose(0, 1)
507 + # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
508 + ql_nope = torch.bmm(q_nope, self.W_UK_T)
509 + # Convert from (N, B, L) to (B, N, L)
510 + ql_nope = ql_nope.transpose(0, 1)
511 +
512 + topk_indices = self.topk_indices_buffer[:num_actual_toks]
513 +
514 + # TODO: handle index / kv_cache correctly
515 + topk_indices_global = triton_convert_req_index_to_global_index(
516 + attn_metadata.req_id_per_token,
517 + attn_metadata.block_table,
518 + topk_indices,
519 + BLOCK_SIZE=attn_metadata.block_size,
520 + NUM_TOPK_TOKENS=attn_metadata.topk_tokens,
521 + )
522 +
523 + q = torch.cat([ql_nope, q_pe], dim=-1)
524 +
525 + # write the latent and rope to kv cache
526 + if kv_cache.numel() > 0:
527 + ops.concat_and_cache_mla(
528 + k_c_normed,
529 + k_pe.squeeze(1),
530 + kv_cache,
531 + attn_metadata.slot_mapping.flatten(),
532 + kv_cache_dtype=self.kv_cache_dtype,
533 + scale=layer._k_scale,
534 + )
535 +
536 + if self.kv_cache_dtype != "fp8_ds_mla":
537 + attn_out = self._forward_bf16_kv(q, kv_cache, topk_indices_global,
538 + attn_metadata)
539 + else:
540 + attn_out = self._forward_fp8_kv(q, kv_cache, topk_indices_global,
541 + attn_metadata)
542 +
543 + self._v_up_proj(attn_out, out=output[:num_actual_toks])
544 + return output
vllm/v1/attention/backends/mla/indexer.py new
+293
@@ -0,0 +1,293 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from dataclasses import dataclass
4 +from typing import ClassVar, Optional
5 +
6 +import torch
7 +
8 +from vllm.attention.backends.abstract import (AttentionBackend,
9 + AttentionMetadata)
10 +from vllm.config import VllmConfig
11 +from vllm.logger import init_logger
12 +from vllm.utils.deep_gemm import get_paged_mqa_logits_metadata
13 +from vllm.v1.attention.backends.utils import (AttentionCGSupport,
14 + AttentionMetadataBuilder,
15 + CommonAttentionMetadata,
16 + split_decodes_and_prefills)
17 +
18 +logger = init_logger(__name__)
19 +
20 +
21 +class DeepseekV32IndexerBackend(AttentionBackend):
22 +
23 + @staticmethod
24 + def get_metadata_cls() -> type["AttentionMetadata"]:
25 + return DeepseekV32IndexerMetadata
26 +
27 + @classmethod
28 + def get_supported_head_sizes(cls) -> list[int]:
29 + return [32, 64, 128]
30 +
31 + @staticmethod
32 + def get_builder_cls() -> type["DeepseekV32IndexerMetadataBuilder"]:
33 + return DeepseekV32IndexerMetadataBuilder
34 +
35 + @staticmethod
36 + def get_kv_cache_shape(
37 + num_blocks: int,
38 + block_size: int,
39 + num_kv_heads: int,
40 + head_size: int,
41 + cache_dtype_str: str = "auto",
42 + ) -> tuple[int, ...]:
43 + assert num_kv_heads == 1
44 + return (num_blocks, block_size, head_size)
45 +
46 + @staticmethod
47 + def get_kv_cache_stride_order() -> tuple[int, ...]:
48 + return (0, 1, 2)
49 +
50 +
51 +@dataclass
52 +class DeepseekV32IndexerPrefillMetadata:
53 + block_table: torch.Tensor
54 + query_start_loc: torch.Tensor
55 + max_query_len: int
56 + cu_seqlen_ks: torch.Tensor
57 + cu_seqlen_ke: torch.Tensor
58 + cu_seq_lens: torch.Tensor
59 + total_seq_lens: int
60 +
61 +
62 +@dataclass
63 +class DeepSeekV32IndexerDecodeMetadata:
64 + block_table: torch.Tensor
65 + seq_lens: torch.Tensor
66 + decode_lens: torch.Tensor
67 + requires_padding: bool
68 + schedule_metadata: torch.Tensor
69 +
70 +
71 +@dataclass
72 +class DeepseekV32IndexerMetadata:
73 +
74 + # FIXME (zyongye)
75 + # hacky way to access the data now, need to be in chunked meta
76 + seq_lens: torch.Tensor
77 +
78 + num_reqs: int
79 + max_query_len: int
80 + max_seq_len: int
81 +
82 + num_actual_tokens: int # Number of tokens excluding padding.
83 + query_start_loc: torch.Tensor
84 + slot_mapping: torch.Tensor
85 + # The dimension of the attention heads
86 + head_dim: int
87 +
88 + # New for MLA (compared to FlashAttention)
89 + # For handling prefill decode split
90 + num_decodes: int
91 + num_decode_tokens: int
92 + num_prefills: int
93 + num_prefill_tokens: int
94 +
95 + decode: Optional[DeepSeekV32IndexerDecodeMetadata] = None
96 + prefill: Optional[DeepseekV32IndexerPrefillMetadata] = None
97 +
98 +
99 +# TODO (zyongye) optimize this, this is now vibe coded
100 +def kv_spans_from_batches(
101 + start_seq_loc: torch.Tensor,
102 + seq_len_per_batch: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
103 + """
104 + Args:
105 + start_seq_loc: 1D long tensor [B+1], cumulative counts of
106 + selected tokens per batch.
107 + Example: [0, 2, 4, 7] ->
108 + batch sizes (selected) [2, 2, 3], N=7 tokens total.
109 + seq_len_per_batch: 1D long tensor [B],
110 + full sequence length (KV length) of each batch.
111 + Example: [5, 9, 4].
112 +
113 + Returns:
114 + start_tensor: 1D long tensor [N], start offset in the
115 + concatenated KV cache for each token's batch.
116 + end_location: 1D long tensor [N],
117 + **exclusive** end = start + token's local position.
118 + (So the attended KV slice is kv[start:end].)
119 +
120 + Assumes each batch contributes its full `seq_len_per_batch[i]`
121 + keys to the KV cache, andthe selected tokens within a batch
122 + are the **last** `counts[i]` positions of that sequence.
123 + """
124 + q = start_seq_loc.to(dtype=torch.long)
125 + L = seq_len_per_batch.to(dtype=torch.long, device=q.device)
126 + assert q.dim() == 1 and L.dim() == 1
127 + assert q.numel() == L.numel() + 1, "start_seq_loc must have length B+1"
128 +
129 + # Selected tokens per batch and totals
130 + counts = q[1:] - q[:-1] # [B]
131 + N = int(q[-1].item()) # total selected tokens
132 + B = L.numel()
133 + device = L.device
134 +
135 + if N == 0:
136 + return (torch.empty(0, dtype=torch.long, device=device),
137 + torch.empty(0, dtype=torch.long, device=device))
138 +
139 + # KV start offsets per batch in the concatenated KV cache
140 + kv_starts_per_batch = torch.cumsum(L, dim=0) - L # [B]
141 +
142 + # For each selected token, which batch does it belong to?
143 + batch_id = torch.repeat_interleave(torch.arange(B, device=device),
144 + counts) # [N]
145 +
146 + # Map batch KV start to each token
147 + start_tensor = kv_starts_per_batch[batch_id] # [N]
148 +
149 + # End-align local positions inside each batch:
150 + # local_pos = L[b] - counts[b] + (1..counts[b]) for each batch b
151 + L_expand = torch.repeat_interleave(L, counts) # [N]
152 + m_expand = torch.repeat_interleave(counts, counts) # [N]
153 + # position within the selected block: 1..counts[b]
154 + pos_within = (torch.arange(N, device=device, dtype=torch.long) -
155 + torch.repeat_interleave(q[:-1], counts) + 1)
156 +
157 + local_pos = L_expand - m_expand + pos_within # [N], 1-based
158 + end_location = start_tensor + local_pos # exclusive end
159 +
160 + return start_tensor.int(), end_location.int()
161 +
162 +
163 +def get_max_prefill_buffer_size(vllm_config: VllmConfig):
164 + max_model_len = vllm_config.model_config.max_model_len
165 + # max_num_batched_tokens = \
166 + # vllm_config.scheduler_config.max_num_batched_tokens
167 + max_num_seq = vllm_config.scheduler_config.max_num_seqs
168 + # NOTE(Chen): an estimated max size of flattened_kv. Need to double check.
169 + return max_model_len * max_num_seq
170 +
171 +
172 +class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
173 + cudagraph_support: ClassVar[AttentionCGSupport] = \
174 + AttentionCGSupport.UNIFORM_BATCH
175 +
176 + reorder_batch_threshold: int = 1
177 +
178 + def __init__(self, *args, **kwargs):
179 + super().__init__(*args, **kwargs)
180 + scheduler_config = self.vllm_config.scheduler_config
181 + #NOTE(Chen):an estimated max size of flattened_kv. Need to double check.
182 + self.max_prefill_buffer_size = get_max_prefill_buffer_size(
183 + self.vllm_config)
184 + self.num_speculative_tokens = (
185 + self.vllm_config.speculative_config.num_speculative_tokens
186 + if self.vllm_config.speculative_config else 0)
187 + # Now deepgemm fp8_paged_mqa_logits does not support next_n > 2
188 + self.reorder_batch_threshold += min(self.num_speculative_tokens, 1)
189 +
190 + props = torch.cuda.get_device_properties(self.device)
191 + sm_count = props.multi_processor_count
192 + self.num_sms = sm_count
193 +
194 + self.decode_lens_buffer = torch.empty(
195 + (scheduler_config.max_num_seqs, ),
196 + dtype=torch.int32,
197 + device=self.device)
198 +
199 + # See: DeepGMM/csrc/apis/attention.hpp
200 + self.scheduler_metadata_buffer = torch.empty((self.num_sms + 1, 2),
201 + dtype=torch.int32,
202 + device=self.device)
203 +
204 + def build(self,
205 + common_prefix_len: int,
206 + common_attn_metadata: CommonAttentionMetadata,
207 + fast_build: bool = False) -> DeepseekV32IndexerMetadata:
208 +
209 + num_reqs = common_attn_metadata.num_reqs
210 + num_tokens = common_attn_metadata.num_actual_tokens
211 +
212 + device = self.device
213 + block_table_tensor = common_attn_metadata.block_table_tensor
214 +
215 + query_start_loc = common_attn_metadata.query_start_loc
216 +
217 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = \
218 + split_decodes_and_prefills(
219 + common_attn_metadata,
220 + decode_threshold=self.reorder_batch_threshold)
221 +
222 + assert num_decodes + num_prefills == num_reqs
223 + assert num_decode_tokens + num_prefill_tokens == num_tokens
224 +
225 + prefill_metadata = None
226 + if num_prefills > 0:
227 + reqs_start = num_decodes
228 + prefill_query_start_loc = query_start_loc[
229 + reqs_start:] - query_start_loc[reqs_start]
230 + cu_seqlen_ks, cu_seqlen_ke = kv_spans_from_batches(
231 + prefill_query_start_loc,
232 + common_attn_metadata.seq_lens[reqs_start:])
233 + total_seq_lens = common_attn_metadata.seq_lens[reqs_start:].sum()
234 + assert total_seq_lens < self.max_prefill_buffer_size
235 + cu_seq_lens = torch.cat([
236 + torch.zeros(1, dtype=torch.int32, device=device),
237 + common_attn_metadata.seq_lens[reqs_start:].cumsum(dim=0)
238 + ]).to(torch.int32).cuda()
239 + prefill_metadata = DeepseekV32IndexerPrefillMetadata(
240 + block_table=block_table_tensor[reqs_start:, ...],
241 + query_start_loc=prefill_query_start_loc,
242 + max_query_len=common_attn_metadata.max_query_len,
243 + cu_seqlen_ks=cu_seqlen_ks,
244 + cu_seqlen_ke=cu_seqlen_ke,
245 + cu_seq_lens=cu_seq_lens,
246 + total_seq_lens=total_seq_lens,
247 + )
248 +
249 + decode_metadata = None
250 + if num_decodes > 0:
251 + torch.diff(common_attn_metadata.query_start_loc[:num_decodes + 1],
252 + out=self.decode_lens_buffer[:num_decodes])
253 + decode_lens = self.decode_lens_buffer[:num_decodes]
254 + decode_lens_cpu = torch.diff(
255 + common_attn_metadata.query_start_loc_cpu[:num_decodes + 1])
256 +
257 + # Use CPU to avoid GPU sync; breaking async scheduling
258 + requires_padding = (decode_lens_cpu.max()
259 + > decode_lens_cpu.min()).item()
260 +
261 + seq_lens = common_attn_metadata.seq_lens[:num_decodes]
262 +
263 + self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(
264 + seq_lens, self.kv_cache_spec.block_size, self.num_sms)
265 + decode_metadata = DeepSeekV32IndexerDecodeMetadata(
266 + block_table=common_attn_metadata.
267 + block_table_tensor[:num_decodes, ...],
268 + seq_lens=common_attn_metadata.seq_lens[:num_decodes],
269 + decode_lens=decode_lens,
270 + requires_padding=requires_padding,
271 + schedule_metadata=self.scheduler_metadata_buffer,
272 + )
273 +
274 + attn_metadata = DeepseekV32IndexerMetadata(
275 + seq_lens=common_attn_metadata.seq_lens,
276 + num_reqs=common_attn_metadata.num_reqs,
277 + max_query_len=common_attn_metadata.max_query_len,
278 + max_seq_len=common_attn_metadata.max_seq_len,
279 + num_actual_tokens=common_attn_metadata.num_actual_tokens,
280 + query_start_loc=common_attn_metadata.query_start_loc,
281 + slot_mapping=common_attn_metadata.slot_mapping,
282 + head_dim=128,
283 + num_decodes=num_decodes,
284 + num_decode_tokens=num_decode_tokens,
285 + num_prefills=num_prefills,
286 + num_prefill_tokens=num_prefill_tokens,
287 + prefill=prefill_metadata,
288 + decode=decode_metadata,
289 + )
290 +
291 + # if get_tensor_model_parallel_rank() == 0:
292 + # logger.info(f"attn_metadata: {attn_metadata}")
293 + return attn_metadata
vllm/v1/attention/backends/pallas.py
+1
@@ -102,6 +102,7 @@ class PallasAttentionBackend(AttentionBackend):
102 block_size: int,
103 num_kv_heads: int,
104 head_size: int,
105 + cache_dtype_str: str = "auto",
106 ) -> tuple[int, ...]:
107 padded_head_size = cdiv(
108 head_size, TPU_HEAD_SIZE_ALIGNMENT) * TPU_HEAD_SIZE_ALIGNMENT
vllm/v1/attention/backends/rocm_aiter_fa.py
+1
@@ -360,6 +360,7 @@ class AiterFlashAttentionBackend(AttentionBackend):
360 block_size: int,
361 num_kv_heads: int,
362 head_size: int,
363 + cache_dtype_str: str = "auto",
364 ) -> tuple[int, ...]:
365 if block_size % 16 != 0:
366 raise ValueError("Block size must be a multiple of 16.")
vllm/v1/attention/backends/tree_attn.py
+1
@@ -68,6 +68,7 @@ class TreeAttentionBackend(AttentionBackend):
68 block_size: int,
69 num_kv_heads: int,
70 head_size: int,
71 + cache_dtype_str: str = "auto",
72 ) -> tuple[int, ...]:
73 if block_size % 16 != 0:
74 raise ValueError("Block size must be a multiple of 16.")
vllm/v1/attention/backends/triton_attn.py
+1
@@ -171,6 +171,7 @@ class TritonAttentionBackend(AttentionBackend):
171 block_size: int,
172 num_kv_heads: int,
173 head_size: int,
174 + cache_dtype_str: str = "auto",
175 ) -> tuple[int, ...]:
176 if block_size % 16 != 0:
177 raise ValueError("Block size must be a multiple of 16.")
vllm/v1/attention/backends/xformers.py
+1
@@ -106,6 +106,7 @@ class XFormersAttentionBackend(AttentionBackend):
106 block_size: int,
107 num_kv_heads: int,
108 head_size: int,
109 + cache_dtype_str: str = "auto",
110 ) -> tuple[int, ...]:
111 if block_size % 16 != 0:
112 raise ValueError("Block size must be a multiple of 16.")
vllm/v1/core/kv_cache_utils.py
+5 -4
@@ -1103,7 +1103,9 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
1103 kv_cache_spec: The kv cache spec of each attention layer in the model
1104 """
1105
1106 - if is_kv_cache_spec_uniform(kv_cache_spec):
1106 + if is_kv_cache_spec_uniform(
1107 + kv_cache_spec) or UniformTypeKVCacheSpecs.is_uniform_type(
1108 + kv_cache_spec):
1109 return
1110
1111 logger.warning(
@@ -1128,7 +1130,6 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
1130 num_kv_heads=spec.num_kv_heads,
1131 head_size=spec.head_size,
1132 dtype=spec.dtype,
1131 - use_mla=spec.use_mla,
1133 sliding_window=spec.sliding_window,
1134 )
1135 elif isinstance(spec, ChunkedLocalAttentionSpec):
@@ -1137,11 +1138,11 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
1138 num_kv_heads=spec.num_kv_heads,
1139 head_size=spec.head_size,
1140 dtype=spec.dtype,
1140 - use_mla=spec.use_mla,
1141 attention_chunk_size=spec.attention_chunk_size,
1142 )
1143
1144 - if not is_kv_cache_spec_uniform(kv_cache_spec):
1144 + if not (is_kv_cache_spec_uniform(kv_cache_spec)
1145 + or UniformTypeKVCacheSpecs.is_uniform_type(kv_cache_spec)):
1146 raise ValueError("Hybrid KV cache manager is disabled but failed to "
1147 "convert the KV cache specs to one unified type.")
1148
vllm/v1/core/single_type_kv_cache_manager.py
+2 -1
@@ -10,7 +10,7 @@ from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock
10 from vllm.v1.kv_cache_interface import (ChunkedLocalAttentionSpec,
11 CrossAttentionSpec, FullAttentionSpec,
12 KVCacheSpec, MambaSpec,
13 - SlidingWindowSpec)
13 + MLAAttentionSpec, SlidingWindowSpec)
14 from vllm.v1.request import Request
15
16
@@ -656,6 +656,7 @@ class CrossAttentionManager(SingleTypeKVCacheManager):
656
657 spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {
658 FullAttentionSpec: FullAttentionManager,
659 + MLAAttentionSpec: FullAttentionManager,
660 SlidingWindowSpec: SlidingWindowManager,
661 ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager,
662 MambaSpec: MambaManager,
vllm/v1/kv_cache_interface.py
+41 -10
@@ -59,13 +59,10 @@ class AttentionSpec(KVCacheSpec):
59 num_kv_heads: int
60 head_size: int
61 dtype: torch.dtype
62 - use_mla: bool
62
63 @property
64 def page_size_bytes(self) -> int:
66 - # For MLA we only store a single latent vector
67 - coef = 1 if self.use_mla else 2
68 - return coef * self.block_size * self.num_kv_heads * self.head_size \
65 + return 2 * self.block_size * self.num_kv_heads * self.head_size \
66 * get_dtype_size(self.dtype)
67
68
@@ -118,12 +115,13 @@ class FullAttentionSpec(AttentionSpec):
115 if spec.sliding_window is not None)
116 attention_chunk_size = set(spec.attention_chunk_size for spec in specs
117 if spec.attention_chunk_size is not None)
118 + assert not any(isinstance(spec, MLAAttentionSpec) for spec in specs), (
119 + "MLAAttentionSpec should be merged in MLAAttentionSpec.merge")
120 merged_spec = cls(
121 block_size=specs[0].block_size,
122 num_kv_heads=specs[0].num_kv_heads,
123 head_size=specs[0].head_size,
124 dtype=specs[0].dtype,
126 - use_mla=specs[0].use_mla,
125 sliding_window=cls.merge_window_sizes(sliding_window),
126 attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
127 )
@@ -140,6 +138,38 @@ class FullAttentionSpec(AttentionSpec):
138 return merged_spec
139
140
141 +@dataclass(frozen=True)
142 +class MLAAttentionSpec(FullAttentionSpec):
143 + # TODO(Lucas/Chen): less hacky way to do this
144 + cache_dtype_str: Optional[str] = None
145 +
146 + @property
147 + def page_size_bytes(self) -> int:
148 + if self.cache_dtype_str == "fp8_ds_mla":
149 + # See `vllm/v1/attention/backends/mla/flashmla_sparse.py`
150 + # for details.
151 + return self.block_size * 656
152 + return self.block_size * self.num_kv_heads * self.head_size \
153 + * get_dtype_size(self.dtype)
154 +
155 + @classmethod
156 + def merge(cls, specs: list[Self]) -> Self:
157 + assert all(isinstance(spec, MLAAttentionSpec) for spec in specs), (
158 + "All attention layers in the same KV cache group must be "
159 + "MLAAttentionSpec.")
160 + cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs)
161 + assert len(cache_dtype_str_set) == 1, (
162 + "All attention layers in the same KV cache group must use the same "
163 + "quantization method.")
164 + return cls(
165 + block_size=specs[0].block_size,
166 + num_kv_heads=specs[0].num_kv_heads,
167 + head_size=specs[0].head_size,
168 + dtype=specs[0].dtype,
169 + cache_dtype_str=cache_dtype_str_set.pop(),
170 + )
171 +
172 +
173 @dataclass(frozen=True)
174 class ChunkedLocalAttentionSpec(AttentionSpec):
175 attention_chunk_size: int
@@ -163,9 +193,6 @@ class ChunkedLocalAttentionSpec(AttentionSpec):
193 class SlidingWindowSpec(AttentionSpec):
194 sliding_window: int
195
166 - def __post_init__(self):
167 - assert not self.use_mla, "MLA is not supported for sliding window"
168 -
196 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
197 assert vllm_config.parallel_config.decode_context_parallel_size == 1, \
198 "DCP not support sliding window."
@@ -266,9 +293,13 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
293 # Different block sizes, not uniform.
294 return False
295 one_spec = next(iter(kv_cache_specs.values()))
269 - if isinstance(one_spec, (FullAttentionSpec, CrossAttentionSpec)):
296 + if isinstance(one_spec, FullAttentionSpec):
297 + return all(
298 + isinstance(spec, FullAttentionSpec)
299 + for spec in kv_cache_specs.values())
300 + elif isinstance(one_spec, CrossAttentionSpec):
301 return all(
271 - isinstance(spec, type(one_spec))
302 + isinstance(spec, CrossAttentionSpec)
303 for spec in kv_cache_specs.values())
304 elif isinstance(one_spec, SlidingWindowSpec):
305 return all(
vllm/v1/spec_decode/eagle.py
+45 -11
@@ -17,6 +17,7 @@ from vllm.forward_context import set_forward_context
17 from vllm.logger import init_logger
18 from vllm.model_executor.model_loader import get_model
19 from vllm.model_executor.models import supports_multimodal
20 +from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache
21 from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
22 from vllm.multimodal import MULTIMODAL_REGISTRY
23 from vllm.platforms import current_platform
@@ -32,6 +33,7 @@ from vllm.v1.sample.metadata import SamplingMetadata
33 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
34 from vllm.v1.utils import CpuGpuBuffer
35 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
36 +from vllm.v1.worker.ubatching import dbo_current_ubatch_id
37
38 logger = init_logger(__name__)
39
@@ -52,6 +54,7 @@ class EagleProposer:
54 self.method = self.speculative_config.method
55
56 self.runner = runner
57 + self.device = device
58 self.dtype = vllm_config.model_config.dtype
59 self.max_model_len = vllm_config.model_config.max_model_len
60 self.block_size = vllm_config.cache_config.block_size
@@ -202,20 +205,30 @@ class EagleProposer:
205
206 assert self.runner is not None
207
205 - # Select the correct attention metadata builders for EAGLE layers.
206 - # Get the attention metadata builders once and reuse for later.
207 - builder = (self._get_attention_metadata_builder()
208 - if self.attn_metadata_builder is None else
209 - self.attn_metadata_builder)
210 - attn_metadata = builder.build_for_drafting( # type: ignore
211 - common_attn_metadata=common_attn_metadata,
212 - draft_index=0)
213 -
208 + # FIXME: need to consider multiple kv_cache_groups
209 + ubatch_id = dbo_current_ubatch_id()
210 + attn_metadata_builder = \
211 + self.runner.attn_groups[0][0].metadata_builders[ubatch_id]
212 + attn_metadata = attn_metadata_builder.build_for_drafting(
213 + common_attn_metadata=common_attn_metadata, draft_index=0)
214 + # FIXME: support hybrid kv for draft model (remove separate indexer)
215 + if self.draft_indexer_metadata_builder:
216 + draft_indexer_metadata = (
217 + self.draft_indexer_metadata_builder.build_for_drafting(
218 + common_attn_metadata=common_attn_metadata,
219 + draft_index=0,
220 + ))
221 + else:
222 + draft_indexer_metadata = None
223 # At this moment, we assume all eagle layers belong to the same KV
224 # cache group, thus using the same attention metadata.
225 per_layer_attn_metadata = {}
226 for layer_name in self.attn_layer_names:
227 per_layer_attn_metadata[layer_name] = attn_metadata
228 + for layer_name in self.indexer_layer_names:
229 + assert draft_indexer_metadata is not None
230 + per_layer_attn_metadata[layer_name] = draft_indexer_metadata
231 +
232 if self.use_cuda_graph and \
233 num_tokens <= self.cudagraph_batch_sizes[-1]:
234 num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
@@ -370,7 +383,7 @@ class EagleProposer:
383 exceeds_max_model_len, PADDING_SLOT_ID)
384
385 # Rebuild attention metadata
373 - attn_metadata = builder.build_for_drafting( # type: ignore
386 + attn_metadata = attn_metadata_builder.build_for_drafting( # type: ignore
387 common_attn_metadata=common_attn_metadata,
388 draft_index=token_index + 1)
389 for layer_name in self.attn_layer_names:
@@ -846,6 +859,10 @@ class EagleProposer:
859 self.vllm_config.speculative_config.draft_model_config
860 target_attn_layer_names = set(
861 get_layers_from_vllm_config(self.vllm_config, Attention).keys())
862 + # FIXME: support hybrid kv for draft model
863 + target_indexer_layer_names = set(
864 + get_layers_from_vllm_config(self.vllm_config,
865 + DeepseekV32IndexerCache).keys())
866
867 from vllm.compilation.backends import set_model_tag
868 with set_model_tag("eagle_head"):
@@ -855,8 +872,25 @@ class EagleProposer:
872 draft_attn_layer_names = (
873 get_layers_from_vllm_config(self.vllm_config, Attention).keys() -
874 target_attn_layer_names)
858 -
875 + indexer_layers = get_layers_from_vllm_config(self.vllm_config,
876 + DeepseekV32IndexerCache)
877 + draft_indexer_layer_names = (indexer_layers.keys() -
878 + target_indexer_layer_names)
879 self.attn_layer_names = list(draft_attn_layer_names)
880 + self.indexer_layer_names = list(draft_indexer_layer_names)
881 +
882 + if self.indexer_layer_names:
883 + first_layer = self.indexer_layer_names[0]
884 + self.draft_indexer_metadata_builder = (
885 + indexer_layers[first_layer].get_attn_backend().get_builder_cls(
886 + )(
887 + indexer_layers[first_layer].get_kv_cache_spec(),
888 + self.indexer_layer_names,
889 + self.vllm_config,
890 + self.device,
891 + ))
892 + else:
893 + self.draft_indexer_metadata_builder = None
894
895 if self.supports_mm_inputs:
896 # Even if the target model is multimodal, we can also use
vllm/v1/worker/gpu_model_runner.py
+29 -16
@@ -40,6 +40,7 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
40 from vllm.model_executor.layers.mamba.abstract import MambaBase
41 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
42 from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader
43 +from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache
44 # yapf conflicts with isort for this block
45 # yapf: disable
46 from vllm.model_executor.models.interfaces import (SupportsMultiModal,
@@ -80,7 +81,8 @@ from vllm.v1.kv_cache_interface import (AttentionSpec,
81 EncoderOnlyAttentionSpec,
82 FullAttentionSpec, KVCacheConfig,
83 KVCacheGroupSpec, KVCacheSpec,
83 - MambaSpec, SlidingWindowSpec,
84 + MambaSpec, MLAAttentionSpec,
85 + SlidingWindowSpec,
86 UniformTypeKVCacheSpecs)
87 # yapf: enable
88 from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput,
@@ -3154,7 +3156,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3156 attn_metadata_i = (attn_group\
3157 .get_metadata_builder(ubatch_id=ubid)\
3158 .build_for_cudagraph_capture(common_attn_metadata))
3157 - for layer_name in kv_cache_group_spec.layer_names:
3159 + for layer_name in attn_group.layer_names:
3160 assert type(attn_metadata) is list
3161 attn_metadata[ubid][
3162 layer_name] = attn_metadata_i
@@ -3162,7 +3164,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3164 assert type(attn_metadata) is dict
3165 attn_metadata_i = attn_group.get_metadata_builder()\
3166 .build_for_cudagraph_capture(common_attn_metadata)
3165 - for layer_name in kv_cache_group_spec.layer_names:
3167 + for layer_name in attn_group.layer_names:
3168 attn_metadata[layer_name] = attn_metadata_i
3169
3170 with self.maybe_dummy_run_with_lora(self.lora_config,
@@ -3915,8 +3917,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3917 if isinstance(kv_cache_spec, AttentionSpec):
3918 has_attn = True
3919 kv_cache_shape = attn_backend.get_kv_cache_shape(
3918 - num_blocks, kv_cache_spec.block_size,
3919 - kv_cache_spec.num_kv_heads, kv_cache_spec.head_size)
3920 + num_blocks,
3921 + kv_cache_spec.block_size,
3922 + kv_cache_spec.num_kv_heads,
3923 + kv_cache_spec.head_size,
3924 + cache_dtype_str=self.cache_config.cache_dtype)
3925 dtype = kv_cache_spec.dtype
3926 try:
3927 kv_cache_stride_order = \
@@ -4101,7 +4106,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4106 Add encoder-only layers to the KV cache config.
4107 """
4108 block_size = self.vllm_config.cache_config.block_size
4104 - use_mla = self.vllm_config.model_config.use_mla
4109 encoder_only_attn_specs: dict[AttentionSpec,
4110 list[str]] = defaultdict(list)
4111 attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention)
@@ -4111,8 +4115,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4115 block_size=block_size,
4116 num_kv_heads=attn_module.num_kv_heads,
4117 head_size=attn_module.head_size,
4114 - dtype=self.kv_cache_dtype,
4115 - use_mla=use_mla)
4118 + dtype=self.kv_cache_dtype)
4119 encoder_only_attn_specs[attn_spec].append(layer_name)
4120 self.runner_only_attn_layers.add(layer_name)
4121 if len(encoder_only_attn_specs) > 0:
@@ -4134,6 +4137,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4137
4138 block_size = self.vllm_config.cache_config.block_size
4139 use_mla = self.vllm_config.model_config.use_mla
4140 + cache_dtype_str = self.vllm_config.cache_config.cache_dtype
4141 kv_cache_spec: dict[str, KVCacheSpec] = {}
4142 attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention)
4143 for layer_name, attn_module in attn_layers.items():
@@ -4153,13 +4157,21 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4157 # the attention backends
4158 if attn_module.attn_type == AttentionType.DECODER:
4159 if attn_module.sliding_window is not None:
4160 + assert not use_mla, "MLA is not supported for sliding" \
4161 + "window"
4162 kv_cache_spec[layer_name] = SlidingWindowSpec(
4163 block_size=block_size,
4164 num_kv_heads=attn_module.num_kv_heads,
4165 head_size=attn_module.head_size,
4166 dtype=self.kv_cache_dtype,
4161 - sliding_window=attn_module.sliding_window,
4162 - use_mla=use_mla)
4167 + sliding_window=attn_module.sliding_window)
4168 + elif use_mla:
4169 + kv_cache_spec[layer_name] = MLAAttentionSpec(
4170 + block_size=block_size,
4171 + num_kv_heads=attn_module.num_kv_heads,
4172 + head_size=attn_module.head_size,
4173 + dtype=self.kv_cache_dtype,
4174 + cache_dtype_str=cache_dtype_str)
4175 elif self.attention_chunk_size is not None \
4176 and isinstance(attn_module, ChunkedLocalAttention):
4177 kv_cache_spec[layer_name] = ChunkedLocalAttentionSpec(
@@ -4167,22 +4179,19 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4179 num_kv_heads=attn_module.num_kv_heads,
4180 head_size=attn_module.head_size,
4181 dtype=self.kv_cache_dtype,
4170 - attention_chunk_size=self.attention_chunk_size,
4171 - use_mla=use_mla)
4182 + attention_chunk_size=self.attention_chunk_size)
4183 else:
4184 kv_cache_spec[layer_name] = FullAttentionSpec(
4185 block_size=block_size,
4186 num_kv_heads=attn_module.num_kv_heads,
4187 head_size=attn_module.head_size,
4177 - dtype=self.kv_cache_dtype,
4178 - use_mla=use_mla)
4188 + dtype=self.kv_cache_dtype)
4189 elif attn_module.attn_type == AttentionType.ENCODER_DECODER:
4190 kv_cache_spec[layer_name] = CrossAttentionSpec(
4191 block_size=block_size,
4192 num_kv_heads=attn_module.num_kv_heads,
4193 head_size=attn_module.head_size,
4184 - dtype=self.kv_cache_dtype,
4185 - use_mla=use_mla)
4194 + dtype=self.kv_cache_dtype)
4195 elif attn_module.attn_type in (AttentionType.ENCODER,
4196 AttentionType.ENCODER_ONLY):
4197 # encoder-only attention does not need KV cache.
@@ -4219,6 +4228,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4228 self.speculative_config.num_speculative_tokens
4229 if self.speculative_config else 0),
4230 )
4231 + ds_indexer_layers = get_layers_from_vllm_config(
4232 + self.vllm_config, DeepseekV32IndexerCache)
4233 + for layer_name, ds_indexer_module in ds_indexer_layers.items():
4234 + kv_cache_spec[layer_name] = ds_indexer_module.get_kv_cache_spec()
4235
4236 return kv_cache_spec
4237
vllm/v1/worker/tpu_model_runner.py
-2
@@ -537,7 +537,6 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
537 head_size=attn_module.head_size,
538 dtype=self.kv_cache_dtype,
539 sliding_window=attn_module.sliding_window,
540 - use_mla=False,
540 )
541 else:
542 kv_cache_spec[layer_name] = FullAttentionSpec(
@@ -545,7 +544,6 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
544 num_kv_heads=attn_module.num_kv_heads,
545 head_size=attn_module.head_size,
546 dtype=self.kv_cache_dtype,
548 - use_mla=False,
547 )
548 elif attn_module.attn_type in (AttentionType.ENCODER,
549 AttentionType.ENCODER_ONLY):