@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 77e10c9cab

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 12, 14:45 UTC

[Perf][Deepseek] optimize gather_and_maybe_dequant_cache kernel's perf for extremely long sequence (#28029)

Signed-off-by: ganyi <ygan@amd.com>

Pleaplusone committed Nov 25, 2025 at 10:05 UTC 77e10c9cab751c83de0b2200977212922cc3776f

6 files changed +132 -106

csrc/cache.h

+6 -5

@@ -41,11 +41,12 @@ void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,

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

43 void gather_and_maybe_dequant_cache(

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

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

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

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

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

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

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

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

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

48 + torch::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS]

49 + int64_t num_tokens, const std::string& kv_cache_dtype,

50 torch::Tensor const& scale,

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

csrc/cache_kernels.cu

+86 -92

@@ -905,91 +905,79 @@ void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,

905 namespace vllm {

906

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

908 -template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>

908 +template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,

909 + int ENTRY_SIZE, int CTA_SIZE>

910 __global__ void gather_and_maybe_dequant_cache(

910 - const cache_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE,

911 - // ENTRIES...]

912 - scalar_t* __restrict__ dst, // [TOT_TOKENS, ENTRIES...]

913 - const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES]

914 - const int32_t* __restrict__ cu_seq_lens, // [BATCH+1]

915 - const int32_t block_size, const int32_t entry_size,

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

912 + // ENTRIES...]

913 + scalar_t* __restrict__ dst, // [TOT_TOKENS, ENTRIES...]

914 + const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES]

915 + const int32_t* __restrict__ cu_seq_lens, // [BATCH+1]

916 + const int32_t* __restrict__ token_to_seq, // [MAX_TOKEN_ACROSS_CHUNK]

917 + const int32_t num_tokens, const int32_t block_size,

918 const int64_t block_table_stride, const int64_t cache_block_stride,

919 const int64_t cache_entry_stride, const int64_t dst_entry_stride,

920 const float* __restrict__ scale,

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

922 // batch

923 + constexpr int vec_size = sizeof(float4) / sizeof(scalar_t);

924 + using ltype = vllm::vec_n_t<cache_t, vec_size>;

925 + using stype = vllm::vec_n_t<scalar_t, vec_size>;

926 + // We are adding this for code readability which will be optimized out when

927 + // build in release.

928 + assert(CTA_SIZE == blockDim.x);

929

922 - const int64_t bid = blockIdx.x; // Batch ID

923 - const int32_t num_splits = gridDim.y;

924 - const int32_t split = blockIdx.y;

925 - const int32_t seq_start = cu_seq_lens[bid];

926 - const int32_t seq_end = cu_seq_lens[bid + 1];

927 - const int32_t seq_len = seq_end - seq_start;

928 - const int32_t tot_blocks = cuda_utils::ceil_div(seq_len, block_size);

929 - const int32_t split_blocks = cuda_utils::ceil_div(tot_blocks, num_splits);

930 -

931 - const int32_t split_start = split * split_blocks;

932 - const int32_t split_end = min((split + 1) * split_blocks, tot_blocks);

933 -

934 - const bool is_active_split = (split_start < tot_blocks);

935 - const bool is_last_split = (split_end == tot_blocks);

936 -

937 - if (!is_active_split) return;

938 -

939 - int32_t full_blocks_end = split_end;

940 - int32_t partial_block_size = 0;

941 -

942 - // Adjust the pointer for the block_table for this batch.

943 - // If seq_starts is provided, compute an offset based on (seq_starts[bid] /

944 - // page_size)

945 - const int32_t batch_offset = bid * block_table_stride;

946 - int32_t offset = 0;

947 - if (seq_starts != nullptr) {

948 - offset = seq_starts[bid] / block_size;

949 - }

950 - const int32_t* batch_block_table = block_table + batch_offset + offset;

951 -

952 - // Adjust dst pointer based on the cumulative sequence lengths.

953 - dst += seq_start * dst_entry_stride;

954 -

955 - if (is_last_split) {

956 - partial_block_size = seq_len % block_size;

957 - if (partial_block_size) full_blocks_end -= 1;

958 - }

930 +#pragma unroll

