@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c76a425278

[Kernel] Optimize long-context MLA cache gathers (#51739)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Yongye Zhu committed Aug 10, 2026 at 22:23 UTC c76a425278bc41e9d661a377d7c166f5be42cf6b
6 files changed +810 -235
.buildkite/test_areas/kernels.yaml
+1 -1
@@ -224,7 +224,7 @@ steps:
224 - label: Kernels Mamba Test
225 device: h200_35gb
226 key: kernels-mamba-test
227 - timeout_in_minutes: 40
227 + timeout_in_minutes: 60
228 source_file_dependencies:
229 - csrc/mamba/
230 - tests/kernels/mamba
benchmarks/kernels/benchmark_cp_gather.py new
+248
@@ -0,0 +1,248 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import math
5 +from collections.abc import Callable
6 +
7 +import torch
8 +
9 +from vllm import _custom_ops as ops
10 +from vllm.triton_utils import triton
11 +from vllm.utils.argparse_utils import FlexibleArgumentParser
12 +
13 +SCENARIOS = {
14 + "single-60k": [60_000],
15 + "single-300k": [300_000],
16 + "skew-2": [60_000, 300_000],
17 + "skew-4": [60_000, 100_000, 180_000, 300_000],
18 + "skew-8": [
19 + 60_000,
20 + 60_000,
21 + 80_000,
22 + 100_000,
23 + 140_000,
24 + 180_000,
25 + 240_000,
26 + 300_000,
27 + ],
28 +}
29 +DTYPES = {
30 + "fp8": torch.float8_e4m3fn,
31 + "bfloat16": torch.bfloat16,
32 + "float32": torch.float32,
33 +}
34 +
35 +
36 +def make_page_table(
37 + seq_lens: list[int],
38 + block_size: int,
39 + seq_starts: list[int] | None = None,
40 +) -> tuple[torch.Tensor, torch.Tensor, int]:
41 + if seq_starts is None:
42 + seq_starts = [0] * len(seq_lens)
43 + blocks_per_req = [
44 + math.ceil((start + length) / block_size)
45 + for start, length in zip(seq_starts, seq_lens)
46 + ]
47 + total_blocks = sum(blocks_per_req)
48 + block_table = torch.zeros(
49 + (len(seq_lens), max(blocks_per_req)), dtype=torch.int32, device="cuda"
50 + )
51 + physical_blocks = torch.randperm(total_blocks, dtype=torch.int32, device="cuda")
52 + block_offset = 0
53 + for req_id, num_blocks in enumerate(blocks_per_req):
54 + block_table[req_id, :num_blocks] = physical_blocks[
55 + block_offset : block_offset + num_blocks
56 + ]
57 + block_offset += num_blocks
58 +
59 + cu_seq_lens = torch.zeros(len(seq_lens) + 1, dtype=torch.int32, device="cuda")
60 + cu_seq_lens[1:] = torch.tensor(seq_lens, dtype=torch.int32, device="cuda").cumsum(
61 + dim=0
62 + )
63 + return block_table, cu_seq_lens, total_blocks
64 +
65 +
66 +def make_cache_gather(
67 + seq_lens: list[int],
68 + block_size: int,
69 + entry_size: int,
70 + dtype: torch.dtype,
71 +) -> tuple[Callable[[], None], int]:
72 + seq_starts_list = [
73 + 13 + (17 * req_id) % block_size for req_id in range(len(seq_lens))
74 + ]
75 + block_table, cu_seq_lens, total_blocks = make_page_table(
76 + seq_lens, block_size, seq_starts_list
77 + )
78 + src_cache = torch.empty(
79 + (total_blocks, block_size, entry_size), dtype=dtype, device="cuda"
80 + )
81 + dst = torch.empty((sum(seq_lens), entry_size), dtype=dtype, device="cuda")
82 + seq_starts = torch.tensor(seq_starts_list, dtype=torch.int32, device="cuda")
83 +
84 + def run() -> None:
85 + ops.cp_gather_cache(
86 + src_cache,
87 + dst,
88 + block_table,
89 + cu_seq_lens,
90 + len(seq_lens),
91 + seq_starts,
92 + )
93 +
94 + bytes_moved = 2 * dst.numel() * dst.element_size()
95 + return run, bytes_moved
96 +
97 +
98 +def make_fp8_upconvert(
99 + seq_lens: list[int],
100 + block_size: int,
101 +) -> tuple[Callable[[], None], int]:
102 + entry_bytes = 656
103 + output_elements = 576
104 + seq_starts_list = [
105 + 13 + (17 * req_id) % block_size for req_id in range(len(seq_lens))
106 + ]
107 + block_table, cu_seq_lens, total_blocks = make_page_table(
108 + seq_lens, block_size, seq_starts_list
109 + )
110 + src_cache = torch.empty(
111 + (total_blocks, block_size, entry_bytes),
112 + dtype=torch.uint8,
113 + device="cuda",
114 + )
115 + dst = torch.empty(
116 + (sum(seq_lens), output_elements), dtype=torch.bfloat16, device="cuda"
117 + )
118 + seq_starts = torch.tensor(seq_starts_list, dtype=torch.int32, device="cuda")
119 +
120 + def run() -> None:
121 + ops.cp_gather_and_upconvert_fp8_kv_cache(
122 + src_cache,
123 + dst,
124 + block_table,
125 + cu_seq_lens[:-1],
126 + len(seq_lens),
127 + seq_starts,
128 + )
129 +
130 + bytes_moved = sum(seq_lens) * (entry_bytes + output_elements * 2)
131 + return run, bytes_moved
132 +
133 +
134 +def make_maybe_dequant_gather(
135 + seq_lens: list[int],
136 + block_size: int,
137 + entry_size: int,
138 +) -> tuple[Callable[[], None], int]:
139 + seq_starts_list = [
140 + 13 + (17 * req_id) % block_size for req_id in range(len(seq_lens))
141 + ]
142 + block_table, cu_seq_lens, total_blocks = make_page_table(
143 + seq_lens, block_size, seq_starts_list
144 + )
145 + src_cache = torch.empty(
146 + (total_blocks, block_size, entry_size),
147 + dtype=torch.float8_e4m3fn,
148 + device="cuda",
149 + )
150 + dst = torch.empty((sum(seq_lens), entry_size), dtype=torch.bfloat16, device="cuda")
151 + token_to_seq = torch.repeat_interleave(
152 + torch.arange(len(seq_lens), dtype=torch.int32, device="cuda"),
153 + torch.tensor(seq_lens, dtype=torch.int32, device="cuda"),
154 + )
155 + seq_starts = torch.tensor(seq_starts_list, dtype=torch.int32, device="cuda")
156 + scale = torch.tensor(0.1, dtype=torch.float32, device="cuda")
157 +
158 + def run() -> None:
159 + ops.gather_and_maybe_dequant_cache(
160 + src_cache,
161 + dst,
162 + block_table,
163 + cu_seq_lens,
164 + token_to_seq,
165 + sum(seq_lens),
166 + "fp8",
167 + scale,
168 + seq_starts,
169 + )
170 +
171 + bytes_moved = sum(seq_lens) * entry_size * 3
172 + return run, bytes_moved
173 +
174 +
175 +@torch.inference_mode()
176 +def run_scenario(
177 + variant: str,
178 + name: str,
179 + seq_lens: list[int],
180 + block_size: int,
181 + entry_size: int,
182 + dtype: torch.dtype,
183 + warmup_ms: int,
184 + rep_ms: int,
185 +) -> None:
186 + if variant == "cache":
187 + run, bytes_moved = make_cache_gather(seq_lens, block_size, entry_size, dtype)
188 + elif variant == "fp8-upconvert":
189 + run, bytes_moved = make_fp8_upconvert(seq_lens, block_size)
190 + else:
191 + run, bytes_moved = make_maybe_dequant_gather(seq_lens, block_size, entry_size)
192 +
193 + latency_ms = triton.testing.do_bench(
194 + run, warmup=warmup_ms, rep=rep_ms, return_mode="median"
195 + )
196 + bandwidth_gbps = bytes_moved / latency_ms / 1e6
197 + lengths = ",".join(str(seq_len) for seq_len in seq_lens)
198 + print(
199 + f"{variant:15s} {name:10s} batch={len(seq_lens):2d} "
200 + f"total={sum(seq_lens):7d} latency={latency_ms * 1e3:9.2f} us "
201 + f"bandwidth={bandwidth_gbps:8.1f} GB/s lengths=[{lengths}]"
202 + )
203 +
204 +
205 +def main() -> None:
206 + parser = FlexibleArgumentParser(description="Benchmark cp_gather variants")
207 + parser.add_argument(
208 + "--variant",
209 + choices=["all", "cache", "fp8-upconvert", "maybe-dequant"],
210 + default="all",
211 + )
212 + parser.add_argument("--scenario", choices=["all", *SCENARIOS], default="all")
213 + parser.add_argument("--dtype", choices=DTYPES, default="fp8")
214 + parser.add_argument("--block-size", type=int, default=64)
215 + parser.add_argument("--entry-size", type=int, default=576)
216 + parser.add_argument("--warmup-ms", type=int, default=25)
217 + parser.add_argument("--rep-ms", type=int, default=100)
218 + parser.add_argument("--seed", type=int, default=0)
219 + args = parser.parse_args()
220 +
221 + torch.manual_seed(args.seed)
222 + variants = (
223 + ("cache", "fp8-upconvert", "maybe-dequant")
224 + if args.variant == "all"
225 + else (args.variant,)
226 + )
227 + scenarios = (
228 + SCENARIOS
229 + if args.scenario == "all"
230 + else {args.scenario: SCENARIOS[args.scenario]}
231 + )
232 + for variant in variants:
233 + for name, seq_lens in scenarios.items():
234 + run_scenario(
235 + variant,
236 + name,
237 + seq_lens,
238 + args.block_size,
239 + args.entry_size,
240 + DTYPES[args.dtype],
241 + args.warmup_ms,
242 + args.rep_ms,
243 + )
244 + torch.accelerator.empty_cache()
245 +
246 +
247 +if __name__ == "__main__":
248 + main()
csrc/libtorch_stable/cache_kernels.cu
+326 -232
@@ -1032,83 +1032,121 @@ void convert_fp8(torch::stable::Tensor& dst_cache,
1032
1033 namespace vllm {
1034
1035 -// grid is launched with dimensions (batch, num_splits)
1035 +struct GatherPageTask {
1036 + int32_t req_id;
1037 + int32_t logical_block;
1038 + int32_t page_token_begin;
1039 + int32_t page_token_end;
1040 + int32_t output_token_begin;
1041 +};
1042 +
1043 +template <bool has_terminal_start>
1044 +__device__ __forceinline__ bool map_gather_page_task(
1045 + int32_t task, const int32_t* __restrict__ output_starts, int32_t num_reqs,
1046 + int32_t total_tokens, int32_t block_size,
1047 + const int32_t* __restrict__ seq_starts, GatherPageTask& page) {
1048 + int32_t relative_page = task;
1049 + for (int32_t req_id = 0; req_id < num_reqs; ++req_id) {
1050 + const int32_t output_begin = min(output_starts[req_id], total_tokens);
1051 + int32_t output_end;
1052 + if constexpr (has_terminal_start) {
1053 + output_end = min(output_starts[req_id + 1], total_tokens);
1054 + } else {
1055 + output_end =
1056 + min(req_id + 1 < num_reqs ? output_starts[req_id + 1] : total_tokens,
1057 + total_tokens);
1058 + }
1059 + const int32_t seq_len = max(output_end - output_begin, 0);
1060 + const int32_t source_begin = seq_starts == nullptr ? 0 : seq_starts[req_id];
1061 + const int32_t first_block = source_begin / block_size;
1062 + const int32_t num_pages =
1063 + cuda_utils::ceil_div(source_begin + seq_len, block_size) - first_block;
1064 + if (relative_page < num_pages) {
1065 + page.req_id = req_id;
1066 + page.logical_block = first_block + relative_page;
1067 + const int32_t page_begin = page.logical_block * block_size;
1068 + const int32_t copy_begin = max(source_begin, page_begin);
1069 + const int32_t copy_end =
1070 + min(source_begin + seq_len, page_begin + block_size);
1071 + page.page_token_begin = copy_begin - page_begin;
1072 + page.page_token_end = copy_end - page_begin;
1073 + page.output_token_begin = output_begin + copy_begin - source_begin;
1074 + return true;
1075 + }
1076 + relative_page -= num_pages;
1077 + }
1078 + return false;
1079 +}
1080 +
1081 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
1037 - int ENTRY_SIZE, int CTA_SIZE>
1038 -__global__ void gather_and_maybe_dequant_cache(
1039 - const cache_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
1040 - // ENTRIES...]
1041 - scalar_t* __restrict__ dst, // [TOT_TOKENS, ENTRIES...]
1042 - const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES]
1043 - const int32_t* __restrict__ cu_seq_lens, // [BATCH+1]
1044 - const int32_t* __restrict__ token_to_seq, // [MAX_TOKEN_ACROSS_CHUNK]
1082 + int ENTRY_SIZE>
1083 +__global__ void gather_and_maybe_dequant_cache_page(
1084 + const cache_t* __restrict__ src_cache, scalar_t* __restrict__ dst,
1085 + const int32_t* __restrict__ block_table,
1086 + const int32_t* __restrict__ cu_seq_lens, const int32_t num_reqs,
1087 const int32_t num_tokens, const int32_t block_size,
1088 const int64_t block_table_stride, const int64_t cache_block_stride,
1089 const int64_t cache_entry_stride, const int64_t dst_entry_stride,
1048 - const float* __restrict__ scale,
1049 - const int32_t* __restrict__ seq_starts) { // Optional: starting offsets per
1050 - // batch
1051 - constexpr int vec_size = sizeof(float4) / sizeof(scalar_t);
1090 + const float* __restrict__ scale, const int32_t* __restrict__ seq_starts) {
1091 + constexpr int32_t vec_size = sizeof(float4) / sizeof(scalar_t);
1092 + constexpr int32_t vec_iter_cnt = ENTRY_SIZE / vec_size;
1093 + static_assert(ENTRY_SIZE % vec_size == 0);
1094 using ltype = vllm::vec_n_t<cache_t, vec_size>;
1095 using stype = vllm::vec_n_t<scalar_t, vec_size>;
1054 - // We are adding this for code readability which will be optimized out when
1055 - // build in release.
1056 - assert(CTA_SIZE == blockDim.x);
1096
1058 -#pragma unroll
1059 - for (int token_id = blockIdx.x; token_id < num_tokens;
1060 - token_id += gridDim.x) {
1061 - int64_t batch_id = token_to_seq[token_id];
1062 - int64_t batch_start = cu_seq_lens[batch_id];
1063 - int64_t batch_end = cu_seq_lens[batch_id + 1];
1064 - int32_t batch_offset = token_id - batch_start;
1065 -
1066 - if (token_id >= batch_end) return;
1067 - int32_t offset = 0;
1068 - if (seq_starts != nullptr) {
1069 - offset = seq_starts[batch_id];
1097 + __shared__ GatherPageTask page;
1098 + __shared__ int32_t physical_block;
1099 + __shared__ bool has_task;
1100 + __shared__ bool copy_task;
1101 + __shared__ float scale_value;
1102 +
1103 + if (threadIdx.x == 0) {
1104 + if constexpr (kv_dt != Fp8KVCacheDataType::kAuto) {
1105 + scale_value = *scale;
1106 }
1071 - batch_offset += offset;
1072 - int32_t block_table_id = batch_offset / block_size;
1073 - int32_t slot_id = batch_offset % block_size;
1074 - // seq_starts may push the block index past the end of the batch's block
1075 - // table row.
1076 - if (block_table_id >= block_table_stride) continue;
1077 - int32_t block_table_offset = batch_id * block_table_stride + block_table_id;
1078 - int32_t block_id = block_table[block_table_offset];
1079 - int64_t cache_offset =
1080 - block_id * cache_block_stride + slot_id * cache_entry_stride;
1081 - constexpr int32_t vec_iter_cnt = ENTRY_SIZE / vec_size;
1082 - scalar_t* dst_ = dst + token_id * dst_entry_stride;
1083 - cache_t* src_ = const_cast<cache_t*>(src_cache) + cache_offset;
1107 + }
1108 + __syncthreads();
1109 +
1110 + if (threadIdx.x == 0) {
1111 + has_task =
1112 + map_gather_page_task<true>(blockIdx.x, cu_seq_lens, num_reqs,
1113 + num_tokens, block_size, seq_starts, page);
1114 + copy_task = has_task && page.logical_block < block_table_stride;
1115 + if (copy_task) {
1116 + physical_block =
1117 + block_table[page.req_id * block_table_stride + page.logical_block];
1118 + }
1119 + }
1120 + __syncthreads();
1121 + if (!has_task) {
1122 + return;
1123 + }
1124 +
1125 + if (copy_task) {
1126 + const int32_t page_vectors =
1127 + (page.page_token_end - page.page_token_begin) * vec_iter_cnt;
1128 + for (int32_t flat_idx = threadIdx.x; flat_idx < page_vectors;
1129 + flat_idx += blockDim.x) {
1130 + const int32_t token_offset = flat_idx / vec_iter_cnt;
1131 + const int32_t idx = flat_idx - token_offset * vec_iter_cnt;
1132 + const int32_t page_token = page.page_token_begin + token_offset;
1133 + const int32_t output_token = page.output_token_begin + token_offset;
1134 + const cache_t* src = src_cache + physical_block * cache_block_stride +
1135 + page_token * cache_entry_stride;
1136 + scalar_t* output = dst + output_token * dst_entry_stride;
1137
1085 -#pragma unroll
1086 - for (int idx = threadIdx.x; idx < vec_iter_cnt; idx += CTA_SIZE) {
1138 if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
1088 - reinterpret_cast<stype*>(dst_)[idx] =
1089 - static_cast<stype>(reinterpret_cast<ltype*>(src_)[idx]);
1139 + reinterpret_cast<stype*>(output)[idx] =
1140 + static_cast<stype>(reinterpret_cast<const ltype*>(src)[idx]);
1141 } else {
1091 - ltype loaded_val = reinterpret_cast<ltype*>(src_)[idx];
1092 - stype store_val;
1142 + const ltype loaded = reinterpret_cast<const ltype*>(src)[idx];
1143 + stype converted;
1144 #pragma unroll
1094 - for (int j = 0; j < vec_size; ++j) {
1095 - store_val.val[j] = fp8::scaled_convert<scalar_t, cache_t, kv_dt>(
1096 - loaded_val.val[j], *scale);
1145 + for (int32_t j = 0; j < vec_size; ++j) {
1146 + converted.val[j] = fp8::scaled_convert<scalar_t, cache_t, kv_dt>(
1147 + loaded.val[j], scale_value);
1148 }
1098 - reinterpret_cast<stype*>(dst_)[idx] = store_val;
1099 - }
1100 - }
1101 - // process tail
1102 - constexpr int32_t tail_cnt = ENTRY_SIZE % vec_size;
1103 - dst_ = dst_ + ENTRY_SIZE - tail_cnt;
1104 - src_ = src_ + ENTRY_SIZE - tail_cnt;
1105 -#pragma unroll
1106 - for (int idx = threadIdx.x; idx < tail_cnt; idx += CTA_SIZE) {
1107 - if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
1108 - dst_[idx] = static_cast<scalar_t>(src_[idx]);
1109 - } else {
1110 - dst_[idx] =
1111 - fp8::scaled_convert<scalar_t, cache_t, kv_dt>(src_[idx], *scale);
1149 + reinterpret_cast<stype*>(output)[idx] = converted;
1150 }
1151 }
1152 }
@@ -1120,18 +1158,17 @@ __global__ void gather_and_maybe_dequant_cache(
1158 // SCALAR_T is the data type of the destination tensor.
1159 // CACHE_T is the stored data type of kv-cache.
1160 // KV_DTYPE is the real data type of kv-cache.
1123 -#define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, ENTRY_SZ) \
1124 - vllm::gather_and_maybe_dequant_cache<SCALAR_T, CACHE_T, KV_DTYPE, ENTRY_SZ, \
1125 - thread_block_size> \
1126 - <<<grid, block, 0, stream>>>( \
1127 - reinterpret_cast<CACHE_T*>(src_cache.data_ptr()), \
1128 - reinterpret_cast<SCALAR_T*>(dst.data_ptr()), \
1129 - block_table.const_data_ptr<int32_t>(), \
1130 - cu_seq_lens.const_data_ptr<int32_t>(), \
1131 - token_to_seq.const_data_ptr<int32_t>(), num_tokens, block_size, \
1132 - block_table_stride, cache_block_stride, cache_entry_stride, \
1133 - dst_entry_stride, reinterpret_cast<const float*>(scale.data_ptr()), \
1134 - seq_starts_ptr);
1161 +#define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, ENTRY_SZ) \
1162 + vllm::gather_and_maybe_dequant_cache_page<SCALAR_T, CACHE_T, KV_DTYPE, \
1163 + ENTRY_SZ> \
1164 + <<<grid, block, 0, stream>>>( \
1165 + reinterpret_cast<CACHE_T*>(src_cache.data_ptr()), \
1166 + reinterpret_cast<SCALAR_T*>(dst.data_ptr()), \
1167 + block_table.const_data_ptr<int32_t>(), \
1168 + cu_seq_lens.const_data_ptr<int32_t>(), num_reqs, \
1169 + static_cast<int32_t>(num_tokens), block_size, block_table_stride, \
1170 + cache_block_stride, cache_entry_stride, dst_entry_stride, \
1171 + reinterpret_cast<const float*>(scale.data_ptr()), seq_starts_ptr)
1172
1173 #define CALL_GATHER_CACHE_576(SCALAR_T, CACHE_T, KV_DTYPE) \
1174 CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 576)
@@ -1142,7 +1179,6 @@ __global__ void gather_and_maybe_dequant_cache(
1179 // Gather sequences from the cache into the destination tensor.
1180 // - cu_seq_lens contains the cumulative sequence lengths for each batch
1181 // - block_table contains the cache block indices for each sequence
1145 -// - token_to_seq contains the back mapping from token_id to batch_id
1182 // - Optionally, seq_starts (if provided) offsets the starting block index by
1183 // (seq_starts[bid] / page_size)
1184 void gather_and_maybe_dequant_cache(
@@ -1158,6 +1194,7 @@ void gather_and_maybe_dequant_cache(
1194 torch::stable::accelerator::DeviceGuard device_guard(
1195 src_cache.get_device_index());
1196 const cudaStream_t stream = get_current_cuda_stream();
1197 + (void)token_to_seq;
1198
1199 int32_t block_size = src_cache.size(1);
1200 int32_t head_dim = dst.size(-1);
@@ -1189,14 +1226,21 @@ void gather_and_maybe_dequant_cache(
1226 "src_cache and seq_starts must be on the same device");
1227 }
1228
1192 - int64_t block_table_stride = block_table.stride(0);
1193 - int64_t cache_block_stride = src_cache.stride(0);
1194 - int64_t cache_entry_stride = src_cache.stride(1);
1195 - int64_t dst_entry_stride = dst.stride(0);
1229 + if (num_tokens == 0) {
1230 + return;
1231 + }
1232
1197 - constexpr int32_t thread_block_size = 64;
1198 - dim3 grid(num_tokens);
1199 - dim3 block(thread_block_size);
1233 + const int32_t num_reqs = cu_seq_lens.size(0) - 1;
1234 + const int64_t block_table_stride = block_table.stride(0);
1235 + const int64_t cache_block_stride = src_cache.stride(0);
1236 + const int64_t cache_entry_stride = src_cache.stride(1);
1237 + const int64_t dst_entry_stride = dst.stride(0);
1238 + const int32_t page_threads = num_tokens >= (1 << 20) ? 128 : 256;
1239 + const int32_t required_blocks =
1240 + cuda_utils::ceil_div(static_cast<int32_t>(num_tokens), block_size) +
1241 + 2 * num_reqs;
1242 + const dim3 grid(required_blocks);
1243 + const dim3 block(page_threads);
1244
1245 const int32_t* seq_starts_ptr =
1246 seq_starts.has_value() ? seq_starts.value().const_data_ptr<int32_t>()
@@ -1213,151 +1257,161 @@ void gather_and_maybe_dequant_cache(
1257
1258 namespace vllm {
1259
1216 -// Gather and upconvert FP8 KV cache tokens to BF16 workspace
1217 -// Similar to cp_gather_cache but specifically for FP8->BF16 conversion
1218 -__global__ void cp_gather_and_upconvert_fp8_kv_cache(
1219 - const uint8_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656]
1220 - __nv_bfloat16* __restrict__ dst, // [total_tokens, 576]
1221 - const int32_t* __restrict__ block_table, // [num_reqs, BLOCK_INDICES]
1222 - const int32_t* __restrict__ workspace_starts, // [num_reqs]
1223 - const int32_t num_reqs, const int32_t block_size,
1224 - const int32_t total_tokens, const int64_t block_table_stride,
1225 - const int64_t cache_block_stride, const int64_t cache_entry_stride,
1226 - const int64_t dst_entry_stride,
1227 - const int32_t* __restrict__ seq_starts) { // Optional source offsets
1228 - const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
1229 - if (flat_warp_id >= total_tokens) return;
1230 - const int lane_id = threadIdx.x & 31;
1231 -
1232 - // Binary search to find which request owns this output token
1233 - int lo = 0, hi = num_reqs - 1;
1234 - while (lo < hi) {
1235 - int mid = (lo + hi + 1) >> 1;
1236 - if (workspace_starts[mid] <= flat_warp_id)
1237 - lo = mid;
1238 - else
1239 - hi = mid - 1;
1240 - }
1241 - const int req_id = lo;
1242 -
1243 - // Compute physical token address via block table
1244 - const int out_token_id = flat_warp_id;
1245 - int token_offset = out_token_id - workspace_starts[req_id];
1246 - if (seq_starts != nullptr) token_offset += seq_starts[req_id];
1247 - const int cache_block_idx = token_offset / block_size;
1248 - const int offset_in_block = token_offset % block_size;
1249 - const int physical_block =
1250 - block_table[req_id * block_table_stride + cache_block_idx];
1251 -
1252 - const uint8_t* token_ptr = src_cache + physical_block * cache_block_stride +
1253 - offset_in_block * cache_entry_stride;
1254 -
1255 - const int4* nope_src = reinterpret_cast<const int4*>(token_ptr);
1256 - const int4 fp8_data = nope_src[lane_id];
1257 -
1258 - const float* scales_ptr = reinterpret_cast<const float*>(token_ptr + 512);
1259 - const float scale = scales_ptr[lane_id >> 3];
1260 +__device__ __forceinline__ void gather_and_upconvert_fp8_token(
1261 + const uint8_t* __restrict__ token_ptr, __nv_bfloat16* __restrict__ dst_ptr,
1262 + int32_t lane_id) {
1263 + const uint2* fp8_src = reinterpret_cast<const uint2*>(token_ptr);
1264 + const float* scales = reinterpret_cast<const float*>(token_ptr + 512);
1265 + int4* nope_dst = reinterpret_cast<int4*>(dst_ptr);
1266
1261 - const uint2 fp8_lo = make_uint2(fp8_data.x, fp8_data.y);
1262 - const uint2 fp8_hi = make_uint2(fp8_data.z, fp8_data.w);
1267 +#pragma unroll
1268 + for (int32_t phase = 0; phase < 2; ++phase) {
1269 + const int32_t chunk = phase * 32 + lane_id;
1270 + const uint2 fp8_data = fp8_src[chunk];
1271 + const float scale = scales[chunk >> 4];
1272 #ifdef USE_ROCM
1264 - const bf16_8_t bf16_lo =
1265 - fp8::scaled_vec_conversion<bf16_8_t, uint2>(fp8_lo, scale);
1266 - const bf16_8_t bf16_hi =
1267 - fp8::scaled_vec_conversion<bf16_8_t, uint2>(fp8_hi, scale);
1273 + const bf16_8_t bf16_data =
1274 + fp8::scaled_vec_conversion<bf16_8_t, uint2>(fp8_data, scale);
1275 #else
1269 - const bf16_8_t bf16_lo =
1270 - fp8::scaled_vec_conversion<bf16_8_t, uint2>(fp8_lo, scale, __NV_E4M3);
1271 - const bf16_8_t bf16_hi =
1272 - fp8::scaled_vec_conversion<bf16_8_t, uint2>(fp8_hi, scale, __NV_E4M3);
1276 + const bf16_8_t bf16_data =
1277 + fp8::scaled_vec_conversion<bf16_8_t, uint2>(fp8_data, scale, __NV_E4M3);
1278 #endif
1274 -
1275 - __nv_bfloat16* dst_ptr = dst + out_token_id * dst_entry_stride;
1276 - int4* nope_dst = reinterpret_cast<int4*>(dst_ptr) + lane_id * 2;
1277 - nope_dst[0] = *reinterpret_cast<const int4*>(&bf16_lo);
1278 - nope_dst[1] = *reinterpret_cast<const int4*>(&bf16_hi);
1279 + nope_dst[chunk] = *reinterpret_cast<const int4*>(&bf16_data);
1280 + }
1281
1282 const int* rope_src = reinterpret_cast<const int*>(token_ptr + 528);
1283 int* rope_dst = reinterpret_cast<int*>(dst_ptr + 512);
1284 rope_dst[lane_id] = rope_src[lane_id];
1285 }
1286
1285 -template <typename scalar_t>
1286 -// Note(hc): The cp_gather_cache allows seq_starts to no longer be divisible by
1287 -// block_size.
1288 -__global__ void cp_gather_cache(
1289 - const scalar_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
1290 - // ENTRY_SIZE]
1291 - scalar_t* __restrict__ dst, // [TOT_TOKENS, ENTRY_SIZE]
1292 - const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES]
1293 - const int32_t* __restrict__ cu_seq_lens, // [BATCH+1]
1294 - const int32_t block_size, const int32_t entry_size,
1287 +__global__ void cp_gather_and_upconvert_fp8_kv_cache_page(
1288 + const uint8_t* __restrict__ src_cache, __nv_bfloat16* __restrict__ dst,
1289 + const int32_t* __restrict__ block_table,
1290 + const int32_t* __restrict__ workspace_starts, const int32_t num_reqs,
1291 + const int32_t block_size, const int32_t total_tokens,
1292 const int64_t block_table_stride, const int64_t cache_block_stride,
1293 const int64_t cache_entry_stride, const int64_t dst_entry_stride,
1297 - const int32_t* __restrict__ seq_starts // Optional: starting offsets per
1298 - // batch
1299 -) {
1300 - const int64_t bid = blockIdx.x; // Batch ID
1301 - const int32_t num_splits = gridDim.y;
1302 - const int32_t split = blockIdx.y;
1303 - const int32_t seq_start = cu_seq_lens[bid];
1304 - const int32_t seq_end = cu_seq_lens[bid + 1];
1305 - const int32_t seq_len = seq_end - seq_start;
1306 - const int32_t tot_slots = seq_len;
1307 - const int32_t split_slots = cuda_utils::ceil_div(tot_slots, num_splits);
1308 -
1309 - const int32_t split_start = split * split_slots;
1310 - const int32_t split_end = min((split + 1) * split_slots, tot_slots);
1311 -
1312 - const bool is_active_split = (split_start < tot_slots);
1313 -
1314 - if (!is_active_split) return;
1315 -
1316 - // Adjust the pointer for the block_table for this batch.
1317 - // If seq_starts is provided, compute an offset based on it
1318 - const int32_t batch_offset = bid * block_table_stride;
1319 - int32_t offset = split_start;
1320 - if (seq_starts != nullptr) {
1321 - offset += seq_starts[bid];
1294 + const int32_t* __restrict__ seq_starts) {
1295 + constexpr int32_t warps_per_cta = 16;
1296 + __shared__ GatherPageTask page;
1297 + __shared__ int32_t physical_block;
1298 + __shared__ bool has_task;
1299 +
1300 + if (threadIdx.x == 0) {
1301 + has_task =
1302 + map_gather_page_task<false>(blockIdx.x, workspace_starts, num_reqs,
1303 + total_tokens, block_size, seq_starts, page);
1304 + if (has_task) {
1305 + physical_block =
1306 + block_table[page.req_id * block_table_stride + page.logical_block];
1307 + }
1308 + }
1309 + __syncthreads();
1310 + if (!has_task) {
1311 + return;
1312 }
1323 - int32_t offset_div = offset / block_size;
1324 - offset = offset % block_size;
1325 - const int32_t* batch_block_table = block_table + batch_offset;
1313
1327 - // Adjust dst pointer based on the cumulative sequence lengths.
1328 - dst += seq_start * dst_entry_stride;
1314 + const int32_t warp_id = threadIdx.x >> 5;
1315 + const int32_t lane_id = threadIdx.x & 31;
1316 + for (int32_t page_token = page.page_token_begin + warp_id;
1317 + page_token < page.page_token_end; page_token += warps_per_cta) {
1318 + const int32_t output_token =
1319 + page.output_token_begin + page_token - page.page_token_begin;
1320 + const uint8_t* token_ptr = src_cache + physical_block * cache_block_stride +
1321 + page_token * cache_entry_stride;
1322 + __nv_bfloat16* dst_ptr = dst + output_token * dst_entry_stride;
1323 + gather_and_upconvert_fp8_token(token_ptr, dst_ptr, lane_id);
1324 + }
1325 +}
1326
1330 - auto copy_entry = [&](const scalar_t* __restrict__ _src,
1331 - scalar_t* __restrict__ _dst) {
1332 - for (int i = threadIdx.x; i < entry_size; i += blockDim.x)
1333 - _dst[i] = _src[i];
1334 - };
1327 +template <bool contiguous_entries, bool vectorized>
1328 +__global__ void cp_gather_cache_page(
1329 + const uint8_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
1330 + // ENTRY_SIZE_BYTES]
1331 + uint8_t* __restrict__ dst, // [TOT_TOKENS, ENTRY_SIZE_BYTES]
1332 + const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES]
1333 + const int32_t* __restrict__ cu_seq_lens, // [BATCH+1]
1334 + const int32_t num_reqs, const int32_t block_size,
1335 + const int32_t entry_size_bytes, const int32_t total_tokens,
1336 + const int64_t block_table_stride, const int64_t cache_block_stride,
1337 + const int64_t cache_entry_stride, const int64_t dst_entry_stride,
1338 + const int32_t* __restrict__ seq_starts) {
1339 + constexpr int32_t threads_per_token = 32;
1340 + constexpr int32_t tokens_per_cta = 256 / threads_per_token;
1341 + __shared__ GatherPageTask page;
1342 + __shared__ int32_t physical_block;
1343 + __shared__ bool has_task;
1344 +
1345 + if (threadIdx.x == 0) {
1346 + has_task =
1347 + map_gather_page_task<true>(blockIdx.x, cu_seq_lens, num_reqs,
1348 + total_tokens, block_size, seq_starts, page);
1349 + if (has_task) {
1350 + physical_block =
1351 + block_table[page.req_id * block_table_stride + page.logical_block];
1352 + }
1353 + }
1354 + __syncthreads();
1355 + if (!has_task) {
1356 + return;
1357 + }
1358
1336 - for (int pid = split_start; pid < split_end; ++pid) {
1337 - auto block_id = batch_block_table[offset_div];
1338 - auto block_start_ptr = src_cache + block_id * cache_block_stride;
1339 - auto block_dst_ptr = dst + pid * dst_entry_stride;
1340 - copy_entry(block_start_ptr + offset * cache_entry_stride, block_dst_ptr);
1341 - offset += 1;
1342 - // bump to next block
1343 - if (offset == block_size) {
1344 - offset_div += 1;
1345 - offset = 0;
1359 + if constexpr (contiguous_entries) {
1360 + const int64_t src_offset = physical_block * cache_block_stride +
1361 + page.page_token_begin * entry_size_bytes;
1362 + const int64_t dst_offset = page.output_token_begin * entry_size_bytes;
1363 + const int32_t copy_bytes =
1364 + (page.page_token_end - page.page_token_begin) * entry_size_bytes;
1365 + if constexpr (vectorized) {
1366 + const int4* src = reinterpret_cast<const int4*>(src_cache + src_offset);
1367 + int4* output = reinterpret_cast<int4*>(dst + dst_offset);
1368 + const int32_t num_vectors = copy_bytes / sizeof(int4);
1369 + int32_t idx = threadIdx.x;
1370 + for (; idx + blockDim.x < num_vectors; idx += 2 * blockDim.x) {
1371 + const int4 first = src[idx];
1372 + const int4 second = src[idx + blockDim.x];
1373 + output[idx] = first;
1374 + output[idx + blockDim.x] = second;
1375 + }
1376 + for (; idx < num_vectors; idx += blockDim.x) {
1377 + output[idx] = src[idx];
1378 + }
1379 + } else {
1380 + const uint8_t* src = src_cache + src_offset;
1381 + uint8_t* output = dst + dst_offset;
1382 + for (int32_t idx = threadIdx.x; idx < copy_bytes; idx += blockDim.x) {
1383 + output[idx] = src[idx];
1384 + }
1385 + }
1386 + } else {
1387 + const int32_t token_group = threadIdx.x / threads_per_token;
1388 + const int32_t lane_id = threadIdx.x % threads_per_token;
1389 + for (int32_t page_token = page.page_token_begin + token_group;
1390 + page_token < page.page_token_end; page_token += tokens_per_cta) {
1391 + const int32_t output_token =
1392 + page.output_token_begin + page_token - page.page_token_begin;
1393 + const uint8_t* src = src_cache + physical_block * cache_block_stride +
1394 + page_token * cache_entry_stride;
1395 + uint8_t* output = dst + output_token * dst_entry_stride;
1396 + if constexpr (vectorized) {
1397 + const int4* src_vec = reinterpret_cast<const int4*>(src);
1398 + int4* dst_vec = reinterpret_cast<int4*>(output);
1399 + const int32_t num_vectors = entry_size_bytes / sizeof(int4);
1400 + for (int32_t idx = lane_id; idx < num_vectors;
1401 + idx += threads_per_token) {
1402 + dst_vec[idx] = src_vec[idx];
1403 + }
1404 + } else {
1405 + for (int32_t idx = lane_id; idx < entry_size_bytes;
1406 + idx += threads_per_token) {
1407 + output[idx] = src[idx];
1408 + }
1409 + }
1410 }
1411 }
1412 }
1413 } // namespace vllm
1414
1351 -// Macro to dispatch the kernel based on the data type.
1352 -#define CALL_CP_GATHER_CACHE(CPY_DTYPE) \
1353 - vllm::cp_gather_cache<CPY_DTYPE><<<grid, block, 0, stream>>>( \
1354 - reinterpret_cast<CPY_DTYPE*>(src_cache.data_ptr()), \
1355 - reinterpret_cast<CPY_DTYPE*>(dst.data_ptr()), \
1356 - block_table.const_data_ptr<int32_t>(), \
1357 - cu_seq_lens.const_data_ptr<int32_t>(), block_size, entry_size, \
1358 - block_table_stride, cache_block_stride, cache_entry_stride, \
1359 - dst_entry_stride, seq_starts_ptr);
1360 -
1415 // Gather sequences from the cache into the destination tensor.
1416 // - cu_seq_lens contains the cumulative sequence lengths for each batch
1417 // - block_table contains the cache block indices for each sequence
@@ -1401,33 +1455,70 @@ void cp_gather_cache(
1455 "src_cache and seq_starts must be on the same device");
1456 }
1457
1404 - int64_t block_table_stride = block_table.stride(0);
1405 - int64_t cache_block_stride = src_cache.stride(0);
1406 - int64_t cache_entry_stride = src_cache.stride(1);
1407 - int64_t dst_entry_stride = dst.stride(0);
1408 -
1409 - // Decide on the number of splits based on the batch size.
1410 - int num_splits = batch_size > 128 ? 2 : batch_size > 64 ? 4 : 16;
1411 - dim3 grid(batch_size, num_splits);
1412 - dim3 block(1024);
1413 -
1458 STD_TORCH_CHECK(src_cache.scalar_type() == dst.scalar_type(),
1459 "src_cache and dst must have the same dtype");
1460 + const int32_t element_size = src_cache.element_size();
1461 + STD_TORCH_CHECK(element_size == 1 || element_size == 2 || element_size == 4,
1462 + "Unsupported data type width: ", element_size * 8);
1463 + STD_TORCH_CHECK(batch_size == cu_seq_lens.size(0) - 1,
1464 + "batch_size must match cu_seq_lens");
1465 +
1466 + const int32_t total_tokens = dst.size(0);
1467 + if (total_tokens == 0) {
1468 + return;
1469 + }
1470
1417 - const int dtype_bits = src_cache.element_size() * 8;
1471 + const int32_t entry_size_bytes = entry_size * element_size;
1472 + const int64_t block_table_stride = block_table.stride(0);
1473 + const int64_t cache_block_stride = src_cache.stride(0) * element_size;
1474 + const int64_t cache_entry_stride = src_cache.stride(1) * element_size;
1475 + const int64_t dst_entry_stride = dst.stride(0) * element_size;
1476 const int32_t* seq_starts_ptr =
1477 seq_starts.has_value() ? seq_starts.value().const_data_ptr<int32_t>()
1478 : nullptr;
1479
1422 - if (dtype_bits == 32) {
1423 - CALL_CP_GATHER_CACHE(uint32_t);
1424 - } else if (dtype_bits == 16) {
1425 - CALL_CP_GATHER_CACHE(uint16_t);
1426 - } else if (dtype_bits == 8) {
1427 - CALL_CP_GATHER_CACHE(uint8_t);
1480 + constexpr std::uintptr_t vector_alignment = alignof(int4);
1481 + const bool vectorized =
1482 + reinterpret_cast<std::uintptr_t>(src_cache.data_ptr()) %
1483 + vector_alignment ==
1484 + 0 &&
1485 + reinterpret_cast<std::uintptr_t>(dst.data_ptr()) % vector_alignment ==
1486 + 0 &&
1487 + entry_size_bytes % vector_alignment == 0 &&
1488 + cache_block_stride % vector_alignment == 0 &&
1489 + cache_entry_stride % vector_alignment == 0 &&
1490 + dst_entry_stride % vector_alignment == 0;
1491 + const bool contiguous_entries = cache_entry_stride == entry_size_bytes &&
1492 + dst_entry_stride == entry_size_bytes;
1493 +
1494 + const int32_t num_reqs = static_cast<int32_t>(batch_size);
1495 + const int32_t required_blocks =
1496 + cuda_utils::ceil_div(total_tokens, block_size) + 2 * num_reqs;
1497 + const dim3 grid(required_blocks);
1498 + const dim3 block(256);
1499 +
1500 +#define CALL_CP_GATHER_CACHE(CONTIGUOUS, VECTORIZED) \
1501 + vllm::cp_gather_cache_page<CONTIGUOUS, VECTORIZED> \
1502 + <<<grid, block, 0, stream>>>( \
1503 + reinterpret_cast<const uint8_t*>(src_cache.data_ptr()), \
1504 + reinterpret_cast<uint8_t*>(dst.data_ptr()), \
1505 + block_table.const_data_ptr<int32_t>(), \
1506 + cu_seq_lens.const_data_ptr<int32_t>(), num_reqs, block_size, \
1507 + entry_size_bytes, total_tokens, block_table_stride, \
1508 + cache_block_stride, cache_entry_stride, dst_entry_stride, \
1509 + seq_starts_ptr)
1510 +
1511 + if (contiguous_entries && vectorized) {
1512 + CALL_CP_GATHER_CACHE(true, true);
1513 + } else if (contiguous_entries) {
1514 + CALL_CP_GATHER_CACHE(true, false);
1515 + } else if (vectorized) {
1516 + CALL_CP_GATHER_CACHE(false, true);
1517 } else {
1429 - STD_TORCH_CHECK(false, "Unsupported data type width: ", dtype_bits);
1518 + CALL_CP_GATHER_CACHE(false, false);
1519 }
1520 +
1521 +#undef CALL_CP_GATHER_CACHE
1522 }
1523
1524 void cp_gather_and_upconvert_fp8_kv_cache(
@@ -1491,15 +1582,18 @@ void cp_gather_and_upconvert_fp8_kv_cache(
1582 }
1583
1584 const int total_tokens = dst.size(0);
1494 - constexpr int warps_per_block = 8;
1495 - const int grid_size = (total_tokens + warps_per_block - 1) / warps_per_block;
1496 - const int block_size_threads = warps_per_block * 32; // 256 threads
1585 + if (total_tokens == 0) {
1586 + return;
1587 + }
1588 + constexpr int warps_per_block = 16;
1589 + const int block_size_threads = warps_per_block * 32;
1590 const int32_t* seq_starts_ptr =
1591 seq_starts.has_value() ? seq_starts.value().const_data_ptr<int32_t>()
1592 : nullptr;
1500 -
1501 - vllm::cp_gather_and_upconvert_fp8_kv_cache<<<grid_size, block_size_threads, 0,
1502 - stream>>>(
1593 + const int required_blocks = cuda_utils::ceil_div(total_tokens, block_size) +
1594 + 2 * static_cast<int32_t>(batch_size);
1595 + vllm::cp_gather_and_upconvert_fp8_kv_cache_page<<<
1596 + required_blocks, block_size_threads, 0, stream>>>(
1597 src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()),
1598 block_table.const_data_ptr<int32_t>(),
1599 workspace_starts.const_data_ptr<int32_t>(),
tests/kernels/attention/test_cache.py
+179
@@ -1152,6 +1152,71 @@ def test_gather_and_maybe_dequant_cache_mla_with_seq_starts(
1152 torch.testing.assert_close(dst, expected)
1153
1154
1155 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
1156 +@pytest.mark.parametrize("use_seq_starts", [False, True])
1157 +@pytest.mark.parametrize("device", CUDA_DEVICES)
1158 +@torch.inference_mode()
1159 +def test_gather_and_maybe_dequant_cache_mla_large_uneven_sequences(
1160 + kv_cache_dtype,
1161 + use_seq_starts,
1162 + device,
1163 +):
1164 + block_size = 64
1165 + entry_size = 576
1166 + num_blocks = 1024
1167 + dtype = torch.bfloat16
1168 + scale = torch.tensor(0.1, dtype=torch.float32, device=device)
1169 + src_cache = _create_mla_cache(
1170 + num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device
1171 + )
1172 + _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype)
1173 +
1174 + starts = torch.tensor([3, 17, 5], dtype=torch.int32, device=device)
1175 + seq_starts = starts if use_seq_starts else None
1176 + seq_lens = torch.tensor([17, 32_768, 71], dtype=torch.int32, device=device)
1177 + batch_size = seq_lens.shape[0]
1178 + total_tokens = seq_lens.sum().item()
1179 + cu_seq_lens = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)
1180 + cu_seq_lens[1:] = seq_lens.cumsum(dim=0)
1181 + token_to_seq = torch.repeat_interleave(
1182 + torch.arange(batch_size, dtype=torch.int32, device=device), seq_lens
1183 + )
1184 + block_table = torch.stack(
1185 + [torch.randperm(num_blocks, device=device) for _ in range(batch_size)]
1186 + ).to(torch.int32)
1187 +
1188 + if kv_cache_dtype == "fp8":
1189 + dequant_src_cache = torch.empty_like(src_cache, dtype=dtype)
1190 + ops.convert_fp8(dequant_src_cache, src_cache, scale.item())
1191 + else:
1192 + dequant_src_cache = src_cache
1193 + expected_batches = []
1194 + for req_id in range(batch_size):
1195 + start = starts[req_id].item() if use_seq_starts else 0
1196 + source_tokens = torch.arange(
1197 + start, start + seq_lens[req_id].item(), device=device
1198 + )
1199 + physical_blocks = block_table[req_id, source_tokens // block_size].long()
1200 + expected_batches.append(
1201 + dequant_src_cache[physical_blocks, source_tokens % block_size]
1202 + )
1203 + expected = torch.cat(expected_batches)
1204 + dst = torch.empty_like(expected)
1205 +
1206 + ops.gather_and_maybe_dequant_cache(
1207 + src_cache,
1208 + dst,
1209 + block_table,
1210 + cu_seq_lens,
1211 + token_to_seq,
1212 + total_tokens,
1213 + kv_cache_dtype,
1214 + scale,
1215 + seq_starts,
1216 + )
1217 + torch.testing.assert_close(dst, expected)
1218 +
1219 +
1220 @pytest.mark.parametrize("kv_lora_rank", [512])
1221 @pytest.mark.parametrize("qk_rope_head_dim", [64])
1222 @pytest.mark.parametrize("block_size", [16])
@@ -1228,6 +1293,120 @@ def test_cp_gather_cache_mla(
1293 torch.testing.assert_close(dst, expected)
1294
1295
1296 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
1297 +@pytest.mark.parametrize("unaligned", [False, True])
1298 +@pytest.mark.parametrize("device", CUDA_DEVICES)
1299 +@torch.inference_mode()
1300 +def test_cp_gather_cache_mla_with_seq_starts(
1301 + kv_cache_dtype,
1302 + unaligned,
1303 + device,
1304 +):
1305 + block_size = 16
1306 + entry_size = 576
1307 + num_blocks = 128
1308 + src_storage = _create_mla_cache(
1309 + num_blocks,
1310 + block_size,
1311 + entry_size + int(unaligned),
1312 + torch.bfloat16,
1313 + kv_cache_dtype,
1314 + device,
1315 + )
1316 + src_cache = src_storage[..., 1:] if unaligned else src_storage
1317 + _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype)
1318 +
1319 + seq_starts = torch.tensor([3, 17, 5, 31], dtype=torch.int32, device=device)
1320 + seq_lens = torch.tensor([0, 1, 63, 4], dtype=torch.int32, device=device)
1321 + batch_size = seq_lens.shape[0]
1322 + cu_seq_lens = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)
1323 + cu_seq_lens[1:] = seq_lens.cumsum(dim=0)
1324 + block_table = torch.stack(
1325 + [torch.randperm(num_blocks, device=device) for _ in range(batch_size)]
1326 + ).to(torch.int32)
1327 +
1328 + expected_rows = []
1329 + for req_id in range(batch_size):
1330 + start = seq_starts[req_id].item()
1331 + for source_token in range(start, start + seq_lens[req_id].item()):
1332 + block_id = block_table[req_id, source_token // block_size]
1333 + expected_rows.append(src_cache[block_id, source_token % block_size])
1334 + expected = torch.stack(expected_rows)
1335 + if unaligned:
1336 + dst_storage = torch.empty(
1337 + (expected.shape[0], entry_size + 1),
1338 + dtype=expected.dtype,
1339 + device=device,
1340 + )
1341 + dst = dst_storage[:, 1:]
1342 + else:
1343 + dst = torch.empty_like(expected)
1344 +
1345 + ops.cp_gather_cache(
1346 + src_cache,
1347 + dst,
1348 + block_table,
1349 + cu_seq_lens,
1350 + batch_size,
1351 + seq_starts,
1352 + )
1353 + torch.testing.assert_close(dst, expected)
1354 +
1355 +
1356 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
1357 +@pytest.mark.parametrize("use_seq_starts", [False, True])
1358 +@pytest.mark.parametrize("device", CUDA_DEVICES)
1359 +@torch.inference_mode()
1360 +def test_cp_gather_cache_mla_large_uneven_sequences(
1361 + kv_cache_dtype,
1362 + use_seq_starts,
1363 + device,
1364 +):
1365 + block_size = 64
1366 + entry_size = 576
1367 + num_blocks = 1024
1368 + src_cache = _create_mla_cache(
1369 + num_blocks,
1370 + block_size,
1371 + entry_size,
1372 + torch.bfloat16,
1373 + kv_cache_dtype,
1374 + device,
1375 + )
1376 + _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype)
1377 +
1378 + starts = torch.tensor([3, 17, 5], dtype=torch.int32, device=device)
1379 + seq_starts = starts if use_seq_starts else None
1380 + seq_lens = torch.tensor([17, 32_768, 71], dtype=torch.int32, device=device)
1381 + batch_size = seq_lens.shape[0]
1382 + cu_seq_lens = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)
1383 + cu_seq_lens[1:] = seq_lens.cumsum(dim=0)
1384 + block_table = torch.stack(
1385 + [torch.randperm(num_blocks, device=device) for _ in range(batch_size)]
1386 + ).to(torch.int32)
1387 +
1388 + expected_batches = []
1389 + for req_id in range(batch_size):
1390 + start = starts[req_id].item() if use_seq_starts else 0
1391 + source_tokens = torch.arange(
1392 + start, start + seq_lens[req_id].item(), device=device
1393 + )
1394 + physical_blocks = block_table[req_id, source_tokens // block_size].long()
1395 + expected_batches.append(src_cache[physical_blocks, source_tokens % block_size])
1396 + expected = torch.cat(expected_batches)
1397 + dst = torch.empty_like(expected)
1398 +
1399 + ops.cp_gather_cache(
1400 + src_cache,
1401 + dst,
1402 + block_table,
1403 + cu_seq_lens,
1404 + batch_size,
1405 + seq_starts,
1406 + )
1407 + torch.testing.assert_close(dst, expected)
1408 +
1409 +
1410 @pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
1411 @pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
1412 @pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA)
tests/kernels/test_cp_gather_fp8.py
+54
@@ -411,3 +411,57 @@ def test_cp_gather_fp8_large_seqlens(seq_lens, block_size):
411 dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2
412 )
413 assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:])
414 +
415 +
416 +def test_cp_gather_fp8_large_uneven_sequences_with_starts():
417 + """Gather uneven long request slices through the page-oriented path."""
418 + gather_seq_lens = [17, 32_768, 71]
419 + seq_starts = [3, 17, 5]
420 + full_seq_lens = [
421 + start + length for start, length in zip(seq_starts, gather_seq_lens)
422 + ]
423 + (
424 + cache,
425 + block_table,
426 + full_workspace_starts,
427 + num_reqs,
428 + _total_tokens,
429 + full_expected,
430 + ) = _build_test_case_fast(full_seq_lens, block_size=64)
431 +
432 + workspace_starts = torch.zeros(num_reqs, dtype=torch.int32, device="cuda")
433 + workspace_starts[1:] = torch.tensor(
434 + gather_seq_lens[:-1], dtype=torch.int32, device="cuda"
435 + ).cumsum(dim=0)
436 + seq_starts_t = torch.tensor(seq_starts, dtype=torch.int32, device="cuda")
437 + dst = torch.empty(
438 + sum(gather_seq_lens),
439 + NOPE_DIM + ROPE_DIM,
440 + dtype=torch.bfloat16,
441 + device="cuda",
442 + )
443 +
444 + ops.cp_gather_and_upconvert_fp8_kv_cache(
445 + cache,
446 + dst,
447 + block_table,
448 + workspace_starts,
449 + num_reqs,
450 + seq_starts_t,
451 + )
452 +
453 + expected = torch.cat(
454 + [
455 + full_expected[
456 + full_workspace_starts[req_id]
457 + + seq_starts[req_id] : full_workspace_starts[req_id]
458 + + seq_starts[req_id]
459 + + gather_seq_lens[req_id]
460 + ]
461 + for req_id in range(num_reqs)
462 + ]
463 + )
464 + torch.testing.assert_close(
465 + dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2
466 + )
467 + assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:])
vllm/model_executor/layers/attention/mla_attention.py
+2 -2
@@ -2574,7 +2574,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2574 # FP8 path: gather cache without dequantization
2575 ops.cp_gather_cache(
2576 src_cache=kv_c_and_k_pe_cache,
2577 - dst=workspace,
2577 + dst=workspace[:toks],
2578 block_table=block_table,
2579 cu_seq_lens=chunk.cu_seq_lens,
2580 batch_size=chunk.num_requests,
@@ -2683,7 +2683,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2683 else:
2684 ops.cp_gather_cache(
2685 src_cache=kv_c_and_k_pe_cache,
2686 - dst=workspace,
2686 + dst=workspace[:toks],
2687 block_table=block_table,
2688 cu_seq_lens=padded_local_cu_seq_lens,
2689 batch_size=chunk.num_requests,