diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 6a99106205..4e73d7620e 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -224,7 +224,7 @@ steps: - label: Kernels Mamba Test device: h200_35gb key: kernels-mamba-test - timeout_in_minutes: 40 + timeout_in_minutes: 60 source_file_dependencies: - csrc/mamba/ - tests/kernels/mamba diff --git a/benchmarks/kernels/benchmark_cp_gather.py b/benchmarks/kernels/benchmark_cp_gather.py new file mode 100644 index 0000000000..5e73d6d6e9 --- /dev/null +++ b/benchmarks/kernels/benchmark_cp_gather.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math +from collections.abc import Callable + +import torch + +from vllm import _custom_ops as ops +from vllm.triton_utils import triton +from vllm.utils.argparse_utils import FlexibleArgumentParser + +SCENARIOS = { + "single-60k": [60_000], + "single-300k": [300_000], + "skew-2": [60_000, 300_000], + "skew-4": [60_000, 100_000, 180_000, 300_000], + "skew-8": [ + 60_000, + 60_000, + 80_000, + 100_000, + 140_000, + 180_000, + 240_000, + 300_000, + ], +} +DTYPES = { + "fp8": torch.float8_e4m3fn, + "bfloat16": torch.bfloat16, + "float32": torch.float32, +} + + +def make_page_table( + seq_lens: list[int], + block_size: int, + seq_starts: list[int] | None = None, +) -> tuple[torch.Tensor, torch.Tensor, int]: + if seq_starts is None: + seq_starts = [0] * len(seq_lens) + blocks_per_req = [ + math.ceil((start + length) / block_size) + for start, length in zip(seq_starts, seq_lens) + ] + total_blocks = sum(blocks_per_req) + block_table = torch.zeros( + (len(seq_lens), max(blocks_per_req)), dtype=torch.int32, device="cuda" + ) + physical_blocks = torch.randperm(total_blocks, dtype=torch.int32, device="cuda") + block_offset = 0 + for req_id, num_blocks in enumerate(blocks_per_req): + block_table[req_id, :num_blocks] = physical_blocks[ + block_offset : block_offset + num_blocks + ] + block_offset += num_blocks + + cu_seq_lens = torch.zeros(len(seq_lens) + 1, dtype=torch.int32, device="cuda") + cu_seq_lens[1:] = torch.tensor(seq_lens, dtype=torch.int32, device="cuda").cumsum( + dim=0 + ) + return block_table, cu_seq_lens, total_blocks + + +def make_cache_gather( + seq_lens: list[int], + block_size: int, + entry_size: int, + dtype: torch.dtype, +) -> tuple[Callable[[], None], int]: + seq_starts_list = [ + 13 + (17 * req_id) % block_size for req_id in range(len(seq_lens)) + ] + block_table, cu_seq_lens, total_blocks = make_page_table( + seq_lens, block_size, seq_starts_list + ) + src_cache = torch.empty( + (total_blocks, block_size, entry_size), dtype=dtype, device="cuda" + ) + dst = torch.empty((sum(seq_lens), entry_size), dtype=dtype, device="cuda") + seq_starts = torch.tensor(seq_starts_list, dtype=torch.int32, device="cuda") + + def run() -> None: + ops.cp_gather_cache( + src_cache, + dst, + block_table, + cu_seq_lens, + len(seq_lens), + seq_starts, + ) + + bytes_moved = 2 * dst.numel() * dst.element_size() + return run, bytes_moved + + +def make_fp8_upconvert( + seq_lens: list[int], + block_size: int, +) -> tuple[Callable[[], None], int]: + entry_bytes = 656 + output_elements = 576 + seq_starts_list = [ + 13 + (17 * req_id) % block_size for req_id in range(len(seq_lens)) + ] + block_table, cu_seq_lens, total_blocks = make_page_table( + seq_lens, block_size, seq_starts_list + ) + src_cache = torch.empty( + (total_blocks, block_size, entry_bytes), + dtype=torch.uint8, + device="cuda", + ) + dst = torch.empty( + (sum(seq_lens), output_elements), dtype=torch.bfloat16, device="cuda" + ) + seq_starts = torch.tensor(seq_starts_list, dtype=torch.int32, device="cuda") + + def run() -> None: + ops.cp_gather_and_upconvert_fp8_kv_cache( + src_cache, + dst, + block_table, + cu_seq_lens[:-1], + len(seq_lens), + seq_starts, + ) + + bytes_moved = sum(seq_lens) * (entry_bytes + output_elements * 2) + return run, bytes_moved + + +def make_maybe_dequant_gather( + seq_lens: list[int], + block_size: int, + entry_size: int, +) -> tuple[Callable[[], None], int]: + seq_starts_list = [ + 13 + (17 * req_id) % block_size for req_id in range(len(seq_lens)) + ] + block_table, cu_seq_lens, total_blocks = make_page_table( + seq_lens, block_size, seq_starts_list + ) + src_cache = torch.empty( + (total_blocks, block_size, entry_size), + dtype=torch.float8_e4m3fn, + device="cuda", + ) + dst = torch.empty((sum(seq_lens), entry_size), dtype=torch.bfloat16, device="cuda") + token_to_seq = torch.repeat_interleave( + torch.arange(len(seq_lens), dtype=torch.int32, device="cuda"), + torch.tensor(seq_lens, dtype=torch.int32, device="cuda"), + ) + seq_starts = torch.tensor(seq_starts_list, dtype=torch.int32, device="cuda") + scale = torch.tensor(0.1, dtype=torch.float32, device="cuda") + + def run() -> None: + ops.gather_and_maybe_dequant_cache( + src_cache, + dst, + block_table, + cu_seq_lens, + token_to_seq, + sum(seq_lens), + "fp8", + scale, + seq_starts, + ) + + bytes_moved = sum(seq_lens) * entry_size * 3 + return run, bytes_moved + + +@torch.inference_mode() +def run_scenario( + variant: str, + name: str, + seq_lens: list[int], + block_size: int, + entry_size: int, + dtype: torch.dtype, + warmup_ms: int, + rep_ms: int, +) -> None: + if variant == "cache": + run, bytes_moved = make_cache_gather(seq_lens, block_size, entry_size, dtype) + elif variant == "fp8-upconvert": + run, bytes_moved = make_fp8_upconvert(seq_lens, block_size) + else: + run, bytes_moved = make_maybe_dequant_gather(seq_lens, block_size, entry_size) + + latency_ms = triton.testing.do_bench( + run, warmup=warmup_ms, rep=rep_ms, return_mode="median" + ) + bandwidth_gbps = bytes_moved / latency_ms / 1e6 + lengths = ",".join(str(seq_len) for seq_len in seq_lens) + print( + f"{variant:15s} {name:10s} batch={len(seq_lens):2d} " + f"total={sum(seq_lens):7d} latency={latency_ms * 1e3:9.2f} us " + f"bandwidth={bandwidth_gbps:8.1f} GB/s lengths=[{lengths}]" + ) + + +def main() -> None: + parser = FlexibleArgumentParser(description="Benchmark cp_gather variants") + parser.add_argument( + "--variant", + choices=["all", "cache", "fp8-upconvert", "maybe-dequant"], + default="all", + ) + parser.add_argument("--scenario", choices=["all", *SCENARIOS], default="all") + parser.add_argument("--dtype", choices=DTYPES, default="fp8") + parser.add_argument("--block-size", type=int, default=64) + parser.add_argument("--entry-size", type=int, default=576) + parser.add_argument("--warmup-ms", type=int, default=25) + parser.add_argument("--rep-ms", type=int, default=100) + parser.add_argument("--seed", type=int, default=0) + args = parser.parse_args() + + torch.manual_seed(args.seed) + variants = ( + ("cache", "fp8-upconvert", "maybe-dequant") + if args.variant == "all" + else (args.variant,) + ) + scenarios = ( + SCENARIOS + if args.scenario == "all" + else {args.scenario: SCENARIOS[args.scenario]} + ) + for variant in variants: + for name, seq_lens in scenarios.items(): + run_scenario( + variant, + name, + seq_lens, + args.block_size, + args.entry_size, + DTYPES[args.dtype], + args.warmup_ms, + args.rep_ms, + ) + torch.accelerator.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu index b16f405a36..2dbd498fb3 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -1032,83 +1032,121 @@ void convert_fp8(torch::stable::Tensor& dst_cache, namespace vllm { -// grid is launched with dimensions (batch, num_splits) +struct GatherPageTask { + int32_t req_id; + int32_t logical_block; + int32_t page_token_begin; + int32_t page_token_end; + int32_t output_token_begin; +}; + +template +__device__ __forceinline__ bool map_gather_page_task( + int32_t task, const int32_t* __restrict__ output_starts, int32_t num_reqs, + int32_t total_tokens, int32_t block_size, + const int32_t* __restrict__ seq_starts, GatherPageTask& page) { + int32_t relative_page = task; + for (int32_t req_id = 0; req_id < num_reqs; ++req_id) { + const int32_t output_begin = min(output_starts[req_id], total_tokens); + int32_t output_end; + if constexpr (has_terminal_start) { + output_end = min(output_starts[req_id + 1], total_tokens); + } else { + output_end = + min(req_id + 1 < num_reqs ? output_starts[req_id + 1] : total_tokens, + total_tokens); + } + const int32_t seq_len = max(output_end - output_begin, 0); + const int32_t source_begin = seq_starts == nullptr ? 0 : seq_starts[req_id]; + const int32_t first_block = source_begin / block_size; + const int32_t num_pages = + cuda_utils::ceil_div(source_begin + seq_len, block_size) - first_block; + if (relative_page < num_pages) { + page.req_id = req_id; + page.logical_block = first_block + relative_page; + const int32_t page_begin = page.logical_block * block_size; + const int32_t copy_begin = max(source_begin, page_begin); + const int32_t copy_end = + min(source_begin + seq_len, page_begin + block_size); + page.page_token_begin = copy_begin - page_begin; + page.page_token_end = copy_end - page_begin; + page.output_token_begin = output_begin + copy_begin - source_begin; + return true; + } + relative_page -= num_pages; + } + return false; +} + template -__global__ void gather_and_maybe_dequant_cache( - const cache_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // ENTRIES...] - scalar_t* __restrict__ dst, // [TOT_TOKENS, ENTRIES...] - const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES] - const int32_t* __restrict__ cu_seq_lens, // [BATCH+1] - const int32_t* __restrict__ token_to_seq, // [MAX_TOKEN_ACROSS_CHUNK] + int ENTRY_SIZE> +__global__ void gather_and_maybe_dequant_cache_page( + const cache_t* __restrict__ src_cache, scalar_t* __restrict__ dst, + const int32_t* __restrict__ block_table, + const int32_t* __restrict__ cu_seq_lens, const int32_t num_reqs, const int32_t num_tokens, const int32_t block_size, const int64_t block_table_stride, const int64_t cache_block_stride, const int64_t cache_entry_stride, const int64_t dst_entry_stride, - const float* __restrict__ scale, - const int32_t* __restrict__ seq_starts) { // Optional: starting offsets per - // batch - constexpr int vec_size = sizeof(float4) / sizeof(scalar_t); + const float* __restrict__ scale, const int32_t* __restrict__ seq_starts) { + constexpr int32_t vec_size = sizeof(float4) / sizeof(scalar_t); + constexpr int32_t vec_iter_cnt = ENTRY_SIZE / vec_size; + static_assert(ENTRY_SIZE % vec_size == 0); using ltype = vllm::vec_n_t; using stype = vllm::vec_n_t; - // We are adding this for code readability which will be optimized out when - // build in release. - assert(CTA_SIZE == blockDim.x); -#pragma unroll - for (int token_id = blockIdx.x; token_id < num_tokens; - token_id += gridDim.x) { - int64_t batch_id = token_to_seq[token_id]; - int64_t batch_start = cu_seq_lens[batch_id]; - int64_t batch_end = cu_seq_lens[batch_id + 1]; - int32_t batch_offset = token_id - batch_start; + __shared__ GatherPageTask page; + __shared__ int32_t physical_block; + __shared__ bool has_task; + __shared__ bool copy_task; + __shared__ float scale_value; - if (token_id >= batch_end) return; - int32_t offset = 0; - if (seq_starts != nullptr) { - offset = seq_starts[batch_id]; + if (threadIdx.x == 0) { + if constexpr (kv_dt != Fp8KVCacheDataType::kAuto) { + scale_value = *scale; } - batch_offset += offset; - int32_t block_table_id = batch_offset / block_size; - int32_t slot_id = batch_offset % block_size; - // seq_starts may push the block index past the end of the batch's block - // table row. - if (block_table_id >= block_table_stride) continue; - int32_t block_table_offset = batch_id * block_table_stride + block_table_id; - int32_t block_id = block_table[block_table_offset]; - int64_t cache_offset = - block_id * cache_block_stride + slot_id * cache_entry_stride; - constexpr int32_t vec_iter_cnt = ENTRY_SIZE / vec_size; - scalar_t* dst_ = dst + token_id * dst_entry_stride; - cache_t* src_ = const_cast(src_cache) + cache_offset; + } + __syncthreads(); + + if (threadIdx.x == 0) { + has_task = + map_gather_page_task(blockIdx.x, cu_seq_lens, num_reqs, + num_tokens, block_size, seq_starts, page); + copy_task = has_task && page.logical_block < block_table_stride; + if (copy_task) { + physical_block = + block_table[page.req_id * block_table_stride + page.logical_block]; + } + } + __syncthreads(); + if (!has_task) { + return; + } + + if (copy_task) { + const int32_t page_vectors = + (page.page_token_end - page.page_token_begin) * vec_iter_cnt; + for (int32_t flat_idx = threadIdx.x; flat_idx < page_vectors; + flat_idx += blockDim.x) { + const int32_t token_offset = flat_idx / vec_iter_cnt; + const int32_t idx = flat_idx - token_offset * vec_iter_cnt; + const int32_t page_token = page.page_token_begin + token_offset; + const int32_t output_token = page.output_token_begin + token_offset; + const cache_t* src = src_cache + physical_block * cache_block_stride + + page_token * cache_entry_stride; + scalar_t* output = dst + output_token * dst_entry_stride; -#pragma unroll - for (int idx = threadIdx.x; idx < vec_iter_cnt; idx += CTA_SIZE) { if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { - reinterpret_cast(dst_)[idx] = - static_cast(reinterpret_cast(src_)[idx]); + reinterpret_cast(output)[idx] = + static_cast(reinterpret_cast(src)[idx]); } else { - ltype loaded_val = reinterpret_cast(src_)[idx]; - stype store_val; + const ltype loaded = reinterpret_cast(src)[idx]; + stype converted; #pragma unroll - for (int j = 0; j < vec_size; ++j) { - store_val.val[j] = fp8::scaled_convert( - loaded_val.val[j], *scale); + for (int32_t j = 0; j < vec_size; ++j) { + converted.val[j] = fp8::scaled_convert( + loaded.val[j], scale_value); } - reinterpret_cast(dst_)[idx] = store_val; - } - } - // process tail - constexpr int32_t tail_cnt = ENTRY_SIZE % vec_size; - dst_ = dst_ + ENTRY_SIZE - tail_cnt; - src_ = src_ + ENTRY_SIZE - tail_cnt; -#pragma unroll - for (int idx = threadIdx.x; idx < tail_cnt; idx += CTA_SIZE) { - if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { - dst_[idx] = static_cast(src_[idx]); - } else { - dst_[idx] = - fp8::scaled_convert(src_[idx], *scale); + reinterpret_cast(output)[idx] = converted; } } } @@ -1120,18 +1158,17 @@ __global__ void gather_and_maybe_dequant_cache( // SCALAR_T is the data type of the destination tensor. // CACHE_T is the stored data type of kv-cache. // KV_DTYPE is the real data type of kv-cache. -#define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, ENTRY_SZ) \ - vllm::gather_and_maybe_dequant_cache \ - <<>>( \ - reinterpret_cast(src_cache.data_ptr()), \ - reinterpret_cast(dst.data_ptr()), \ - block_table.const_data_ptr(), \ - cu_seq_lens.const_data_ptr(), \ - token_to_seq.const_data_ptr(), num_tokens, block_size, \ - block_table_stride, cache_block_stride, cache_entry_stride, \ - dst_entry_stride, reinterpret_cast(scale.data_ptr()), \ - seq_starts_ptr); +#define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, ENTRY_SZ) \ + vllm::gather_and_maybe_dequant_cache_page \ + <<>>( \ + reinterpret_cast(src_cache.data_ptr()), \ + reinterpret_cast(dst.data_ptr()), \ + block_table.const_data_ptr(), \ + cu_seq_lens.const_data_ptr(), num_reqs, \ + static_cast(num_tokens), block_size, block_table_stride, \ + cache_block_stride, cache_entry_stride, dst_entry_stride, \ + reinterpret_cast(scale.data_ptr()), seq_starts_ptr) #define CALL_GATHER_CACHE_576(SCALAR_T, CACHE_T, KV_DTYPE) \ CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 576) @@ -1142,7 +1179,6 @@ __global__ void gather_and_maybe_dequant_cache( // Gather sequences from the cache into the destination tensor. // - cu_seq_lens contains the cumulative sequence lengths for each batch // - block_table contains the cache block indices for each sequence -// - token_to_seq contains the back mapping from token_id to batch_id // - Optionally, seq_starts (if provided) offsets the starting block index by // (seq_starts[bid] / page_size) void gather_and_maybe_dequant_cache( @@ -1158,6 +1194,7 @@ void gather_and_maybe_dequant_cache( torch::stable::accelerator::DeviceGuard device_guard( src_cache.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(); + (void)token_to_seq; int32_t block_size = src_cache.size(1); int32_t head_dim = dst.size(-1); @@ -1189,14 +1226,21 @@ void gather_and_maybe_dequant_cache( "src_cache and seq_starts must be on the same device"); } - int64_t block_table_stride = block_table.stride(0); - int64_t cache_block_stride = src_cache.stride(0); - int64_t cache_entry_stride = src_cache.stride(1); - int64_t dst_entry_stride = dst.stride(0); + if (num_tokens == 0) { + return; + } - constexpr int32_t thread_block_size = 64; - dim3 grid(num_tokens); - dim3 block(thread_block_size); + const int32_t num_reqs = cu_seq_lens.size(0) - 1; + const int64_t block_table_stride = block_table.stride(0); + const int64_t cache_block_stride = src_cache.stride(0); + const int64_t cache_entry_stride = src_cache.stride(1); + const int64_t dst_entry_stride = dst.stride(0); + const int32_t page_threads = num_tokens >= (1 << 20) ? 128 : 256; + const int32_t required_blocks = + cuda_utils::ceil_div(static_cast(num_tokens), block_size) + + 2 * num_reqs; + const dim3 grid(required_blocks); + const dim3 block(page_threads); const int32_t* seq_starts_ptr = seq_starts.has_value() ? seq_starts.value().const_data_ptr() @@ -1213,151 +1257,161 @@ void gather_and_maybe_dequant_cache( namespace vllm { -// Gather and upconvert FP8 KV cache tokens to BF16 workspace -// Similar to cp_gather_cache but specifically for FP8->BF16 conversion -__global__ void cp_gather_and_upconvert_fp8_kv_cache( - const uint8_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656] - __nv_bfloat16* __restrict__ dst, // [total_tokens, 576] - const int32_t* __restrict__ block_table, // [num_reqs, BLOCK_INDICES] - const int32_t* __restrict__ workspace_starts, // [num_reqs] - const int32_t num_reqs, const int32_t block_size, - const int32_t total_tokens, const int64_t block_table_stride, - const int64_t cache_block_stride, const int64_t cache_entry_stride, - const int64_t dst_entry_stride, - const int32_t* __restrict__ seq_starts) { // Optional source offsets - const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5; - if (flat_warp_id >= total_tokens) return; - const int lane_id = threadIdx.x & 31; +__device__ __forceinline__ void gather_and_upconvert_fp8_token( + const uint8_t* __restrict__ token_ptr, __nv_bfloat16* __restrict__ dst_ptr, + int32_t lane_id) { + const uint2* fp8_src = reinterpret_cast(token_ptr); + const float* scales = reinterpret_cast(token_ptr + 512); + int4* nope_dst = reinterpret_cast(dst_ptr); - // Binary search to find which request owns this output token - int lo = 0, hi = num_reqs - 1; - while (lo < hi) { - int mid = (lo + hi + 1) >> 1; - if (workspace_starts[mid] <= flat_warp_id) - lo = mid; - else - hi = mid - 1; - } - const int req_id = lo; - - // Compute physical token address via block table - const int out_token_id = flat_warp_id; - int token_offset = out_token_id - workspace_starts[req_id]; - if (seq_starts != nullptr) token_offset += seq_starts[req_id]; - const int cache_block_idx = token_offset / block_size; - const int offset_in_block = token_offset % block_size; - const int physical_block = - block_table[req_id * block_table_stride + cache_block_idx]; - - const uint8_t* token_ptr = src_cache + physical_block * cache_block_stride + - offset_in_block * cache_entry_stride; - - const int4* nope_src = reinterpret_cast(token_ptr); - const int4 fp8_data = nope_src[lane_id]; - - const float* scales_ptr = reinterpret_cast(token_ptr + 512); - const float scale = scales_ptr[lane_id >> 3]; - - const uint2 fp8_lo = make_uint2(fp8_data.x, fp8_data.y); - const uint2 fp8_hi = make_uint2(fp8_data.z, fp8_data.w); +#pragma unroll + for (int32_t phase = 0; phase < 2; ++phase) { + const int32_t chunk = phase * 32 + lane_id; + const uint2 fp8_data = fp8_src[chunk]; + const float scale = scales[chunk >> 4]; #ifdef USE_ROCM - const bf16_8_t bf16_lo = - fp8::scaled_vec_conversion(fp8_lo, scale); - const bf16_8_t bf16_hi = - fp8::scaled_vec_conversion(fp8_hi, scale); + const bf16_8_t bf16_data = + fp8::scaled_vec_conversion(fp8_data, scale); #else - const bf16_8_t bf16_lo = - fp8::scaled_vec_conversion(fp8_lo, scale, __NV_E4M3); - const bf16_8_t bf16_hi = - fp8::scaled_vec_conversion(fp8_hi, scale, __NV_E4M3); + const bf16_8_t bf16_data = + fp8::scaled_vec_conversion(fp8_data, scale, __NV_E4M3); #endif - - __nv_bfloat16* dst_ptr = dst + out_token_id * dst_entry_stride; - int4* nope_dst = reinterpret_cast(dst_ptr) + lane_id * 2; - nope_dst[0] = *reinterpret_cast(&bf16_lo); - nope_dst[1] = *reinterpret_cast(&bf16_hi); + nope_dst[chunk] = *reinterpret_cast(&bf16_data); + } const int* rope_src = reinterpret_cast(token_ptr + 528); int* rope_dst = reinterpret_cast(dst_ptr + 512); rope_dst[lane_id] = rope_src[lane_id]; } -template -// Note(hc): The cp_gather_cache allows seq_starts to no longer be divisible by -// block_size. -__global__ void cp_gather_cache( - const scalar_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // ENTRY_SIZE] - scalar_t* __restrict__ dst, // [TOT_TOKENS, ENTRY_SIZE] - const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES] - const int32_t* __restrict__ cu_seq_lens, // [BATCH+1] - const int32_t block_size, const int32_t entry_size, +__global__ void cp_gather_and_upconvert_fp8_kv_cache_page( + const uint8_t* __restrict__ src_cache, __nv_bfloat16* __restrict__ dst, + const int32_t* __restrict__ block_table, + const int32_t* __restrict__ workspace_starts, const int32_t num_reqs, + const int32_t block_size, const int32_t total_tokens, const int64_t block_table_stride, const int64_t cache_block_stride, const int64_t cache_entry_stride, const int64_t dst_entry_stride, - const int32_t* __restrict__ seq_starts // Optional: starting offsets per - // batch -) { - const int64_t bid = blockIdx.x; // Batch ID - const int32_t num_splits = gridDim.y; - const int32_t split = blockIdx.y; - const int32_t seq_start = cu_seq_lens[bid]; - const int32_t seq_end = cu_seq_lens[bid + 1]; - const int32_t seq_len = seq_end - seq_start; - const int32_t tot_slots = seq_len; - const int32_t split_slots = cuda_utils::ceil_div(tot_slots, num_splits); + const int32_t* __restrict__ seq_starts) { + constexpr int32_t warps_per_cta = 16; + __shared__ GatherPageTask page; + __shared__ int32_t physical_block; + __shared__ bool has_task; - const int32_t split_start = split * split_slots; - const int32_t split_end = min((split + 1) * split_slots, tot_slots); - - const bool is_active_split = (split_start < tot_slots); - - if (!is_active_split) return; - - // Adjust the pointer for the block_table for this batch. - // If seq_starts is provided, compute an offset based on it - const int32_t batch_offset = bid * block_table_stride; - int32_t offset = split_start; - if (seq_starts != nullptr) { - offset += seq_starts[bid]; + if (threadIdx.x == 0) { + has_task = + map_gather_page_task(blockIdx.x, workspace_starts, num_reqs, + total_tokens, block_size, seq_starts, page); + if (has_task) { + physical_block = + block_table[page.req_id * block_table_stride + page.logical_block]; + } + } + __syncthreads(); + if (!has_task) { + return; } - int32_t offset_div = offset / block_size; - offset = offset % block_size; - const int32_t* batch_block_table = block_table + batch_offset; - // Adjust dst pointer based on the cumulative sequence lengths. - dst += seq_start * dst_entry_stride; + const int32_t warp_id = threadIdx.x >> 5; + const int32_t lane_id = threadIdx.x & 31; + for (int32_t page_token = page.page_token_begin + warp_id; + page_token < page.page_token_end; page_token += warps_per_cta) { + const int32_t output_token = + page.output_token_begin + page_token - page.page_token_begin; + const uint8_t* token_ptr = src_cache + physical_block * cache_block_stride + + page_token * cache_entry_stride; + __nv_bfloat16* dst_ptr = dst + output_token * dst_entry_stride; + gather_and_upconvert_fp8_token(token_ptr, dst_ptr, lane_id); + } +} - auto copy_entry = [&](const scalar_t* __restrict__ _src, - scalar_t* __restrict__ _dst) { - for (int i = threadIdx.x; i < entry_size; i += blockDim.x) - _dst[i] = _src[i]; - }; +template +__global__ void cp_gather_cache_page( + const uint8_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // ENTRY_SIZE_BYTES] + uint8_t* __restrict__ dst, // [TOT_TOKENS, ENTRY_SIZE_BYTES] + const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES] + const int32_t* __restrict__ cu_seq_lens, // [BATCH+1] + const int32_t num_reqs, const int32_t block_size, + const int32_t entry_size_bytes, const int32_t total_tokens, + const int64_t block_table_stride, const int64_t cache_block_stride, + const int64_t cache_entry_stride, const int64_t dst_entry_stride, + const int32_t* __restrict__ seq_starts) { + constexpr int32_t threads_per_token = 32; + constexpr int32_t tokens_per_cta = 256 / threads_per_token; + __shared__ GatherPageTask page; + __shared__ int32_t physical_block; + __shared__ bool has_task; - for (int pid = split_start; pid < split_end; ++pid) { - auto block_id = batch_block_table[offset_div]; - auto block_start_ptr = src_cache + block_id * cache_block_stride; - auto block_dst_ptr = dst + pid * dst_entry_stride; - copy_entry(block_start_ptr + offset * cache_entry_stride, block_dst_ptr); - offset += 1; - // bump to next block - if (offset == block_size) { - offset_div += 1; - offset = 0; + if (threadIdx.x == 0) { + has_task = + map_gather_page_task(blockIdx.x, cu_seq_lens, num_reqs, + total_tokens, block_size, seq_starts, page); + if (has_task) { + physical_block = + block_table[page.req_id * block_table_stride + page.logical_block]; + } + } + __syncthreads(); + if (!has_task) { + return; + } + + if constexpr (contiguous_entries) { + const int64_t src_offset = physical_block * cache_block_stride + + page.page_token_begin * entry_size_bytes; + const int64_t dst_offset = page.output_token_begin * entry_size_bytes; + const int32_t copy_bytes = + (page.page_token_end - page.page_token_begin) * entry_size_bytes; + if constexpr (vectorized) { + const int4* src = reinterpret_cast(src_cache + src_offset); + int4* output = reinterpret_cast(dst + dst_offset); + const int32_t num_vectors = copy_bytes / sizeof(int4); + int32_t idx = threadIdx.x; + for (; idx + blockDim.x < num_vectors; idx += 2 * blockDim.x) { + const int4 first = src[idx]; + const int4 second = src[idx + blockDim.x]; + output[idx] = first; + output[idx + blockDim.x] = second; + } + for (; idx < num_vectors; idx += blockDim.x) { + output[idx] = src[idx]; + } + } else { + const uint8_t* src = src_cache + src_offset; + uint8_t* output = dst + dst_offset; + for (int32_t idx = threadIdx.x; idx < copy_bytes; idx += blockDim.x) { + output[idx] = src[idx]; + } + } + } else { + const int32_t token_group = threadIdx.x / threads_per_token; + const int32_t lane_id = threadIdx.x % threads_per_token; + for (int32_t page_token = page.page_token_begin + token_group; + page_token < page.page_token_end; page_token += tokens_per_cta) { + const int32_t output_token = + page.output_token_begin + page_token - page.page_token_begin; + const uint8_t* src = src_cache + physical_block * cache_block_stride + + page_token * cache_entry_stride; + uint8_t* output = dst + output_token * dst_entry_stride; + if constexpr (vectorized) { + const int4* src_vec = reinterpret_cast(src); + int4* dst_vec = reinterpret_cast(output); + const int32_t num_vectors = entry_size_bytes / sizeof(int4); + for (int32_t idx = lane_id; idx < num_vectors; + idx += threads_per_token) { + dst_vec[idx] = src_vec[idx]; + } + } else { + for (int32_t idx = lane_id; idx < entry_size_bytes; + idx += threads_per_token) { + output[idx] = src[idx]; + } + } } } } } // namespace vllm -// Macro to dispatch the kernel based on the data type. -#define CALL_CP_GATHER_CACHE(CPY_DTYPE) \ - vllm::cp_gather_cache<<>>( \ - reinterpret_cast(src_cache.data_ptr()), \ - reinterpret_cast(dst.data_ptr()), \ - block_table.const_data_ptr(), \ - cu_seq_lens.const_data_ptr(), block_size, entry_size, \ - block_table_stride, cache_block_stride, cache_entry_stride, \ - dst_entry_stride, seq_starts_ptr); - // Gather sequences from the cache into the destination tensor. // - cu_seq_lens contains the cumulative sequence lengths for each batch // - block_table contains the cache block indices for each sequence @@ -1401,33 +1455,70 @@ void cp_gather_cache( "src_cache and seq_starts must be on the same device"); } - int64_t block_table_stride = block_table.stride(0); - int64_t cache_block_stride = src_cache.stride(0); - int64_t cache_entry_stride = src_cache.stride(1); - int64_t dst_entry_stride = dst.stride(0); - - // Decide on the number of splits based on the batch size. - int num_splits = batch_size > 128 ? 2 : batch_size > 64 ? 4 : 16; - dim3 grid(batch_size, num_splits); - dim3 block(1024); - STD_TORCH_CHECK(src_cache.scalar_type() == dst.scalar_type(), "src_cache and dst must have the same dtype"); + const int32_t element_size = src_cache.element_size(); + STD_TORCH_CHECK(element_size == 1 || element_size == 2 || element_size == 4, + "Unsupported data type width: ", element_size * 8); + STD_TORCH_CHECK(batch_size == cu_seq_lens.size(0) - 1, + "batch_size must match cu_seq_lens"); - const int dtype_bits = src_cache.element_size() * 8; + const int32_t total_tokens = dst.size(0); + if (total_tokens == 0) { + return; + } + + const int32_t entry_size_bytes = entry_size * element_size; + const int64_t block_table_stride = block_table.stride(0); + const int64_t cache_block_stride = src_cache.stride(0) * element_size; + const int64_t cache_entry_stride = src_cache.stride(1) * element_size; + const int64_t dst_entry_stride = dst.stride(0) * element_size; const int32_t* seq_starts_ptr = seq_starts.has_value() ? seq_starts.value().const_data_ptr() : nullptr; - if (dtype_bits == 32) { - CALL_CP_GATHER_CACHE(uint32_t); - } else if (dtype_bits == 16) { - CALL_CP_GATHER_CACHE(uint16_t); - } else if (dtype_bits == 8) { - CALL_CP_GATHER_CACHE(uint8_t); + constexpr std::uintptr_t vector_alignment = alignof(int4); + const bool vectorized = + reinterpret_cast(src_cache.data_ptr()) % + vector_alignment == + 0 && + reinterpret_cast(dst.data_ptr()) % vector_alignment == + 0 && + entry_size_bytes % vector_alignment == 0 && + cache_block_stride % vector_alignment == 0 && + cache_entry_stride % vector_alignment == 0 && + dst_entry_stride % vector_alignment == 0; + const bool contiguous_entries = cache_entry_stride == entry_size_bytes && + dst_entry_stride == entry_size_bytes; + + const int32_t num_reqs = static_cast(batch_size); + const int32_t required_blocks = + cuda_utils::ceil_div(total_tokens, block_size) + 2 * num_reqs; + const dim3 grid(required_blocks); + const dim3 block(256); + +#define CALL_CP_GATHER_CACHE(CONTIGUOUS, VECTORIZED) \ + vllm::cp_gather_cache_page \ + <<>>( \ + reinterpret_cast(src_cache.data_ptr()), \ + reinterpret_cast(dst.data_ptr()), \ + block_table.const_data_ptr(), \ + cu_seq_lens.const_data_ptr(), num_reqs, block_size, \ + entry_size_bytes, total_tokens, block_table_stride, \ + cache_block_stride, cache_entry_stride, dst_entry_stride, \ + seq_starts_ptr) + + if (contiguous_entries && vectorized) { + CALL_CP_GATHER_CACHE(true, true); + } else if (contiguous_entries) { + CALL_CP_GATHER_CACHE(true, false); + } else if (vectorized) { + CALL_CP_GATHER_CACHE(false, true); } else { - STD_TORCH_CHECK(false, "Unsupported data type width: ", dtype_bits); + CALL_CP_GATHER_CACHE(false, false); } + +#undef CALL_CP_GATHER_CACHE } void cp_gather_and_upconvert_fp8_kv_cache( @@ -1491,15 +1582,18 @@ void cp_gather_and_upconvert_fp8_kv_cache( } const int total_tokens = dst.size(0); - constexpr int warps_per_block = 8; - const int grid_size = (total_tokens + warps_per_block - 1) / warps_per_block; - const int block_size_threads = warps_per_block * 32; // 256 threads + if (total_tokens == 0) { + return; + } + constexpr int warps_per_block = 16; + const int block_size_threads = warps_per_block * 32; const int32_t* seq_starts_ptr = seq_starts.has_value() ? seq_starts.value().const_data_ptr() : nullptr; - - vllm::cp_gather_and_upconvert_fp8_kv_cache<<>>( + const int required_blocks = cuda_utils::ceil_div(total_tokens, block_size) + + 2 * static_cast(batch_size); + vllm::cp_gather_and_upconvert_fp8_kv_cache_page<<< + required_blocks, block_size_threads, 0, stream>>>( src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()), block_table.const_data_ptr(), workspace_starts.const_data_ptr(), diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py index c48a1d2338..b6a029ce08 100644 --- a/tests/kernels/attention/test_cache.py +++ b/tests/kernels/attention/test_cache.py @@ -1152,6 +1152,71 @@ def test_gather_and_maybe_dequant_cache_mla_with_seq_starts( torch.testing.assert_close(dst, expected) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize("use_seq_starts", [False, True]) +@pytest.mark.parametrize("device", CUDA_DEVICES) +@torch.inference_mode() +def test_gather_and_maybe_dequant_cache_mla_large_uneven_sequences( + kv_cache_dtype, + use_seq_starts, + device, +): + block_size = 64 + entry_size = 576 + num_blocks = 1024 + dtype = torch.bfloat16 + scale = torch.tensor(0.1, dtype=torch.float32, device=device) + src_cache = _create_mla_cache( + num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device + ) + _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype) + + starts = torch.tensor([3, 17, 5], dtype=torch.int32, device=device) + seq_starts = starts if use_seq_starts else None + seq_lens = torch.tensor([17, 32_768, 71], dtype=torch.int32, device=device) + batch_size = seq_lens.shape[0] + total_tokens = seq_lens.sum().item() + cu_seq_lens = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + cu_seq_lens[1:] = seq_lens.cumsum(dim=0) + token_to_seq = torch.repeat_interleave( + torch.arange(batch_size, dtype=torch.int32, device=device), seq_lens + ) + block_table = torch.stack( + [torch.randperm(num_blocks, device=device) for _ in range(batch_size)] + ).to(torch.int32) + + if kv_cache_dtype == "fp8": + dequant_src_cache = torch.empty_like(src_cache, dtype=dtype) + ops.convert_fp8(dequant_src_cache, src_cache, scale.item()) + else: + dequant_src_cache = src_cache + expected_batches = [] + for req_id in range(batch_size): + start = starts[req_id].item() if use_seq_starts else 0 + source_tokens = torch.arange( + start, start + seq_lens[req_id].item(), device=device + ) + physical_blocks = block_table[req_id, source_tokens // block_size].long() + expected_batches.append( + dequant_src_cache[physical_blocks, source_tokens % block_size] + ) + expected = torch.cat(expected_batches) + dst = torch.empty_like(expected) + + ops.gather_and_maybe_dequant_cache( + src_cache, + dst, + block_table, + cu_seq_lens, + token_to_seq, + total_tokens, + kv_cache_dtype, + scale, + seq_starts, + ) + torch.testing.assert_close(dst, expected) + + @pytest.mark.parametrize("kv_lora_rank", [512]) @pytest.mark.parametrize("qk_rope_head_dim", [64]) @pytest.mark.parametrize("block_size", [16]) @@ -1228,6 +1293,120 @@ def test_cp_gather_cache_mla( torch.testing.assert_close(dst, expected) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize("unaligned", [False, True]) +@pytest.mark.parametrize("device", CUDA_DEVICES) +@torch.inference_mode() +def test_cp_gather_cache_mla_with_seq_starts( + kv_cache_dtype, + unaligned, + device, +): + block_size = 16 + entry_size = 576 + num_blocks = 128 + src_storage = _create_mla_cache( + num_blocks, + block_size, + entry_size + int(unaligned), + torch.bfloat16, + kv_cache_dtype, + device, + ) + src_cache = src_storage[..., 1:] if unaligned else src_storage + _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype) + + seq_starts = torch.tensor([3, 17, 5, 31], dtype=torch.int32, device=device) + seq_lens = torch.tensor([0, 1, 63, 4], dtype=torch.int32, device=device) + batch_size = seq_lens.shape[0] + cu_seq_lens = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + cu_seq_lens[1:] = seq_lens.cumsum(dim=0) + block_table = torch.stack( + [torch.randperm(num_blocks, device=device) for _ in range(batch_size)] + ).to(torch.int32) + + expected_rows = [] + for req_id in range(batch_size): + start = seq_starts[req_id].item() + for source_token in range(start, start + seq_lens[req_id].item()): + block_id = block_table[req_id, source_token // block_size] + expected_rows.append(src_cache[block_id, source_token % block_size]) + expected = torch.stack(expected_rows) + if unaligned: + dst_storage = torch.empty( + (expected.shape[0], entry_size + 1), + dtype=expected.dtype, + device=device, + ) + dst = dst_storage[:, 1:] + else: + dst = torch.empty_like(expected) + + ops.cp_gather_cache( + src_cache, + dst, + block_table, + cu_seq_lens, + batch_size, + seq_starts, + ) + torch.testing.assert_close(dst, expected) + + +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize("use_seq_starts", [False, True]) +@pytest.mark.parametrize("device", CUDA_DEVICES) +@torch.inference_mode() +def test_cp_gather_cache_mla_large_uneven_sequences( + kv_cache_dtype, + use_seq_starts, + device, +): + block_size = 64 + entry_size = 576 + num_blocks = 1024 + src_cache = _create_mla_cache( + num_blocks, + block_size, + entry_size, + torch.bfloat16, + kv_cache_dtype, + device, + ) + _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype) + + starts = torch.tensor([3, 17, 5], dtype=torch.int32, device=device) + seq_starts = starts if use_seq_starts else None + seq_lens = torch.tensor([17, 32_768, 71], dtype=torch.int32, device=device) + batch_size = seq_lens.shape[0] + cu_seq_lens = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + cu_seq_lens[1:] = seq_lens.cumsum(dim=0) + block_table = torch.stack( + [torch.randperm(num_blocks, device=device) for _ in range(batch_size)] + ).to(torch.int32) + + expected_batches = [] + for req_id in range(batch_size): + start = starts[req_id].item() if use_seq_starts else 0 + source_tokens = torch.arange( + start, start + seq_lens[req_id].item(), device=device + ) + physical_blocks = block_table[req_id, source_tokens // block_size].long() + expected_batches.append(src_cache[physical_blocks, source_tokens % block_size]) + expected = torch.cat(expected_batches) + dst = torch.empty_like(expected) + + ops.cp_gather_cache( + src_cache, + dst, + block_table, + cu_seq_lens, + batch_size, + seq_starts, + ) + torch.testing.assert_close(dst, expected) + + @pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS) @pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS) @pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA) diff --git a/tests/kernels/test_cp_gather_fp8.py b/tests/kernels/test_cp_gather_fp8.py index 36e7f0c9d4..35afe33808 100644 --- a/tests/kernels/test_cp_gather_fp8.py +++ b/tests/kernels/test_cp_gather_fp8.py @@ -411,3 +411,57 @@ def test_cp_gather_fp8_large_seqlens(seq_lens, block_size): dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2 ) assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:]) + + +def test_cp_gather_fp8_large_uneven_sequences_with_starts(): + """Gather uneven long request slices through the page-oriented path.""" + gather_seq_lens = [17, 32_768, 71] + seq_starts = [3, 17, 5] + full_seq_lens = [ + start + length for start, length in zip(seq_starts, gather_seq_lens) + ] + ( + cache, + block_table, + full_workspace_starts, + num_reqs, + _total_tokens, + full_expected, + ) = _build_test_case_fast(full_seq_lens, block_size=64) + + workspace_starts = torch.zeros(num_reqs, dtype=torch.int32, device="cuda") + workspace_starts[1:] = torch.tensor( + gather_seq_lens[:-1], dtype=torch.int32, device="cuda" + ).cumsum(dim=0) + seq_starts_t = torch.tensor(seq_starts, dtype=torch.int32, device="cuda") + dst = torch.empty( + sum(gather_seq_lens), + NOPE_DIM + ROPE_DIM, + dtype=torch.bfloat16, + device="cuda", + ) + + ops.cp_gather_and_upconvert_fp8_kv_cache( + cache, + dst, + block_table, + workspace_starts, + num_reqs, + seq_starts_t, + ) + + expected = torch.cat( + [ + full_expected[ + full_workspace_starts[req_id] + + seq_starts[req_id] : full_workspace_starts[req_id] + + seq_starts[req_id] + + gather_seq_lens[req_id] + ] + for req_id in range(num_reqs) + ] + ) + torch.testing.assert_close( + dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2 + ) + assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:]) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index b21f967732..1e7647d163 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -2574,7 +2574,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]): # FP8 path: gather cache without dequantization ops.cp_gather_cache( src_cache=kv_c_and_k_pe_cache, - dst=workspace, + dst=workspace[:toks], block_table=block_table, cu_seq_lens=chunk.cu_seq_lens, batch_size=chunk.num_requests, @@ -2683,7 +2683,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]): else: ops.cp_gather_cache( src_cache=kv_c_and_k_pe_cache, - dst=workspace, + dst=workspace[:toks], block_table=block_table, cu_seq_lens=padded_local_cu_seq_lens, batch_size=chunk.num_requests,