931 + for (int token_id = blockIdx.x; token_id < num_tokens;

932 + token_id += gridDim.x) {

933 + int64_t batch_id = token_to_seq[token_id];

934 + int64_t batch_start = cu_seq_lens[batch_id];

935 + int64_t batch_end = cu_seq_lens[batch_id + 1];

936 + int32_t batch_offset = token_id - batch_start;

937 +

938 + if (token_id >= batch_end) return;

939 + int32_t offset = 0;

940 + if (seq_starts != nullptr) {

941 + offset = seq_starts[batch_id];

942 + }

943 + batch_offset += offset;

944 + int32_t block_table_id = batch_offset / block_size;

945 + int32_t slot_id = batch_offset % block_size;

946 + int32_t block_table_offset = batch_id * block_table_stride + block_table_id;

947 + int32_t block_id = block_table[block_table_offset];

948 + int64_t cache_offset =

949 + block_id * cache_block_stride + slot_id * cache_entry_stride;

950 + constexpr int32_t vec_iter_cnt = ENTRY_SIZE / vec_size;

951 + scalar_t* dst_ = dst + token_id * dst_entry_stride;

952 + cache_t* src_ = const_cast<cache_t*>(src_cache) + cache_offset;

953

960 - auto copy_entry = [&](const cache_t* __restrict__ _src,

961 - scalar_t* __restrict__ _dst) {

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

954 +#pragma unroll

955 + for (int idx = threadIdx.x; idx < vec_iter_cnt; idx += CTA_SIZE) {

956 if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {

964 - _dst[i] = static_cast<scalar_t>(_src[i]);

957 + reinterpret_cast<stype*>(dst_)[idx] =

958 + static_cast<stype>(reinterpret_cast<ltype*>(src_)[idx]);

959 } else {

966 - _dst[i] =

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

960 + ltype loaded_val = reinterpret_cast<ltype*>(src_)[idx];

961 + stype store_val;

962 +#pragma unroll

963 + for (int j = 0; j < vec_size; ++j) {

964 + store_val.val[j] = fp8::scaled_convert<scalar_t, cache_t, kv_dt>(

965 + loaded_val.val[j], *scale);

966 + }

967 + reinterpret_cast<stype*>(dst_)[idx] = store_val;

968 }

969 }

970 - };

971 -

972 - const auto loop_end =

973 - std::min((int64_t)full_blocks_end, block_table_stride - offset);

974 - for (int pid = split_start; pid < loop_end; ++pid) {

975 - auto block_id = batch_block_table[pid];

976 - auto block_start_ptr = src_cache + block_id * cache_block_stride;

977 - auto block_dst_ptr = dst + pid * block_size * dst_entry_stride;

978 - for (int eid = 0; eid < block_size; ++eid) {

979 - copy_entry(block_start_ptr + eid * cache_entry_stride,

980 - block_dst_ptr + eid * dst_entry_stride);

981 - }

982 - }

983 -

984 - if (partial_block_size) {

985 - if (offset + full_blocks_end < block_table_stride) {

986 - auto block_id = batch_block_table[full_blocks_end];

987 - auto block_start_ptr = src_cache + block_id * cache_block_stride;

988 - auto block_dst_ptr =

989 - dst + full_blocks_end * block_size * dst_entry_stride;

990 - for (int eid = 0; eid < partial_block_size; ++eid) {

991 - copy_entry(block_start_ptr + eid * cache_entry_stride,

992 - block_dst_ptr + eid * dst_entry_stride);

970 + // process tail

971 + constexpr int32_t tail_cnt = ENTRY_SIZE % vec_size;

972 + dst_ = dst_ + ENTRY_SIZE - tail_cnt;

973 + src_ = src_ + ENTRY_SIZE - tail_cnt;

974 +#pragma unroll

975 + for (int idx = threadIdx.x; idx < tail_cnt; idx += CTA_SIZE) {

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

977 + dst_[idx] = static_cast<scalar_t>(src_[idx]);

978 + } else {

979 + dst_[idx] =

980 + fp8::scaled_convert<scalar_t, cache_t, kv_dt>(src_[idx], *scale);

981 }

982 }

983 }

@@ -1001,34 +989,38 @@ __global__ void gather_and_maybe_dequant_cache(

989 // SCALAR_T is the data type of the destination tensor.

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

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

1004 -#define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE) \

1005 - vllm::gather_and_maybe_dequant_cache<SCALAR_T, CACHE_T, KV_DTYPE> \

1006 - <<<grid, block, 0, stream>>>( \

1007 - reinterpret_cast<CACHE_T*>(src_cache.data_ptr()), \

1008 - reinterpret_cast<SCALAR_T*>(dst.data_ptr()), \

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

1010 - block_size, entry_size, block_table_stride, cache_block_stride, \

1011 - cache_entry_stride, dst_entry_stride, \

1012 - reinterpret_cast<const float*>(scale.data_ptr()), seq_starts_ptr);

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

993 + vllm::gather_and_maybe_dequant_cache<SCALAR_T, CACHE_T, KV_DTYPE, 576, \

994 + thread_block_size> \

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

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

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

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

999 + token_to_seq.data_ptr<int32_t>(), num_tokens, block_size, \

1000 + block_table_stride, cache_block_stride, cache_entry_stride, \

1001 + dst_entry_stride, reinterpret_cast<const float*>(scale.data_ptr()), \

1002 + seq_starts_ptr);

1003

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

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

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

1007 +// - token_to_seq contains the back mapping from token_id to batch_id

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

1009 // (seq_starts[bid] / page_size)

1010 void gather_and_maybe_dequant_cache(

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

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

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

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

1024 - int64_t batch_size, const std::string& kv_cache_dtype,

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

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

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

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

1015 + torch::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS]

1016 + int64_t num_tokens, const std::string& kv_cache_dtype,

1017 torch::Tensor const& scale,

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

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

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

1021

1022 int32_t block_size = src_cache.size(1);

1031 - int32_t entry_size = src_cache.flatten(2, -1).size(2);

1023 + int32_t head_dim = dst.size(-1);

1024

1025 TORCH_CHECK(block_table.dtype() == torch::kInt32,

1026 "block_table must be int32");

@@ -1038,6 +1030,9 @@ void gather_and_maybe_dequant_cache(

1030 TORCH_CHECK(seq_starts.value().dtype() == torch::kInt32,

1031 "seq_starts must be int32");

1032 }

1033 + TORCH_CHECK(head_dim == 576,

1034 + "gather_and_maybe_dequant_cache only support the head_dim to 576 "

1035 + "for better performance")

1036

1037 TORCH_CHECK(src_cache.device() == dst.device(),

1038 "src_cache and dst must be on the same device");

@@ -1055,10 +1050,9 @@ void gather_and_maybe_dequant_cache(

1050 int64_t cache_entry_stride = src_cache.stride(1);

1051 int64_t dst_entry_stride = dst.stride(0);

1052

1058 - // Decide on the number of splits based on the batch size.

1059 - int num_splits = batch_size > 128 ? 2 : batch_size > 64 ? 4 : 16;

1060 - dim3 grid(batch_size, num_splits);

1061 - dim3 block(1024);

1053 + constexpr int32_t thread_block_size = 64;

1054 + dim3 grid(num_tokens);

1055 + dim3 block(thread_block_size);

1056

1057 const int32_t* seq_starts_ptr =

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

csrc/torch_bindings.cpp

+2 -1

@@ -695,7 +695,8 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {

695 cache_ops.def(

696 "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, "

697 " Tensor block_table, Tensor cu_seq_lens, "

698 - " int batch_size, "

698 + " Tensor token_to_seq, "

699 + " int num_tokens, "

700 " str kv_cache_dtype, "

701 " Tensor scale, Tensor? seq_starts) -> ()");

702 cache_ops.impl("gather_and_maybe_dequant_cache", torch::kCUDA,

tests/kernels/attention/test_cache.py

+9 -3

@@ -921,12 +921,16 @@ def test_gather_and_maybe_dequant_cache_mla(

921 )

922 _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype)

923

924 - seq_len_tensor = torch.randint(0, max_seq_len + 1, (batch_size,), device=device)

924 + seq_len_tensor = torch.randint(

925 + max_seq_len, max_seq_len + 1, (batch_size,), device=device

926 + )

927

928 total_tokens = seq_len_tensor.sum()

929 cu_seq_lens = torch.empty((batch_size + 1), dtype=torch.int32, device=device)

930 cu_seq_lens[0] = 0

931 cu_seq_lens[1:] = seq_len_tensor.cumsum(dim=0).to(dtype=torch.int32)

932 + token_to_seq = torch.arange(0, batch_size, dtype=torch.int32, device=device)

933 + token_to_seq = torch.repeat_interleave(token_to_seq, seq_len_tensor)

934 print("seq_len_tensor", seq_len_tensor)

935

936 tot_blocks_tensor = (seq_len_tensor + block_size - 1) // block_size

@@ -977,7 +981,8 @@ def test_gather_and_maybe_dequant_cache_mla(

981 dst,

982 block_table,

983 cu_seq_lens,

980 - batch_size,

984 + token_to_seq,

985 + total_tokens,

986 kv_cache_dtype,

987 scale,

988 None,

@@ -990,7 +995,8 @@ def test_gather_and_maybe_dequant_cache_mla(

995 dst,

996 block_table,

997 cu_seq_lens,

993 - batch_size,

998 + token_to_seq,

999 + total_tokens,

1000 kv_cache_dtype,

1001 scale,

1002 None,

vllm/_custom_ops.py

+4 -2

@@ -2201,7 +2201,8 @@ def gather_and_maybe_dequant_cache(

2201 dst: torch.Tensor,

2202 block_table: torch.Tensor,

2203 cu_seq_lens: torch.Tensor,

2204 - batch_size: int,

2204 + token_to_seq: torch.Tensor,

2205 + num_tokens: int,

2206 kv_cache_dtype: str,

2207 scale: torch.Tensor,

2208 seq_starts: torch.Tensor | None = None,

@@ -2211,7 +2212,8 @@ def gather_and_maybe_dequant_cache(

2212 dst,

2213 block_table,

2214 cu_seq_lens,

2214 - batch_size,

2215 + token_to_seq,

2216 + num_tokens,

2217 kv_cache_dtype,

2218 scale,

2219 seq_starts,

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

+25 -3

@@ -340,6 +340,8 @@ class MLACommonPrefillMetadata:

340 max_seq_lens: list[int]

341 seq_lens: torch.Tensor

342 workspace: torch.Tensor

343 + token_to_seq: torch.Tensor

344 + chunk_total_token: list[int]

345

346 # for mla DCP

347 padded_local_chunk_seq_lens: list[list[int]] | None = None

@@ -839,6 +841,19 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

841 torch.cumsum(

842 chunk_seq_lens, dim=1, out=cu_seq_lens_cpu[:, 1:], dtype=torch.int32

843 )

844 + chunk_total_token = cu_seq_lens_cpu[:, -1]

845 +

846 + max_token_num_over_chunk = chunk_total_token.max().item()

847 + token_to_seq_tensor_cpu = torch.zeros(

848 + [num_chunks, max_token_num_over_chunk], dtype=torch.int32

849 + )

850 + range_idx = torch.arange(num_prefills, dtype=torch.int32)

851 + for i in range(num_chunks):

852 + chunk_token_to_seq_tensor = torch.repeat_interleave(

853 + range_idx, chunk_seq_lens[i]

854 + )

855 + chunk_len = chunk_token_to_seq_tensor.shape[0]

856 + token_to_seq_tensor_cpu[i, :chunk_len] = chunk_token_to_seq_tensor

857

858 if self.dcp_world_size > 1:

859 local_context_lens_allranks = get_dcp_local_seq_lens(

@@ -906,6 +921,10 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

921 seq_tot=padded_local_chunk_seq_lens.sum(dim=1).tolist(),

922 max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),

923 seq_lens=chunk_seq_lens,

924 + token_to_seq=token_to_seq_tensor_cpu.to(

925 + device, non_blocking=True

926 + ),

927 + chunk_total_token=chunk_total_token.tolist(),

928 workspace=self.chunked_prefill_workspace,

929 padded_local_chunk_seq_lens=padded_local_chunk_seq_lens.tolist(),

930 local_context_lens_allranks=local_context_lens_allranks.tolist(),

@@ -922,6 +941,10 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

941 seq_tot=chunk_seq_lens.sum(dim=1).tolist(),

942 max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),

943 seq_lens=chunk_seq_lens,

944 + token_to_seq=token_to_seq_tensor_cpu.to(

945 + device, non_blocking=True

946 + ),

947 + chunk_total_token=chunk_total_token,

948 workspace=self.chunked_prefill_workspace,

949 )

950

@@ -1638,16 +1661,15 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1661 output = None

1662 iters = len(prefill_metadata.chunked_context.seq_tot)

1663 workspace = prefill_metadata.chunked_context.workspace

1641 -

1664 for i in range(iters):

1665 toks = prefill_metadata.chunked_context.seq_tot[i]

1644 -

1666 ops.gather_and_maybe_dequant_cache(

1667 src_cache=kv_c_and_k_pe_cache,

1668 dst=workspace,

1669 block_table=prefill_metadata.block_table,

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

1650 - batch_size=attn_metadata.num_prefills,

1671 + token_to_seq=prefill_metadata.chunked_context.token_to_seq[i],

1672 + num_tokens=prefill_metadata.chunked_context.chunk_total_token[i],

1673 kv_cache_dtype=self.kv_cache_dtype,

1674 scale=k_scale,

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