diff --git a/CMakeLists.txt b/CMakeLists.txt index a7d40ae586..018138f54c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -476,6 +476,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_STABLE_EXT_SRC + "csrc/libtorch_stable/hisparse_kernels.cu" "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu" "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu" "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu") diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu index f42d61a4f4..887c83f3af 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -1458,7 +1458,11 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache_page( 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) { + const int32_t* __restrict__ seq_starts, + const uint8_t* __restrict__ host_cache, + const int32_t* __restrict__ host_row_ids, + const int32_t* __restrict__ device_row_ids, + const int64_t host_entry_stride) { constexpr int32_t warps_per_cta = 16; __shared__ GatherPageTask page; __shared__ int32_t physical_block; @@ -1484,8 +1488,22 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache_page( 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; + const int32_t source_row = physical_block * block_size + page_token; + const uint8_t* token_ptr; + if (device_row_ids == nullptr) { + token_ptr = src_cache + physical_block * cache_block_stride + + page_token * cache_entry_stride; + } else { + const int32_t device_row = device_row_ids[source_row]; + if (device_row >= 0) { + token_ptr = src_cache + (device_row / block_size) * cache_block_stride + + (device_row % block_size) * cache_entry_stride; + } else { + token_ptr = + host_cache + + static_cast(host_row_ids[source_row]) * host_entry_stride; + } + } __nv_bfloat16* dst_ptr = dst + output_token * dst_entry_stride; gather_and_upconvert_fp8_token(token_ptr, dst_ptr, lane_id); } @@ -1694,7 +1712,10 @@ void cp_gather_and_upconvert_fp8_kv_cache( torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] torch::stable::Tensor const& workspace_starts, // [BATCH] int64_t batch_size, - std::optional seq_starts = std::nullopt) { + std::optional seq_starts = std::nullopt, + std::optional host_cache = std::nullopt, + std::optional host_row_ids = std::nullopt, + std::optional device_row_ids = std::nullopt) { torch::stable::accelerator::DeviceGuard device_guard( src_cache.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(); @@ -1714,6 +1735,12 @@ void cp_gather_and_upconvert_fp8_kv_cache( "seq_starts must be int32"); } + const bool has_host_rows = host_cache.has_value(); + STD_TORCH_CHECK( + has_host_rows == host_row_ids.has_value() && + has_host_rows == device_row_ids.has_value(), + "host_cache, host_row_ids, and device_row_ids must be provided together"); + STD_TORCH_CHECK(src_cache.device() == dst.device(), "src_cache and dst must be on the same device"); STD_TORCH_CHECK(src_cache.device() == block_table.device(), @@ -1724,6 +1751,43 @@ void cp_gather_and_upconvert_fp8_kv_cache( STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(), "src_cache and seq_starts must be on the same device"); } + const uint8_t* host_cache_ptr = nullptr; + const int32_t* host_row_ids_ptr = nullptr; + const int32_t* device_row_ids_ptr = nullptr; + int64_t host_entry_stride = 0; + if (has_host_rows) { + auto const& host = host_cache.value(); + auto const& host_rows = host_row_ids.value(); + auto const& device_rows = device_row_ids.value(); + STD_TORCH_CHECK( + host.device().is_cpu() && + host.scalar_type() == torch::headeronly::ScalarType::Byte && + host.dim() == 2 && host.is_contiguous() && + host.size(1) == src_cache.size(2), + "host_cache must be a contiguous uint8 CPU row matrix matching " + "src_cache"); + cudaPointerAttributes attributes{}; + const cudaError_t pointer_status = + cudaPointerGetAttributes(&attributes, host.const_data_ptr()); + if (pointer_status != cudaSuccess) { + cudaGetLastError(); + } + STD_TORCH_CHECK( + pointer_status == cudaSuccess && attributes.type == cudaMemoryTypeHost, + "host_cache must be CUDA-accessible pinned memory"); + STD_TORCH_CHECK( + host_rows.is_cuda() && device_rows.is_cuda() && + host_rows.scalar_type() == torch::headeronly::ScalarType::Int && + device_rows.scalar_type() == torch::headeronly::ScalarType::Int && + host_rows.is_contiguous() && device_rows.is_contiguous() && + host_rows.numel() == device_rows.numel() && host_rows.dim() <= 2 && + device_rows.dim() <= 2, + "host/device row maps must be matching contiguous CUDA int32 tensors"); + host_cache_ptr = reinterpret_cast(host.const_data_ptr()); + host_row_ids_ptr = host_rows.const_data_ptr(); + device_row_ids_ptr = device_rows.const_data_ptr(); + host_entry_stride = host.stride(0); + } auto dtype = src_cache.scalar_type(); STD_TORCH_CHECK( dtype == torch::headeronly::ScalarType::Byte || // uint8 @@ -1766,7 +1830,8 @@ void cp_gather_and_upconvert_fp8_kv_cache( workspace_starts.const_data_ptr(), static_cast(batch_size), block_size, total_tokens, block_table_stride, cache_block_stride, cache_entry_stride, - dst_entry_stride, seq_starts_ptr); + dst_entry_stride, seq_starts_ptr, host_cache_ptr, host_row_ids_ptr, + device_row_ids_ptr, host_entry_stride); } void cp_gather_and_upconvert_nvfp4_kv_cache( diff --git a/csrc/libtorch_stable/hisparse_kernels.cu b/csrc/libtorch_stable/hisparse_kernels.cu new file mode 100644 index 0000000000..63d8457627 --- /dev/null +++ b/csrc/libtorch_stable/hisparse_kernels.cu @@ -0,0 +1,1226 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// HiSparse decode hot-buffer kernels for sparse MLA (GLM-5 / DeepSeek-V3.2). +// +// The swap-in algorithm is a port of SGLang's hisparse +// load_cache_to_device_buffer kernel (sgl jit_kernel/csrc/hisparse.cuh), +// adapted to vLLM addressing: +// - tokens are keyed by their global KV slot id (block_table-converted +// indexer output) instead of in-request positions, so no per-request +// host-location table is needed: host pool row i is global slot i. +// - each batch row owns a fixed region of `region_stride` hot rows; +// slots [0, hot_size) are LRU-managed, slot `hot_size` holds the row's +// newest token (written directly by the KV-cache update). +// +// The kernels only move bytes; they are dtype agnostic. + +#include "torch_utils.h" +#include "ops.h" +#include "../cuda_utils.h" + +#include +#include +#include + +#include +#include +#include + +namespace { + +constexpr int kWarpSize = 32; +// Sentinel in the shared top-k scratch: entry already resolved (hit / +// newest / invalid), no miss handling needed. +constexpr int32_t kTokenDone = -1; +constexpr int32_t kHashEmpty = -1; + +bool is_pinned_cpu_tensor(const torch::stable::Tensor& tensor) { + cudaPointerAttributes attributes{}; + const auto status = + cudaPointerGetAttributes(&attributes, tensor.const_data_ptr()); + if (status != cudaSuccess) { + cudaGetLastError(); + return false; + } + return attributes.type == cudaMemoryTypeHost; +} + +__device__ __forceinline__ int32_t hash_slot(int32_t key, int32_t hash_size) { + // Knuth multiplicative hash for the open-addressing table. + return static_cast((static_cast(key) * 2654435761u) % + static_cast(hash_size)); +} + +// Copy one row of `row_bytes` bytes with a single warp. MLA rows take the +// 16-byte vectorized path; packed indexer rows use the scalar tail-safe path. +__device__ __forceinline__ void copy_row_warp(int lane_id, const char* src, + char* dst, int64_t row_bytes) { + const auto alignment = reinterpret_cast(src) | + reinterpret_cast(dst) | + static_cast(row_bytes); + if ((alignment & 15) == 0) { + const int64_t num_vec = row_bytes / 16; + const uint4* src4 = reinterpret_cast(src); + uint4* dst4 = reinterpret_cast(dst); + for (int64_t j = lane_id; j < num_vec; j += kWarpSize) { + __stcg(dst4 + j, __ldcg(src4 + j)); + } + return; + } + + if ((alignment & 3) == 0) { + const int64_t num_words = row_bytes / 4; + const unsigned int* src_words = reinterpret_cast(src); + unsigned int* dst_words = reinterpret_cast(dst); + for (int64_t j = lane_id; j < num_words; j += kWarpSize) { + __stcg(dst_words + j, __ldcg(src_words + j)); + } + return; + } + + for (int64_t j = lane_id; j < row_bytes; j += kWarpSize) { + __stcg(dst + j, __ldcg(src + j)); + } +} + +// Zero one row with a single warp (16B vectorized, L2-only stores), same +// addressing contract as copy_row_warp. +__device__ __forceinline__ void zero_row_warp(int lane_id, char* dst, + int64_t row_bytes) { + const auto alignment = + reinterpret_cast(dst) | static_cast(row_bytes); + if ((alignment & 15) == 0) { + const int64_t num_vec = row_bytes / 16; + uint64_t* dst8 = reinterpret_cast(dst); + for (int64_t j = lane_id; j < num_vec; j += kWarpSize) { + uint64_t* d = dst8 + j * 2; + asm volatile("st.global.cg.v2.b64 [%0],{%1,%2};" ::"l"(d), "l"(0ULL), + "l"(0ULL) + : "memory"); + } + return; + } + + if ((alignment & 3) == 0) { + const int64_t num_words = row_bytes / 4; + unsigned int* dst_words = reinterpret_cast(dst); + for (int64_t j = lane_id; j < num_words; j += kWarpSize) { + __stcg(dst_words + j, 0u); + } + return; + } + + for (int64_t j = lane_id; j < row_bytes; j += kWarpSize) { + __stcg(dst + j, static_cast(0)); + } +} + +__device__ __forceinline__ char* cache_row_ptr(char* cache, int64_t row, + int32_t block_size, + int64_t block_stride, + int64_t row_bytes) { + return cache + (row / block_size) * block_stride + + (row % block_size) * row_bytes; +} + +__device__ __forceinline__ void zero_cache_row_warp(int lane_id, char* cache, + int64_t row, + int32_t block_size, + int64_t block_stride, + int64_t row_bytes) { + zero_row_warp(lane_id, + cache_row_ptr(cache, row, block_size, block_stride, row_bytes), + row_bytes); +} + +// In-place inclusive scan over s_data[offset, count) performed by warp 0, +// carrying `accumulator` across calls. Returns the running total. +__device__ __forceinline__ int warp_inclusive_scan(int32_t* s_data, int lane_id, + int offset, int count, + int accumulator) { + int idx = lane_id + offset; + int val = (idx < count) ? s_data[idx] : 0; +#pragma unroll + for (int i = 1; i < 32; i *= 2) { + int n = __shfl_up_sync(0xffffffff, val, i); + if (lane_id >= i) val += n; + } + val += accumulator; + if (idx < count) { + s_data[idx] = val; + } + return __shfl_sync(0xffffffff, val, 31); +} + +__device__ __forceinline__ int64_t +get_physical_hot_row(const int32_t* hot_block_table, int32_t row, + int64_t table_stride, int32_t block_size, int32_t slot) { + const int32_t block = + hot_block_table[static_cast(row) * table_stride + + slot / block_size]; + return static_cast(block) * block_size + slot % block_size; +} + +__device__ __forceinline__ void store_hot_index( + int32_t* hot_indices, int32_t* attention_indices, int32_t index, + int32_t physical_row, int32_t hot_block_size, + int64_t attention_block_stride) { + hot_indices[index] = physical_row; + if (attention_indices != nullptr) { + attention_indices[index] = + physical_row < 0 + ? -1 + : (physical_row / hot_block_size) * attention_block_stride + + physical_row % hot_block_size; + } +} + +// One block per batch row. +// +// Shared memory layout (int32 region followed by int16 region): +// s_topk[top_k] top-k global ids; reused as miss scratch +// s_chunk_off[nbc + 1] prefix sums for hit (then miss) compaction +// s_evict_off[nbc + 1] prefix sums for evictable compaction +// s_hash_keys[hash_size] open addressing: global id -> top-k index +// s_counters[3] hits, phase-1 resolved, valid count +// s_lru_out[hot_size] int16, compacted slots: [hits fwd | evict bwd] +// s_hash_vals[hash_size] int16 hash values (top-k index) +// Valid global ids must be unique within each row. +__global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( + const int32_t* __restrict__ hot_block_table, // [max_rows, hot_blocks] + const int32_t* __restrict__ global_indices, // global or request-relative + const int32_t* __restrict__ request_ids, // [num_rows] or nullptr + const int32_t* __restrict__ source_block_table, // [num_reqs, max_blocks] + const int32_t* __restrict__ resident_block_table, + int32_t* __restrict__ resolved_global_indices, // [num_rows, top_k] + int32_t* __restrict__ valid_counts, // [num_rows] or nullptr + int32_t* __restrict__ swap_host_physical_rows, // [num_rows, top_k] + int32_t* __restrict__ swap_device_physical_rows, // [num_rows, top_k] + int32_t* __restrict__ swap_counts, // [num_rows] + int32_t* __restrict__ hot_indices, // [num_rows, top_k] + int32_t* __restrict__ attention_indices, // [num_rows, top_k] + int32_t* __restrict__ miss_mask, // [num_rows, top_k] or nullptr + int32_t* __restrict__ device_global_indices, // [max_rows, region_stride] + int16_t* __restrict__ lru_slots, // [max_rows, hot_size] + const int32_t* __restrict__ request_state_indices, // [num_requests] or + // nullptr + const int32_t request_state_count, const int64_t host_rows, + const int64_t hot_table_stride, const int32_t hot_block_size, + const int32_t top_k, const int32_t hot_size, const int32_t hash_size, + const int64_t region_stride, const int64_t attention_block_stride, + const int64_t source_bt_stride, const int32_t source_num_reqs, + const int32_t source_num_blocks, const int32_t source_block_size, + const int64_t resident_bt_stride, const int32_t resident_num_reqs, + const int32_t resident_num_blocks, const int32_t resident_block_size, + const int32_t resident_null_block, const int64_t input_row_stride, + const int64_t attention_row_stride, const int64_t valid_count_stride) { + const int NUM_WARPS = blockDim.x / kWarpSize; + const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize; + const int num_token_chunks = (top_k + kWarpSize - 1) / kWarpSize; + + const int batch_row = blockIdx.x; + const int request_row = + request_ids != nullptr ? request_ids[batch_row] : batch_row; + const int state_row = + request_state_indices != nullptr && request_row >= 0 && + request_row < request_state_count + ? request_state_indices[request_row] + : (request_state_indices == nullptr ? request_row : -1); + // V2 publishes -1 for CUDA-graph padding rows. + if (state_row < 0) { + for (int i = threadIdx.x; i < top_k; i += blockDim.x) { + const int64_t index = static_cast(batch_row) * top_k + i; + hot_indices[index] = -1; + if (attention_indices != nullptr) { + attention_indices[static_cast(batch_row) * + attention_row_stride + + i] = -1; + } + if (resolved_global_indices != nullptr) { + resolved_global_indices[index] = -1; + } + } + if (valid_counts != nullptr && threadIdx.x == 0) { + valid_counts[static_cast(batch_row) * valid_count_stride] = 0; + } + if (swap_counts != nullptr && threadIdx.x == 0) { + swap_counts[batch_row] = 0; + } + return; + } + const int tid = threadIdx.x; + const int warp_id = tid / kWarpSize; + const int lane_id = tid % kWarpSize; + const unsigned int lanes_before = ((unsigned int)1 << lane_id) - 1; + + const int32_t* row_topk = + global_indices + static_cast(batch_row) * input_row_stride; + int32_t* row_out = hot_indices + static_cast(batch_row) * top_k; + int32_t* row_attention = + attention_indices != nullptr + ? attention_indices + + static_cast(batch_row) * attention_row_stride + : nullptr; + int32_t* row_miss = (miss_mask != nullptr) + ? miss_mask + static_cast(batch_row) * top_k + : nullptr; + int32_t* row_dgi = + device_global_indices + static_cast(state_row) * region_stride; + int16_t* row_lru = lru_slots + static_cast(state_row) * hot_size; + + extern __shared__ char smem_raw[]; + int32_t* s_topk = reinterpret_cast(smem_raw); + int32_t* s_chunk_off = s_topk + top_k; + int32_t* s_evict_off = s_chunk_off + (num_buffer_chunks + 1); + int32_t* s_hash_keys = s_evict_off + (num_buffer_chunks + 1); + int32_t* s_counters = s_hash_keys + hash_size; + int16_t* s_lru_out = reinterpret_cast(s_counters + 3); + int16_t* s_hash_vals = s_lru_out + hot_size; + + if (tid < 3) { + s_counters[tid] = 0; + } + __syncthreads(); + + // Phase 1: translate request-relative positions and resolve resident rows. + for (int i = tid; i < top_k; i += blockDim.x) { + const int32_t token_index = row_topk[i]; + int32_t g = token_index; + int32_t resident_row = -1; + if (source_block_table != nullptr) { + const int32_t request_id = request_ids[batch_row]; + const int32_t source_block = + token_index >= 0 ? token_index / source_block_size : -1; + if (request_id >= 0 && request_id < source_num_reqs && + source_block >= 0 && source_block < source_num_blocks) { + const int32_t physical_block = + source_block_table[static_cast(request_id) * + source_bt_stride + + source_block]; + g = physical_block > 0 ? physical_block * source_block_size + + token_index % source_block_size + : -1; + } else { + g = -1; + } + if (resident_block_table != nullptr) { + const int32_t resident_block = + token_index >= 0 ? token_index / resident_block_size : -1; + if (request_id >= 0 && request_id < resident_num_reqs && + resident_block >= 0 && resident_block < resident_num_blocks) { + const int32_t physical_block = + resident_block_table[static_cast(request_id) * + resident_bt_stride + + resident_block]; + if (physical_block != resident_null_block && physical_block >= 0) { + resident_row = physical_block * resident_block_size + + token_index % resident_block_size; + } + } + } + } + if (g >= host_rows) { + g = -1; + } + if (resolved_global_indices != nullptr) { + resolved_global_indices[static_cast(batch_row) * top_k + i] = g; + } + if (resident_row >= 0 || g >= 0) atomicAdd(&s_counters[2], 1); + if (row_miss != nullptr) row_miss[i] = 0; + if (resident_row >= 0) { + store_hot_index(row_out, row_attention, i, resident_row, hot_block_size, + attention_block_stride); + s_topk[i] = kTokenDone; + atomicAdd(&s_counters[1], 1); + } else if (g < 0) { + store_hot_index(row_out, row_attention, i, -1, hot_block_size, + attention_block_stride); + s_topk[i] = kTokenDone; + atomicAdd(&s_counters[1], 1); + } else { + s_topk[i] = g; + } + } + __syncthreads(); + if (valid_counts != nullptr && tid == 0) { + valid_counts[static_cast(batch_row) * valid_count_stride] = + s_counters[2]; + } + // Fully resident rows need only request-relative page translation. Avoid + // scanning or rewriting the hot LRU when no selected row can consult it. + if (s_counters[1] == top_k) { + if (tid == 0 && swap_counts != nullptr) { + swap_counts[batch_row] = 0; + } + return; + } + + for (int i = tid; i < hash_size; i += blockDim.x) { + s_hash_keys[i] = kHashEmpty; + } + for (int i = tid; i < num_buffer_chunks + 1; i += blockDim.x) { + s_chunk_off[i] = 0; + s_evict_off[i] = 0; + } + __syncthreads(); + + for (int i = tid; i < top_k; i += blockDim.x) { + const int32_t g = s_topk[i]; + if (g == kTokenDone) { + continue; + } + int slot = hash_slot(g, hash_size); + while (true) { + const int32_t old = atomicCAS(&s_hash_keys[slot], kHashEmpty, g); + if (old == kHashEmpty || old == g) { + s_hash_vals[slot] = static_cast(i); + break; + } + slot = (slot + 1) % hash_size; + } + } + __syncthreads(); + + // Phase 2: walk hot slots in LRU order, classify hit / evictable, and + // compact them (hits forward, evictables backward) into s_lru_out. + const int iters_buffer = (num_buffer_chunks + NUM_WARPS - 1) / NUM_WARPS; + int total_hit_count = 0; + int total_evict_count = 0; + for (int iter = 0; iter < iters_buffer; iter++) { + const int chunk_idx = warp_id + iter * NUM_WARPS; + const bool has_valid_chunk = chunk_idx < num_buffer_chunks; + + const int pos = chunk_idx * kWarpSize + lane_id; + const bool has_valid_pos = has_valid_chunk && (pos < hot_size); + const int16_t slot = has_valid_pos ? row_lru[pos] : int16_t(-1); + // Corruption tripwire: lru/dgi are long-lived device state; if some + // external writer (e.g. stray RDMA into reused VRAM) corrupts a slot + // out of [0, region_stride), treat it as no-hit so it degrades to a re-miss + // instead of an unbounded dgi read (phases 3/5 bound the write side). + const int32_t cached_g = + (slot >= 0 && slot < region_stride) ? row_dgi[slot] : -1; + + int found_topk_idx = -1; + if (cached_g >= 0) { + int h = hash_slot(cached_g, hash_size); + while (true) { + const int32_t k = s_hash_keys[h]; + if (k == cached_g) { + found_topk_idx = static_cast(s_hash_vals[h]); + break; + } + if (k == kHashEmpty) break; + h = (h + 1) % hash_size; + } + } + const bool is_hit = found_topk_idx >= 0; + const bool is_evictable = has_valid_pos && !is_hit; + + if (is_hit) { + s_topk[found_topk_idx] = kTokenDone; + store_hot_index(row_out, row_attention, found_topk_idx, + static_cast(get_physical_hot_row( + hot_block_table, request_row, hot_table_stride, + hot_block_size, slot)), + hot_block_size, attention_block_stride); + } + + int local_hit_off = 0; + int local_evict_off = 0; + if (has_valid_chunk) { + const unsigned int hit_mask = __ballot_sync(0xFFFFFFFF, is_hit); + const unsigned int evict_mask = __ballot_sync(0xFFFFFFFF, is_evictable); + local_hit_off = __popc(hit_mask & lanes_before); + local_evict_off = __popc(evict_mask & lanes_before); + if (lane_id == 0) { + s_chunk_off[chunk_idx + 1] = __popc(hit_mask); + s_evict_off[chunk_idx + 1] = __popc(evict_mask); + } + } + __syncthreads(); + + if (warp_id == 0) { + total_hit_count = + warp_inclusive_scan(s_chunk_off, lane_id, chunk_idx + 1, + num_buffer_chunks + 1, total_hit_count); + total_evict_count = + warp_inclusive_scan(s_evict_off, lane_id, chunk_idx + 1, + num_buffer_chunks + 1, total_evict_count); + if (tid == 0) { + s_counters[0] = total_hit_count; + } + } + __syncthreads(); + + if (is_hit) { + const int off = s_chunk_off[chunk_idx] + local_hit_off; + s_lru_out[off] = slot; + } + if (is_evictable) { + const int off = s_evict_off[chunk_idx] + local_evict_off; + s_lru_out[hot_size - 1 - off] = slot; + } + } + __syncthreads(); + + // Reset prefix sums for the miss compaction (token chunks <= buffer + // chunks because hot_size >= top_k). + for (int i = tid; i < num_token_chunks + 1; i += blockDim.x) { + s_chunk_off[i] = 0; + } + __syncthreads(); + + // Phase 3: compact misses, assign them eviction slots (oldest first) and + // record the new ownership in device_global_indices. + const int iters_token = (num_token_chunks + NUM_WARPS - 1) / NUM_WARPS; + int miss_running_total = 0; + for (int iter = 0; iter < iters_token; iter++) { + const int chunk_idx = warp_id + iter * NUM_WARPS; + const bool has_valid_chunk = chunk_idx < num_token_chunks; + + const int i = chunk_idx * kWarpSize + lane_id; + const bool has_valid_token = has_valid_chunk && (i < top_k); + + int32_t g = 0; + bool is_miss = false; + if (has_valid_token) { + is_miss = s_topk[i] != kTokenDone; + if (is_miss) { + g = s_topk[i]; + } + } + + int local_miss_off = 0; + if (has_valid_chunk) { + const unsigned int miss_mask = __ballot_sync(0xFFFFFFFF, is_miss); + local_miss_off = __popc(miss_mask & lanes_before); + if (lane_id == 0) { + s_chunk_off[chunk_idx + 1] = __popc(miss_mask); + } + } + __syncthreads(); + + if (warp_id == 0) { + miss_running_total = + warp_inclusive_scan(s_chunk_off, lane_id, chunk_idx + 1, + num_token_chunks + 1, miss_running_total); + } + __syncthreads(); + + if (is_miss) { + const int m = s_chunk_off[chunk_idx] + local_miss_off; + const int16_t evict_slot = s_lru_out[hot_size - 1 - m]; + if (evict_slot < 0 || evict_slot >= region_stride) { + // Corruption tripwire: an out-of-range slot (corrupted lru state, + // see phase 2) must not become a hot-cache/dgi write. Resolve the + // entry as invalid (-1, masked by attention, not in the miss set); + // it re-misses on a later step. Phase 5 re-checks the same + // s_lru_out value, so its copy is skipped consistently. + store_hot_index(row_out, row_attention, i, -1, hot_block_size, + attention_block_stride); + if (swap_host_physical_rows != nullptr) { + const int64_t compact_index = + static_cast(batch_row) * top_k + m; + swap_host_physical_rows[compact_index] = g; + swap_device_physical_rows[compact_index] = -1; + } + } else { + // Reuse s_topk as compacted miss scratch: m < i always holds (done + // entries are skipped), so writes never overrun pending reads. + s_topk[m] = g; + const int32_t physical_row = static_cast( + get_physical_hot_row(hot_block_table, request_row, hot_table_stride, + hot_block_size, evict_slot)); + store_hot_index(row_out, row_attention, i, physical_row, hot_block_size, + attention_block_stride); + if (swap_host_physical_rows != nullptr) { + const int64_t compact_index = + static_cast(batch_row) * top_k + m; + swap_host_physical_rows[compact_index] = g; + swap_device_physical_rows[compact_index] = physical_row; + } + if (row_miss != nullptr) row_miss[i] = 1; + row_dgi[evict_slot] = g; + } + } + } + __syncthreads(); + + const int total_hits = s_counters[0]; + const int total_misses = top_k - total_hits - s_counters[1]; + if (swap_counts != nullptr && tid == 0) { + swap_counts[batch_row] = total_misses; + } + // Phase 4: write back the LRU order: stale evictables at the front, + // freshly loaded misses next, then hits at MRU. + const int total_evictable = hot_size - total_hits; + const int remaining_evictable = total_evictable - total_misses; + for (int i = tid; i < hot_size; i += blockDim.x) { + if (i < remaining_evictable) { + row_lru[i] = s_lru_out[hot_size - 1 - total_misses - i]; + } else if (i < remaining_evictable + total_misses) { + row_lru[i] = s_lru_out[hot_size - 1 - (i - remaining_evictable)]; + } else { + row_lru[i] = s_lru_out[i - remaining_evictable - total_misses]; + } + } +} + +// Shared-layer plan replay. Given a plan (hot_indices + miss_mask) already +// computed by the group's index-producing layer, gather THIS +// layer's own missed KV rows into the planned hot slots. No LRU resolution: +// index-sharing shared layers see identical global_indices/hot_indices, so the +// slot assignment is identical -- only the per-layer bytes differ. Fixed shape +// (num_rows x top_k), so it is CUDA-graph-capture safe. +__global__ void hisparse_gather_plan_kernel( + const char* __restrict__ host_cache, // [host_rows, row_bytes] + char* __restrict__ hot_cache, // [hot_rows, row_bytes] + const int32_t* __restrict__ global_indices, // [num_rows, top_k] + const int32_t* __restrict__ hot_indices, // [num_rows, top_k] abs hot rows + const int32_t* __restrict__ miss_mask, // [num_rows, top_k] + int32_t* __restrict__ attention_indices, // [num_rows, top_k] + const int32_t* __restrict__ request_state_indices, // [num_rows] or nullptr + const int64_t host_rows, const int64_t hot_rows, const int64_t row_bytes, + const int64_t hot_block_stride, const int32_t hot_block_size, + const int32_t top_k, const int64_t attention_block_stride) { + const int NUM_WARPS = blockDim.x / kWarpSize; + // Columns are interleaved across gridDim.y blocks per row so few-row + // launches (local-prefill staging's single-row layout, small decode + // batches) still fill the device with outstanding host reads instead of + // starving on one block per row. + const int row = blockIdx.x; + const bool is_padding = + request_state_indices != nullptr && request_state_indices[row] < 0; + const int warp_id = threadIdx.x / kWarpSize; + const int lane_id = threadIdx.x % kWarpSize; + const int col_start = blockIdx.y * NUM_WARPS + warp_id; + const int col_stride = gridDim.y * NUM_WARPS; + const int64_t base = static_cast(row) * top_k; + for (int col = col_start; col < top_k; col += col_stride) { + const int32_t dst = hot_indices[base + col]; + if (lane_id == 0 && attention_indices != nullptr) { + attention_indices[base + col] = + is_padding + ? -1 + : (dst < 0 ? -1 + : (dst / hot_block_size) * attention_block_stride + + dst % hot_block_size); + } + if (is_padding) { + continue; + } + if (miss_mask[base + col] == 0) { + continue; + } + const int32_t g = global_indices[base + col]; + if (g < 0 || dst < 0 || dst >= hot_rows) { + continue; + } + if (g < host_rows) { + copy_row_warp(lane_id, host_cache + static_cast(g) * row_bytes, + cache_row_ptr(hot_cache, dst, hot_block_size, + hot_block_stride, row_bytes), + row_bytes); + } else { + // No source row for g: zero the planned slot rather than serving + // whatever bytes it held (see the swap-in kernel's phase 5). + zero_cache_row_warp(lane_id, hot_cache, dst, hot_block_size, + hot_block_stride, row_bytes); + } + } +} + +// Copy the compact swap rows produced by the residency resolver. +__global__ void hisparse_gather_compact_kernel( + const char* __restrict__ host_cache, char* __restrict__ hot_cache, + const int32_t* __restrict__ miss_global_indices, + const int32_t* __restrict__ miss_hot_indices, + const int32_t* __restrict__ miss_counts, const int64_t host_rows, + const int64_t hot_rows, const int64_t row_bytes, + const int64_t hot_block_stride, const int32_t hot_block_size, + const int32_t top_k) { + const int NUM_WARPS = blockDim.x / kWarpSize; + const int row = blockIdx.x; + const int warp_id = threadIdx.x / kWarpSize; + const int lane_id = threadIdx.x % kWarpSize; + const int miss_count = min(max(miss_counts[row], 0), top_k); + const int col_start = blockIdx.y * NUM_WARPS + warp_id; + const int col_stride = gridDim.y * NUM_WARPS; + const int64_t base = static_cast(row) * top_k; + for (int col = col_start; col < miss_count; col += col_stride) { + const int32_t g = miss_global_indices[base + col]; + const int32_t dst = miss_hot_indices[base + col]; + if (g < 0 || dst < 0 || dst >= hot_rows) { + continue; + } + if (g < host_rows) { + copy_row_warp(lane_id, host_cache + static_cast(g) * row_bytes, + cache_row_ptr(hot_cache, dst, hot_block_size, + hot_block_stride, row_bytes), + row_bytes); + } else { + zero_cache_row_warp(lane_id, hot_cache, dst, hot_block_size, + hot_block_stride, row_bytes); + } + } +} + +__global__ void hisparse_invalidate_written_slots_kernel( + int32_t* __restrict__ device_global_indices, + const int32_t* __restrict__ request_state_indices, + const int32_t* __restrict__ req_id_per_token, + const int64_t* __restrict__ written_slots, const int64_t num_tokens, + const int64_t num_request_ids, const int64_t num_state_rows, + const int64_t region_stride) { + const int64_t token_idx = blockIdx.x; + if (token_idx >= num_tokens) { + return; + } + const int32_t req_idx = req_id_per_token[token_idx]; + if (req_idx < 0 || req_idx >= num_request_ids) { + return; + } + const int32_t state_idx = request_state_indices[req_idx]; + const int64_t written_slot = written_slots[token_idx]; + if (state_idx < 0 || state_idx >= num_state_rows || written_slot < 0) { + return; + } + int32_t* row = device_global_indices + state_idx * region_stride; + for (int64_t offset = threadIdx.x; offset < region_stride; + offset += blockDim.x) { + if (row[offset] == written_slot) { + row[offset] = -1; + } + } +} + +int64_t check_2d_rows(const torch::stable::Tensor& t, const char* name, + int64_t row_bytes) { + STD_TORCH_CHECK(t.dim() == 2, name, " must be 2D"); + STD_TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); + STD_TORCH_CHECK(t.size(1) * t.element_size() == row_bytes, name, + " row width mismatch"); + return t.size(0); +} + +} // namespace + +void hisparse_invalidate_written_slots( + torch::stable::Tensor& device_global_indices, + torch::stable::Tensor const& request_state_indices, + torch::stable::Tensor const& req_id_per_token, + torch::stable::Tensor const& written_slots) { + STD_TORCH_CHECK(device_global_indices.is_cuda() && + request_state_indices.is_cuda() && + req_id_per_token.is_cuda() && written_slots.is_cuda(), + "HiSparse invalidation tensors must be on CUDA"); + STD_TORCH_CHECK( + device_global_indices.scalar_type() == + torch::headeronly::ScalarType::Int && + device_global_indices.dim() == 2 && + device_global_indices.is_contiguous(), + "device_global_indices must be a contiguous 2D int32 CUDA tensor"); + STD_TORCH_CHECK( + request_state_indices.scalar_type() == + torch::headeronly::ScalarType::Int && + request_state_indices.dim() == 1 && + request_state_indices.is_contiguous(), + "request_state_indices must be a contiguous 1D int32 CUDA tensor"); + STD_TORCH_CHECK( + req_id_per_token.scalar_type() == torch::headeronly::ScalarType::Int && + req_id_per_token.dim() == 1 && req_id_per_token.is_contiguous(), + "req_id_per_token must be a contiguous 1D int32 CUDA tensor"); + STD_TORCH_CHECK( + written_slots.scalar_type() == torch::headeronly::ScalarType::Long && + written_slots.dim() == 1 && written_slots.is_contiguous(), + "written_slots must be a contiguous 1D int64 CUDA tensor"); + STD_TORCH_CHECK(req_id_per_token.numel() == written_slots.numel(), + "req_id_per_token and written_slots must have equal length"); + const int device_index = device_global_indices.get_device_index(); + STD_TORCH_CHECK( + request_state_indices.get_device_index() == device_index && + req_id_per_token.get_device_index() == device_index && + written_slots.get_device_index() == device_index, + "HiSparse invalidation tensors must be on the same CUDA device"); + + const int64_t num_tokens = written_slots.numel(); + const int64_t num_state_rows = device_global_indices.size(0); + const int64_t region_stride = device_global_indices.size(1); + if (num_tokens == 0 || num_state_rows == 0 || region_stride == 0) { + return; + } + STD_TORCH_CHECK(num_tokens <= INT32_MAX, + "HiSparse invalidation grid exceeds CUDA limits"); + + constexpr int kBlockSize = 256; + const torch::stable::accelerator::DeviceGuard device_guard(device_index); + const cudaStream_t stream = get_current_cuda_stream(); + hisparse_invalidate_written_slots_kernel<<(num_tokens), + kBlockSize, 0, stream>>>( + device_global_indices.mutable_data_ptr(), + request_state_indices.const_data_ptr(), + req_id_per_token.const_data_ptr(), + written_slots.const_data_ptr(), num_tokens, + request_state_indices.numel(), num_state_rows, region_stride); + const cudaError_t launch_error = cudaGetLastError(); + STD_TORCH_CHECK(launch_error == cudaSuccess, + "HiSparse invalidation kernel launch failed: ", + cudaGetErrorString(launch_error)); +} + +void hisparse_resolve_residency( + torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& hot_block_table, + torch::stable::Tensor const& global_indices, + torch::stable::Tensor& hot_indices, + torch::stable::Tensor& device_global_indices, + torch::stable::Tensor& lru_slots, + std::optional const& request_state_indices, + int64_t region_stride, + std::optional const& miss_mask, + std::optional const& attention_indices, + int64_t attention_block_stride, + std::optional const& request_ids, + std::optional const& source_block_table, + int64_t source_block_size, + std::optional const& resolved_global_indices, + std::optional const& valid_counts, + std::optional const& swap_host_physical_rows, + std::optional const& swap_device_physical_rows, + std::optional const& swap_counts, + std::optional const& resident_block_table, + int64_t resident_block_size, int64_t resident_null_block) { + STD_TORCH_CHECK( + host_cache.device().is_cpu() && is_pinned_cpu_tensor(host_cache), + "host_cache must be pinned CPU memory"); + STD_TORCH_CHECK(hot_cache.is_cuda(), "hot_cache must be on CUDA"); + STD_TORCH_CHECK( + hot_block_table.is_cuda() && + hot_block_table.scalar_type() == torch::headeronly::ScalarType::Int && + hot_block_table.dim() == 2, + "hot_block_table must be a 2D int32 CUDA tensor"); + STD_TORCH_CHECK(global_indices.is_cuda() && hot_indices.is_cuda() && + device_global_indices.is_cuda() && lru_slots.is_cuda(), + "index tensors must be on CUDA"); + STD_TORCH_CHECK( + global_indices.scalar_type() == torch::headeronly::ScalarType::Int && + hot_indices.scalar_type() == torch::headeronly::ScalarType::Int, + "global_indices/hot_indices must be int32"); + STD_TORCH_CHECK( + device_global_indices.scalar_type() == torch::headeronly::ScalarType::Int, + "device_global_indices must be int32"); + STD_TORCH_CHECK( + lru_slots.scalar_type() == torch::headeronly::ScalarType::Short, + "lru_slots must be int16"); + STD_TORCH_CHECK(global_indices.dim() == 2 && global_indices.stride(1) == 1, + "global_indices must be row-major 2D"); + const int64_t num_rows = global_indices.size(0); + STD_TORCH_CHECK(num_rows >= 0 && num_rows <= INT32_MAX, + "num_rows must fit the CUDA grid"); + STD_TORCH_CHECK(hot_indices.size(0) == num_rows && + hot_indices.size(1) == global_indices.size(1) && + hot_indices.is_contiguous(), + "hot_indices must have one contiguous row per launch row"); + STD_TORCH_CHECK( + device_global_indices.dim() == 2 && device_global_indices.is_contiguous(), + "device_global_indices must be contiguous 2D"); + STD_TORCH_CHECK( + lru_slots.dim() == 2 && + lru_slots.size(0) == device_global_indices.size(0) && + lru_slots.is_contiguous(), + "lru_slots must be contiguous with one row per request state"); + STD_TORCH_CHECK(hot_cache.dim() == 3, + "hot_cache must be [num_blocks, block_size, row_width]"); + const int64_t row_bytes = hot_cache.size(-1) * hot_cache.element_size(); + STD_TORCH_CHECK(row_bytes % 16 == 0, "KV rows must be 16-byte aligned"); + STD_TORCH_CHECK(hot_cache.stride(1) * hot_cache.element_size() == row_bytes, + "hot-cache rows must be contiguous"); + const int64_t hot_block_size = hot_cache.size(1); + const int64_t hot_rows = hot_cache.size(0) * hot_block_size; + const int64_t hot_block_stride = + hot_cache.stride(0) * hot_cache.element_size(); + const int64_t host_rows = check_2d_rows(host_cache, "host_cache", row_bytes); + const auto launch_rows = static_cast(num_rows); + const auto top_k = static_cast(global_indices.size(1)); + const auto hot_size = static_cast(lru_slots.size(1)); + STD_TORCH_CHECK(hot_size >= top_k, "hot buffer size must be >= top_k"); + STD_TORCH_CHECK(hot_size <= 32768, "hot buffer size must fit int16 slots"); + STD_TORCH_CHECK(region_stride == hot_size, + "region_stride must match the LRU size"); + STD_TORCH_CHECK(device_global_indices.size(1) == region_stride, + "device_global_indices must cover the full hot region"); + STD_TORCH_CHECK(hot_block_table.size(1) >= + (region_stride + hot_block_size - 1) / hot_block_size, + "hot_block_table has too few columns"); + STD_TORCH_CHECK(hot_rows < INT32_MAX, "hot indices must fit int32"); + + const int32_t* request_ids_ptr = nullptr; + const int32_t* source_block_table_ptr = nullptr; + int64_t source_bt_stride = 0; + int32_t source_num_reqs = 0; + int32_t source_num_blocks = 0; + STD_TORCH_CHECK( + request_ids.has_value() == source_block_table.has_value(), + "request_ids and source_block_table must be provided together"); + if (source_block_table.has_value()) { + auto const& req = request_ids.value(); + auto const& table = source_block_table.value(); + STD_TORCH_CHECK( + req.is_cuda() && + req.scalar_type() == torch::headeronly::ScalarType::Int && + req.numel() >= launch_rows && req.is_contiguous(), + "request_ids must be contiguous int32 on CUDA with one entry per row"); + STD_TORCH_CHECK( + table.is_cuda() && + table.scalar_type() == torch::headeronly::ScalarType::Int && + table.dim() == 2 && table.stride(1) == 1, + "source_block_table must be a row-major 2D int32 CUDA tensor"); + STD_TORCH_CHECK(source_block_size > 0, + "source_block_size must be positive"); + request_ids_ptr = req.const_data_ptr(); + source_block_table_ptr = table.const_data_ptr(); + source_bt_stride = table.stride(0); + source_num_reqs = static_cast(table.size(0)); + source_num_blocks = static_cast(table.size(1)); + } + const int32_t required_hot_rows = + request_ids_ptr != nullptr ? source_num_reqs : launch_rows; + STD_TORCH_CHECK(hot_block_table.size(0) >= required_hot_rows, + "hot_block_table has too few rows"); + + const int32_t* resident_block_table_ptr = nullptr; + int64_t resident_bt_stride = 0; + int32_t resident_num_reqs = 0; + int32_t resident_num_blocks = 0; + if (resident_block_table.has_value()) { + auto const& table = resident_block_table.value(); + STD_TORCH_CHECK( + source_block_table.has_value() && request_ids.has_value(), + "resident lookup requires request_ids and source_block_table"); + STD_TORCH_CHECK( + table.is_cuda() && + table.scalar_type() == torch::headeronly::ScalarType::Int && + table.dim() == 2 && table.stride(1) == 1, + "resident_block_table must be a row-major 2D int32 CUDA tensor"); + STD_TORCH_CHECK(resident_block_size == hot_block_size, + "resident and hot block sizes must match"); + STD_TORCH_CHECK(resident_null_block >= 0, + "resident null block must be non-negative"); + resident_block_table_ptr = table.const_data_ptr(); + resident_bt_stride = table.stride(0); + resident_num_reqs = static_cast(table.size(0)); + resident_num_blocks = static_cast(table.size(1)); + } + + int32_t* resolved_global_indices_ptr = nullptr; + if (resolved_global_indices.has_value()) { + auto const& resolved = resolved_global_indices.value(); + STD_TORCH_CHECK( + resolved.is_cuda() && + resolved.scalar_type() == torch::headeronly::ScalarType::Int && + resolved.dim() == 2 && resolved.size(0) == launch_rows && + resolved.size(1) == global_indices.size(1) && + resolved.is_contiguous(), + "resolved_global_indices must be contiguous int32 matching indices"); + resolved_global_indices_ptr = resolved.mutable_data_ptr(); + } + + int32_t* valid_counts_ptr = nullptr; + if (valid_counts.has_value()) { + auto const& counts = valid_counts.value(); + STD_TORCH_CHECK( + counts.is_cuda() && + counts.scalar_type() == torch::headeronly::ScalarType::Int && + counts.dim() == 1 && counts.size(0) == launch_rows, + "valid_counts must be int32 on CUDA with one entry per row"); + valid_counts_ptr = counts.mutable_data_ptr(); + } + + int32_t* swap_host_physical_rows_ptr = nullptr; + int32_t* swap_device_physical_rows_ptr = nullptr; + int32_t* swap_counts_ptr = nullptr; + const bool has_compact_swaps = swap_host_physical_rows.has_value(); + STD_TORCH_CHECK(has_compact_swaps == swap_device_physical_rows.has_value() && + has_compact_swaps == swap_counts.has_value(), + "compact swap tensors must be provided together"); + if (has_compact_swaps) { + auto const& globals = swap_host_physical_rows.value(); + auto const& hots = swap_device_physical_rows.value(); + auto const& counts = swap_counts.value(); + STD_TORCH_CHECK( + globals.is_cuda() && hots.is_cuda() && counts.is_cuda() && + globals.scalar_type() == torch::headeronly::ScalarType::Int && + hots.scalar_type() == torch::headeronly::ScalarType::Int && + counts.scalar_type() == torch::headeronly::ScalarType::Int && + globals.dim() == 2 && hots.dim() == 2 && + globals.size(0) == launch_rows && + globals.size(1) == global_indices.size(1) && + hots.size(0) == launch_rows && + hots.size(1) == global_indices.size(1) && + counts.numel() >= launch_rows && globals.is_contiguous() && + hots.is_contiguous() && counts.is_contiguous(), + "compact swaps must be contiguous int32 matching indices"); + swap_host_physical_rows_ptr = globals.mutable_data_ptr(); + swap_device_physical_rows_ptr = hots.mutable_data_ptr(); + swap_counts_ptr = counts.mutable_data_ptr(); + } + + const int32_t* request_state_ptr = nullptr; + int32_t request_state_count = 0; + if (request_state_indices.has_value()) { + auto const& state_indices = request_state_indices.value(); + STD_TORCH_CHECK( + state_indices.is_cuda() && + state_indices.scalar_type() == torch::headeronly::ScalarType::Int && + state_indices.dim() == 1 && state_indices.is_contiguous() && + state_indices.numel() <= INT32_MAX, + "request_state_indices must be contiguous 1D int32 on CUDA"); + request_state_ptr = state_indices.const_data_ptr(); + request_state_count = static_cast(state_indices.numel()); + } else { + STD_TORCH_CHECK(device_global_indices.size(0) >= launch_rows, + "device_global_indices has too few rows"); + } + + // Optional output: 1 at columns requiring a host-to-device swap, 0 elsewhere. + int32_t* miss_mask_ptr = nullptr; + if (miss_mask.has_value()) { + auto const& mm = miss_mask.value(); + STD_TORCH_CHECK( + mm.is_cuda() && + mm.scalar_type() == torch::headeronly::ScalarType::Int && + mm.dim() == 2 && mm.is_contiguous() && mm.size(0) == launch_rows && + mm.size(1) == global_indices.size(1), + "miss_mask must be a contiguous int32 CUDA tensor matching " + "global_indices"); + miss_mask_ptr = mm.mutable_data_ptr(); + } + + int32_t* attention_indices_ptr = nullptr; + if (attention_indices.has_value()) { + auto const& indices = attention_indices.value(); + STD_TORCH_CHECK( + indices.is_cuda() && + indices.scalar_type() == torch::headeronly::ScalarType::Int && + indices.dim() == 2 && indices.size(0) == launch_rows && + indices.stride(1) == 1 && indices.size(1) == global_indices.size(1), + "attention_indices must be row-major int32 matching global_indices"); + STD_TORCH_CHECK(attention_block_stride >= hot_block_size, + "attention block stride must cover one hot block"); + attention_indices_ptr = indices.mutable_data_ptr(); + } + + if (launch_rows == 0 || top_k == 0) { + return; + } + + constexpr int kBlockSize = 1024; + const int hash_size = 2 * top_k; + const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize; + const size_t smem_bytes = + sizeof(int32_t) * (top_k + 2 * (num_buffer_chunks + 1) + hash_size + 3) + + sizeof(int16_t) * (hot_size + hash_size); + + const torch::stable::accelerator::DeviceGuard device_guard( + hot_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + const int64_t attention_row_stride = + attention_indices.has_value() ? attention_indices.value().stride(0) : 0; + const int64_t valid_count_stride = + valid_counts.has_value() ? valid_counts.value().stride(0) : 0; + auto kernel = hisparse_resolve_residency_kernel; + if (smem_bytes > 48 * 1024) { + const cudaError_t attribute_error = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); + STD_TORCH_CHECK(attribute_error == cudaSuccess, + "failed to configure HiSparse swap-in shared memory: ", + cudaGetErrorString(attribute_error)); + } + kernel<<>>( + hot_block_table.const_data_ptr(), + global_indices.const_data_ptr(), request_ids_ptr, + source_block_table_ptr, resident_block_table_ptr, + resolved_global_indices_ptr, valid_counts_ptr, + swap_host_physical_rows_ptr, swap_device_physical_rows_ptr, + swap_counts_ptr, hot_indices.mutable_data_ptr(), + attention_indices_ptr, miss_mask_ptr, + device_global_indices.mutable_data_ptr(), + lru_slots.mutable_data_ptr(), request_state_ptr, + request_state_count, host_rows, hot_block_table.stride(0), hot_block_size, + top_k, hot_size, hash_size, region_stride, attention_block_stride, + source_bt_stride, source_num_reqs, source_num_blocks, + static_cast(source_block_size), resident_bt_stride, + resident_num_reqs, resident_num_blocks, + static_cast(resident_block_size), + static_cast(resident_null_block), global_indices.stride(0), + attention_row_stride, valid_count_stride); + const cudaError_t launch_error = cudaGetLastError(); + STD_TORCH_CHECK(launch_error == cudaSuccess, + "HiSparse residency kernel launch failed: ", + cudaGetErrorString(launch_error)); +} + +void hisparse_gather_plan( + torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& global_indices, + torch::stable::Tensor const& hot_indices, + torch::stable::Tensor const& miss_mask, + std::optional const& request_state_indices, + std::optional const& attention_indices, + int64_t attention_block_stride) { + STD_TORCH_CHECK( + host_cache.device().is_cpu() && is_pinned_cpu_tensor(host_cache), + "host_cache must be pinned CPU memory"); + STD_TORCH_CHECK(hot_cache.is_cuda(), "hot_cache must be on CUDA"); + STD_TORCH_CHECK( + global_indices.is_cuda() && hot_indices.is_cuda() && miss_mask.is_cuda(), + "plan tensors must be on CUDA"); + STD_TORCH_CHECK( + global_indices.scalar_type() == torch::headeronly::ScalarType::Int && + hot_indices.scalar_type() == torch::headeronly::ScalarType::Int && + miss_mask.scalar_type() == torch::headeronly::ScalarType::Int, + "plan tensors must be int32"); + STD_TORCH_CHECK(global_indices.dim() == 2 && global_indices.is_contiguous(), + "global_indices must be contiguous 2D"); + STD_TORCH_CHECK( + hot_indices.size(0) == global_indices.size(0) && + hot_indices.size(1) == global_indices.size(1) && + miss_mask.size(0) == global_indices.size(0) && + miss_mask.size(1) == global_indices.size(1) && + hot_indices.is_contiguous() && miss_mask.is_contiguous(), + "hot_indices/miss_mask must match contiguous 2D global_indices"); + + STD_TORCH_CHECK(hot_cache.dim() == 2 || hot_cache.dim() == 3, + "hot_cache must be a 2D staging buffer or paged 3D cache"); + const int64_t row_bytes = hot_cache.size(-1) * hot_cache.element_size(); + STD_TORCH_CHECK(row_bytes % 16 == 0, "KV rows must be 16-byte aligned"); + STD_TORCH_CHECK( + hot_cache.stride(hot_cache.dim() - 2) * hot_cache.element_size() == + row_bytes, + "hot-cache rows must be contiguous"); + const int64_t host_rows = check_2d_rows(host_cache, "host_cache", row_bytes); + + const auto num_rows = static_cast(global_indices.size(0)); + const auto top_k = static_cast(global_indices.size(1)); + + const int32_t* request_state_ptr = nullptr; + if (request_state_indices.has_value()) { + auto const& state_indices = request_state_indices.value(); + STD_TORCH_CHECK( + state_indices.is_cuda() && + state_indices.scalar_type() == torch::headeronly::ScalarType::Int && + state_indices.numel() >= num_rows, + "request_state_indices must be int32 on CUDA with one entry per row"); + request_state_ptr = state_indices.const_data_ptr(); + } + + int32_t* attention_indices_ptr = nullptr; + if (attention_indices.has_value()) { + auto const& indices = attention_indices.value(); + STD_TORCH_CHECK( + indices.is_cuda() && + indices.scalar_type() == torch::headeronly::ScalarType::Int && + indices.dim() == 2 && indices.is_contiguous() && + indices.size(0) == global_indices.size(0) && + indices.size(1) == global_indices.size(1), + "attention_indices must be a contiguous int32 CUDA tensor matching " + "global_indices"); + attention_indices_ptr = indices.mutable_data_ptr(); + } + + if (num_rows == 0 || top_k == 0) { + return; + } + + // Match the swap-in kernel's block size: the gather serves 3 of every 4 + // layers' misses (index-sharing replay), so per-row copy parallelism is + // the throughput limiter on cold rows. + constexpr int kBlockSize = 1024; + constexpr int kNumWarps = kBlockSize / kWarpSize; + // Interleave columns over enough blocks per row to cover the device even + // for few-row launches (local-prefill staging's single-row layout, small + // decode batches), keeping >= 1 column per warp. + constexpr int kTargetBlocks = 256; + const int max_chunks = std::max(1, (top_k + kNumWarps - 1) / kNumWarps); + const int num_chunks = + std::min(max_chunks, std::max(1, kTargetBlocks / num_rows)); + const dim3 grid(num_rows, num_chunks); + const int32_t hot_block_size = + hot_cache.dim() == 3 ? static_cast(hot_cache.size(1)) : 1; + if (attention_indices.has_value()) { + STD_TORCH_CHECK(attention_block_stride >= hot_block_size, + "attention block stride must cover one hot block"); + } + const int64_t hot_rows = hot_cache.size(0) * hot_block_size; + const int64_t hot_block_stride = + hot_cache.stride(0) * hot_cache.element_size(); + const torch::stable::accelerator::DeviceGuard device_guard( + hot_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + hisparse_gather_plan_kernel<<>>( + static_cast(host_cache.const_data_ptr()), + static_cast(hot_cache.mutable_data_ptr()), + global_indices.const_data_ptr(), + hot_indices.const_data_ptr(), + miss_mask.const_data_ptr(), attention_indices_ptr, + request_state_ptr, host_rows, hot_rows, row_bytes, hot_block_stride, + hot_block_size, top_k, attention_block_stride); +} + +void hisparse_gather_compact(torch::stable::Tensor const& host_cache, + torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& miss_global_indices, + torch::stable::Tensor const& miss_hot_indices, + torch::stable::Tensor const& miss_counts) { + STD_TORCH_CHECK( + host_cache.device().is_cpu() && is_pinned_cpu_tensor(host_cache), + "host_cache must be pinned CPU memory"); + STD_TORCH_CHECK(hot_cache.is_cuda(), "hot_cache must be on CUDA"); + STD_TORCH_CHECK( + miss_global_indices.is_cuda() && miss_hot_indices.is_cuda() && + miss_counts.is_cuda() && + miss_global_indices.scalar_type() == + torch::headeronly::ScalarType::Int && + miss_hot_indices.scalar_type() == + torch::headeronly::ScalarType::Int && + miss_counts.scalar_type() == torch::headeronly::ScalarType::Int && + miss_global_indices.dim() == 2 && miss_hot_indices.dim() == 2 && + miss_global_indices.size(0) == miss_hot_indices.size(0) && + miss_global_indices.size(1) == miss_hot_indices.size(1) && + miss_counts.numel() >= miss_global_indices.size(0) && + miss_global_indices.is_contiguous() && + miss_hot_indices.is_contiguous() && miss_counts.is_contiguous(), + "compact miss plan must be matching contiguous int32 CUDA tensors"); + STD_TORCH_CHECK(hot_cache.dim() == 3, + "hot_cache must be [num_blocks, block_size, row_width]"); + const int64_t row_bytes = hot_cache.size(-1) * hot_cache.element_size(); + STD_TORCH_CHECK(row_bytes % 16 == 0, "KV rows must be 16-byte aligned"); + STD_TORCH_CHECK(hot_cache.stride(1) * hot_cache.element_size() == row_bytes, + "hot-cache rows must be contiguous"); + const int64_t host_rows = check_2d_rows(host_cache, "host_cache", row_bytes); + const auto num_rows = static_cast(miss_global_indices.size(0)); + const auto top_k = static_cast(miss_global_indices.size(1)); + if (num_rows == 0 || top_k == 0) { + return; + } + + constexpr int kBlockSize = 512; + constexpr int kTargetBlocks = 64; + const int num_chunks = std::min(8, std::max(1, kTargetBlocks / num_rows)); + const dim3 grid(num_rows, num_chunks); + const int32_t hot_block_size = static_cast(hot_cache.size(1)); + const int64_t hot_rows = hot_cache.size(0) * hot_block_size; + const int64_t hot_block_stride = + hot_cache.stride(0) * hot_cache.element_size(); + const torch::stable::accelerator::DeviceGuard device_guard( + hot_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + hisparse_gather_compact_kernel<<>>( + static_cast(host_cache.const_data_ptr()), + static_cast(hot_cache.mutable_data_ptr()), + miss_global_indices.const_data_ptr(), + miss_hot_indices.const_data_ptr(), + miss_counts.const_data_ptr(), host_rows, hot_rows, row_bytes, + hot_block_stride, hot_block_size, top_k); +} diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 13429a54e6..6893c6913d 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -661,6 +661,54 @@ void concat_and_cache_mla_grouped( std::optional kv_scales, const std::string& kv_cache_dtype); +#ifndef USE_ROCM +// HiSparse kernels use raw PTX in the row copy; CUDA-only. +void hisparse_resolve_residency( + torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& hot_block_table, + torch::stable::Tensor const& global_indices, + torch::stable::Tensor& hot_indices, + torch::stable::Tensor& device_global_indices, + torch::stable::Tensor& lru_slots, + std::optional const& request_state_indices, + int64_t region_stride, + std::optional const& miss_mask, + std::optional const& attention_indices, + int64_t attention_block_stride, + std::optional const& request_ids, + std::optional const& source_block_table, + int64_t source_block_size, + std::optional const& resolved_global_indices, + std::optional const& valid_counts, + std::optional const& swap_host_physical_rows, + std::optional const& swap_device_physical_rows, + std::optional const& swap_counts, + std::optional const& resident_block_table, + int64_t resident_block_size, int64_t resident_null_block); + +void hisparse_invalidate_written_slots( + torch::stable::Tensor& device_global_indices, + torch::stable::Tensor const& request_state_indices, + torch::stable::Tensor const& req_id_per_token, + torch::stable::Tensor const& written_slots); + +void hisparse_gather_plan( + torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& global_indices, + torch::stable::Tensor const& hot_indices, + torch::stable::Tensor const& miss_mask, + std::optional const& request_state_indices, + std::optional const& attention_indices, + int64_t attention_block_stride); + +void hisparse_gather_compact(torch::stable::Tensor const& host_cache, + torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& miss_global_indices, + torch::stable::Tensor const& miss_hot_indices, + torch::stable::Tensor const& miss_counts); + +#endif // !USE_ROCM + // NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla void concat_and_cache_mla_rope_fused( torch::stable::Tensor& positions, torch::stable::Tensor& q_pe, @@ -704,7 +752,10 @@ void cp_gather_and_upconvert_fp8_kv_cache( torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] torch::stable::Tensor const& workspace_starts, // [BATCH] int64_t batch_size, - std::optional seq_starts = std::nullopt); + std::optional seq_starts = std::nullopt, + std::optional host_cache = std::nullopt, + std::optional host_row_ids = std::nullopt, + std::optional device_row_ids = std::nullopt); // Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace void cp_gather_and_upconvert_nvfp4_kv_cache( diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 336c495113..cb6f463396 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -987,6 +987,57 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { " Tensor? kv_scales=None," " str kv_cache_dtype='auto') -> ()"); +#ifndef USE_ROCM + ops.def( + "hisparse_resolve_residency(Tensor host_cache," + " Tensor! hot_cache," + " Tensor hot_block_table," + " Tensor global_indices," + " Tensor! hot_indices," + " Tensor! device_global_indices," + " Tensor! lru_slots," + " Tensor? request_state_indices," + " int region_stride," + " Tensor(a!)? miss_mask=None," + " Tensor(b!)? attention_indices=None," + " int attention_block_stride=0," + " Tensor? request_ids=None," + " Tensor? source_block_table=None," + " int source_block_size=0," + " Tensor(c!)? resolved_global_indices=None," + " Tensor(d!)? valid_counts=None," + " Tensor(e!)? swap_host_physical_rows=None," + " Tensor(f!)? swap_device_physical_rows=None," + " Tensor(g!)? swap_counts=None," + " Tensor? resident_block_table=None," + " int resident_block_size=0," + " int resident_null_block=0) -> ()"); + + ops.def( + "hisparse_invalidate_written_slots(Tensor! device_global_indices," + " Tensor request_state_indices," + " Tensor req_id_per_token," + " Tensor written_slots) -> ()"); + + ops.def( + "hisparse_gather_plan(Tensor host_cache," + " Tensor! hot_cache," + " Tensor global_indices," + " Tensor hot_indices," + " Tensor miss_mask," + " Tensor? request_state_indices," + " Tensor(a!)? attention_indices=None," + " int attention_block_stride=0) -> ()"); + + ops.def( + "hisparse_gather_compact(Tensor host_cache," + " Tensor! hot_cache," + " Tensor miss_global_indices," + " Tensor miss_hot_indices," + " Tensor miss_counts) -> ()"); + +#endif // !USE_ROCM + // Rotate Q and K, then write to kv cache for MLA ops.def( "concat_and_cache_mla_rope_fused(" @@ -1023,7 +1074,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { ops.def( "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, " "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? " - "seq_starts) -> ()"); + "seq_starts, Tensor? host_cache=None, Tensor? host_row_ids=None, Tensor? " + "device_row_ids=None) -> ()"); ops.def( "cp_gather_and_upconvert_nvfp4_kv_cache(Tensor src_cache, Tensor! dst, " @@ -1091,6 +1143,15 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) { ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla)); ops.impl("concat_and_cache_mla_grouped", TORCH_BOX(&concat_and_cache_mla_grouped)); + +#ifndef USE_ROCM + ops.impl("hisparse_resolve_residency", + TORCH_BOX(&hisparse_resolve_residency)); + ops.impl("hisparse_invalidate_written_slots", + TORCH_BOX(&hisparse_invalidate_written_slots)); + ops.impl("hisparse_gather_plan", TORCH_BOX(&hisparse_gather_plan)); + ops.impl("hisparse_gather_compact", TORCH_BOX(&hisparse_gather_compact)); +#endif // !USE_ROCM ops.impl("concat_and_cache_mla_rope_fused", TORCH_BOX(&concat_and_cache_mla_rope_fused)); ops.impl("convert_fp8", TORCH_BOX(&convert_fp8)); diff --git a/docs/design/hisparse.md b/docs/design/hisparse.md new file mode 100644 index 0000000000..965ab81d8d --- /dev/null +++ b/docs/design/hisparse.md @@ -0,0 +1,244 @@ +# HiSparse local KV offload architecture + +Status: experimental + +## The short version + +There are three different jobs: + +1. The normal KV cache system manages GPU block pools and tables. +2. `HiSparseCoordinator` manages logical host blocks, source-prefix identity, and + host/GPU residency transitions. +3. `HiSparseConnector` carries residency work between scheduler and worker; + `HiSparseWorker` coordinates transfers, while per-cache + `HiSparseRuntime` objects own host/hot views and GPU replacement state. + +Neither worker-side object allocates or frees logical blocks. HMA provides the +GPU allocation shared by the resident and hot groups; it does not manage CPU +memory or KV identity. + +```text +request ────► HiSparseCoordinator ── source blocks + residency policy + │ + ├──► KV cache manager ── resident/hot GPU leases (HMA) + │ + └──► HiSparseConnector ── host bytes + copies + GPU LRU +``` + +This is a local KV connector. When P/D or another offload connector is also +configured, `MultiConnector` composes it with `HiSparseConnector`. + +`host_pool_gib` is configured on `HiSparseConnector` and is the usable +host-cache capacity per data-parallel replica, not a node-wide memory budget. +Tensor-parallel ranks hold replicated views of that logical cache. Those views +may use private per-rank backing or one shared +physical allocation without changing the configured capacity. Physical host +memory consumption is therefore topology- and implementation-dependent. The +realized capacity may be slightly smaller because the budget is rounded down +to complete host blocks. + +## Ownership + +| Thing | Owner | What “owner” means | +| --- | --- | --- | +| HiSparse source and prefix identity | `HiSparseCoordinator` | maps tokens to logical host blocks | +| Resident GPU block leases | normal KV cache manager | allocates and frees HMA blocks | +| Resident block tables | normal KV cache manager | tells attention where resident pages are | +| Residency transitions | `HiSparseCoordinator` | plans spill-before-free transactions | +| Logical host block allocation | `HiSparseCoordinator` | owns the separate CPU block pool and its lifecycle | +| Pinned host-pool lifecycle | `HiSparseWorker` | worker-wide backing and teardown | +| Per-cache host view and hot contents | `HiSparseRuntime` | binds host/hot storage and fills cache-manager-provided hot leases | +| Hot row map and LRU | `HiSparseRuntime` | resolves hits and chooses victims on GPU | +| Resident-cache route | `HiSparseCacheHandle` | exposes resident or host/hot resolution to attention | +| Sparse attention | attention backend | consumes a device cache and physical row IDs | +| HMA | allocator | provides GPU capacity; owns no KV meaning | + +The key distinction is logical allocation versus contents. `HiSparseCoordinator` +owns host block IDs and their request/prefix associations. `HiSparseWorker` and +its per-cache runtimes own the corresponding bytes. The normal cache manager +sees only device pools. +The source group has `block_pool_id=None`; device-pool consumers must narrow it +before indexing, so host ownership cannot masquerade as a numeric GPU pool. + +## Code boundary + +```text +scheduler process worker process + +HiSparseConnector HiSparseConnector + └─ HiSparseCoordinator └─ HiSparseWorker + │ │ + │ connector metadata ├─ host bytes + │ - page transfers ├─ copy scheduling + │ - block-table replacements └─ per-layer hot state + └───────────────────────────────────────────────►│ + ◄──────── connector worker metadata ─────────────┘ + enqueued and completed transfer IDs +``` + +The command travels in `kv_connector_metadata`; transfer updates return in +`KVConnectorOutput.kv_connector_worker_meta`. The model runner does not +interpret page transfers. Enqueue acknowledgements let the scheduler release +source leases in stream order; completion acknowledgements publish the copied +host pages. + +## Resident device pages + +Resident pages are intentionally outside `HiSparseRuntime`. + +KV-cache initialization binds cache-manager allocations to the attention-facing +`HiSparseCacheHandle` before constructing `HiSparseWorker`. That same +handle's runtime retains the resident source index needed by a transfer plan. +There is no second resident object or registration wrapper. + +```text +KV cache setup + │ + ├─ bind resident allocation ──► HiSparseCacheHandle + │ cache + block table + slot mapping + │ + ├─ bind host/hot allocation ──► HiSparseRuntime + │ host + hot + GPU LRU + │ + └─ register cache handles ────► HiSparseWorker + step-level transfers + +HiSparseWorker registers the same HiSparseCacheHandle objects directly +``` + +Attention construction links each layer to the most recent layer that actually +owns an indexer. This releases a follower's duplicate LRU tensors before GPU +memory profiling. Cache binding only attaches storage; it does not infer +semantic groups from the physical packed-tensor order. The construction cursor +is discarded with the worker's pinned state. + +Every HiSparse decode batch uses the same fused resolver. It checks resident +pages first, then hot rows, then pinned host memory. A resident hit exits inside +the kernel before hot-LRU lookup or host copying; there is no framework-level +residency route or separate CUDA graph. No CPU decision is added to the decode +path. The resolver consumes the existing graph-stable request mapping from +attention metadata; neither the worker nor individual cache handles keep a +duplicate mapping. + +Speculative decoding resolves and consumes each verification step in order. +Each step receives distinct replayable plan rows while sharing the request's +hot-cache state, so a later step cannot reuse a hot row before an earlier step +has consumed it. + +## P/D import target + +The decoder chooses the landing target once per request from the normal cache +admission calculation. If the complete imported prefix fits the device pools, +NIXL transfers it directly into resident GPU pages. Otherwise, if the fixed +host-backed GPU footprint and host source blocks fit, the request imports into +the host tier. There is no context-length threshold or other heuristic, and a +request waiting for capacity retains its choice across admission retries. + +A host import reads through a bounded decoder-GPU staging pool before copying +into registered host memory. Pages needed immediately are mirrored into their +resident destinations during that copy. Both landing targets then use the same +fused decode resolver described above. + +## Indexer KV offloading + +HiSparse does not keep a private CPU copy of indexer KV. The indexer remains a +normal prefix-cacheable GPU cache group. If `OffloadingConnector` is configured +with HiSparse, it stores and restores that group through the generic KV +offloading path; HiSparse continues to own only the sparse MLA host tier. + +The two prefix sources can have different hit lengths. When the HiSparse host +prefix extends beyond the GPU-resident indexer prefix, the scheduler asks +`OffloadingConnector` to restore only the missing indexer suffix, capped at the +host prefix boundary. If that suffix is unavailable, all groups fall back to +the shorter prefix they share. NIXL P/D transfers continue to place indexer KV +directly in its GPU group. + +## Spill transaction + +A resident block cannot be reused until its contents have been handed to the +worker. + +```text +HiSparseCoordinator HiSparseWorker + │ │ + │ pin source and destination leases │ + │── SparseKVPageTransfer ─────────────►│ + │ │ enqueue GPU-to-host copy + │◄── enqueued transfer ID ────────────│ + │ replace resident table entry │ + │ release resident lease to HMA │ + │ │ copy reaches its event + │◄── completed transfer ID ───────────│ + │ mark host page valid │ + │ release destination host lease │ +``` + +“Enqueued” means the copy has entered the worker stream. Stream ordering makes +it safe to reuse the resident GPU block for later work, but the host page is +not yet published. “Completed” means the worker has observed the copy's event; +only then does the coordinator publish the host page for prefix reuse and +release its destination lease. A host-write event separately protects direct +CPU readers from writes already queued on the accelerator. + +The worker transfer contains only its transfer ID and physical copy +coordinates. Request identity and logical page state remain in the scheduler. + +## Hot lookup and LRU + +The NVIDIA CUDA path keeps replacement entirely on the accelerator: + +```text +top-K logical positions + │ + ▼ +resident page? ── yes ──► resident physical row + │ no + ▼ +hot row? ──────── yes ──► existing hot physical row + update GPU LRU + │ no + ▼ +choose GPU LRU victim ──► copy pinned host row ──► hot physical row +``` + +ROCm is not currently supported because the fused HiSparse cache operations are +implemented only by CUDA kernels. A future platform-specific worker may provide +the same command, output, and cache-resolution boundaries. + +## Main classes + +| Class | Inherits / implements | Responsibility | +| --- | --- | --- | +| `HiSparseCoordinator` | plain scheduler component | host allocation, source prefixes, resident leases, and spill state machine | +| `HiSparseConnector` | `KVConnectorBase_V1`, `SupportsHMA` | scheduler/worker metadata and lifecycle boundary | +| `HiSparseResidentManager` | `SingleTypeKVCacheManager` | normal block-pool bookkeeping with host-backed holes | +| `PagedCacheView` | immutable data object | shared resident/hot HMA tensor binding | +| `HiSparseWorker` | connector-owned worker component | worker-wide transfer scheduling and host-pool lifecycle | +| `HiSparseRuntime` | plain worker-owned component | per-cache host/hot tensors, GPU LRU, and fused resolution | +| `HiSparseCacheHandle` | plain attention component | resident view and fused cache resolution | +| `SparseKVOffloadCommand` | dataclass | opaque scheduler-to-worker work | + +## Performance invariants + +- Resident hits bypass hot-LRU lookup and host copies inside the fused resolver. +- Hot lookup, victim selection, and LRU updates stay on the GPU. +- A hot miss still copies directly from registered pinned host memory. +- Top-K resolution stays inside the attention invocation and remains graph + capturable. +- Compatible layers still share one miss plan. +- Index-sharing followers release their private LRU state before memory sizing. +- Indexer KV is untouched by HiSparse unless a generic KV offloader is configured. +- Resident and hot leases can still share one packed HMA allocation. +- No device scalar readback or CPU/device metadata round trip is added. +- The abstraction wraps the fused kernel; it does not add another kernel + launch. +- When HiSparse is disabled, the scheduler does not construct an offload + command or empty update table. + +## What remains platform-specific + +The command/result and attention-layer boundaries can be shared. The host +allocator, copy implementation, hot layout, and replacement policy should stay +platform-specific. NVIDIA uses the current accelerator LRU and fused host/hot +kernel. ROCm is not currently supported; AMD or other accelerator backends can +implement their own worker without forcing NVIDIA's policy into the shared +boundary. diff --git a/tests/evals/gsm8k/gsm8k_eval.py b/tests/evals/gsm8k/gsm8k_eval.py index b5a1e6a9cc..83e54afd9e 100644 --- a/tests/evals/gsm8k/gsm8k_eval.py +++ b/tests/evals/gsm8k/gsm8k_eval.py @@ -364,6 +364,12 @@ def main() -> None: type=int, help="Maximum number of concurrent requests", ) + parser.add_argument( + "--request-timeout-seconds", + type=float, + default=600, + help="Timeout for each request, including time waiting for a connection", + ) parser.add_argument("--save-results", type=str, help="Save results to JSON file") args = parser.parse_args() @@ -377,6 +383,7 @@ def main() -> None: temperature=args.temperature, seed=args.seed, max_concurrency=args.max_concurrency, + request_timeout_seconds=args.request_timeout_seconds, ) # Print results to terminal diff --git a/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py b/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py index ac46073559..ffdb76e176 100644 --- a/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py +++ b/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py @@ -217,7 +217,8 @@ def test_flashinfer_mla_dense_decode_unified_slot_view(): assert (out_ref - out_view).abs().max().item() == 0.0 -def test_flashmla_fp8_sparse_decode_unified_slot_view(): +@pytest.mark.parametrize("next_n", [1, 2, 4]) +def test_flashmla_fp8_sparse_decode_unified_slot_view(next_n: int): """FlashMLA fp8 sparse decode (DeepSeek V3.2/V4 DSA path) must read a unified-slot block-major view bit-identically to a contiguous fp8_ds_mla cache, with finite nonzero output.""" @@ -240,7 +241,10 @@ def test_flashmla_fp8_sparse_decode_unified_slot_view(): n_layers = 3 layer = 1 - q = torch.randn(batch, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1 + q = ( + torch.randn(batch, next_n, h_q, head_dim, device=dev, dtype=torch.bfloat16) + * 0.1 + ) # Structurally valid fp8 ds_mla payload: 512B fp8 + 16B f32 scales + 128B # bf16 rope (random bytes corrupt the scale region and yield NaNs). @@ -273,13 +277,13 @@ def test_flashmla_fp8_sparse_decode_unified_slot_view(): # Sparse indices: each batch uses its own disjoint blocks. blocks_per_batch = num_blocks // batch - idx = torch.full((batch, 1, topk), -1, device=dev, dtype=torch.int32) + idx = torch.full((batch, next_n, topk), -1, device=dev, dtype=torch.int32) for b in range(batch): slots: list[int] = [] for blk in range(b * blocks_per_batch, (b + 1) * blocks_per_batch): slots.extend(blk * page + off for off in range(page)) slots_t = torch.tensor(slots[:topk], device=dev, dtype=torch.int32) - idx[b, 0, : slots_t.numel()] = slots_t + idx[b, :, : slots_t.numel()] = slots_t def run(kc): meta, num_splits = fm.get_mla_metadata() diff --git a/tests/kernels/test_cp_gather_fp8.py b/tests/kernels/test_cp_gather_fp8.py index 35afe33808..fe539abb4f 100644 --- a/tests/kernels/test_cp_gather_fp8.py +++ b/tests/kernels/test_cp_gather_fp8.py @@ -316,6 +316,48 @@ def test_cp_gather_fp8_shuffled_blocks(): assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:]) +def test_cp_gather_fp8_mixes_resident_and_pinned_host_rows(): + seq_lens = [7, 5] + block_size = 4 + ( + cache, + block_table, + workspace_starts, + num_reqs, + total_tokens, + expected, + ) = _build_test_case(seq_lens, block_size) + + device_cache = cache.clone() + host_cache = cache.cpu().pin_memory().view(-1, ENTRY_BYTES) + device_rows = torch.arange(host_cache.shape[0], dtype=torch.int32, device="cuda") + use_device = device_rows % 2 == 0 + host_rows = device_rows.clone() + device_rows.masked_fill_(~use_device, -1) + + host_cache[use_device.cpu()] = 0 + device_cache.view(-1, ENTRY_BYTES)[~use_device] = 0 + dst = torch.empty( + total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda" + ) + + ops.cp_gather_and_upconvert_fp8_kv_cache( + device_cache, + dst, + block_table, + workspace_starts, + num_reqs, + host_cache=host_cache, + host_row_ids=host_rows, + device_row_ids=device_rows, + ) + + 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:]) + + @pytest.mark.parametrize( "gather_seq_lens,seq_starts", [ diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index d0562f6d2d..fbe22af2f8 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -234,6 +234,76 @@ def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str assert (topk == -1).all(), "topk buffer not cleared on indexer layer" +def test_fused_norm_rope_packed_indexer_block_stride(): + """Indexer writes use their HMA block stride and independent slot map.""" + torch.manual_seed(7) + dev = "cuda" + num_tokens = block_size = 2 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) + q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16) + kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16) + k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16) + qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16) + kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16) + ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16) + ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32) + ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32) + cos_sin = make_cos_sin(32, ROPE_DIM, dev) + + idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4 + packed_block_stride = block_size * idx_row + 64 + backing = torch.zeros(2 * packed_block_stride, device=dev, dtype=torch.uint8) + idx_cache = torch.as_strided( + backing, + (2, block_size, idx_row), + (packed_block_stride, idx_row, 1), + ) + mla_cache = torch.zeros( + 1, + block_size, + KV_LORA + ROPE_DIM, + device=dev, + dtype=torch.bfloat16, + ) + topk = torch.zeros(num_tokens, 8, device=dev, dtype=torch.int32) + + K.fused_norm_rope( + pos, + q_c, + qw, + EPS, + kv_c, + kvw, + EPS, + k_pe, + cos_sin, + ik, + ikw, + ikb, + EPS, + cos_sin, + topk, + slot_mapping=torch.arange(num_tokens, device=dev, dtype=torch.int64), + indexer_slot_mapping=torch.arange( + block_size, block_size + num_tokens, device=dev, dtype=torch.int64 + ), + indexer_k_cache=idx_cache, + mla_kv_cache=mla_cache, + has_indexer=True, + ) + + ik_ref = rope(layer_norm(ik, ikw, ikb), pos, cos_sin, interleave=False) + q_ref, s_ref = ue8m0_quant(ik_ref) + packed = idx_cache[1].reshape(-1) + values = ( + packed[: block_size * INDEX_HEAD_DIM].view(FP8).view(block_size, INDEX_HEAD_DIM) + ) + scales = packed[block_size * INDEX_HEAD_DIM :].view(torch.float32) + assert_fp8(values, q_ref, "packed indexer-K fp8") + torch.testing.assert_close(scales, s_ref, rtol=0, atol=0) + assert (backing[block_size * idx_row : packed_block_stride] == 0).all() + + def test_fused_norm_rope_materializes_qk_without_local_cache_slots(): """DCP non-owner ranks still need valid query shards for query AllGather.""" torch.manual_seed(7) diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index de6518ebec..f8e7a89cca 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -7,6 +7,7 @@ import pytest import torch import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer +import vllm.models.deepseek_v32.attention as deepseek_attention from vllm.config import CUDAGraphMode from vllm.models.deepseek_v32 import attention as deepseek_v32_attention from vllm.models.deepseek_v32.attention import DeepseekV32Attention @@ -16,6 +17,51 @@ INDEXER_LAYER = "model.layers.0.self_attn.indexer.k_cache" MLA_LAYER = "model.layers.0.self_attn.attn" +def test_sparse_attention_refreshes_batch_state_inside_eager_segment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + metadata = SimpleNamespace(num_actual_tokens=1) + + class BatchStateRefreshed(Exception): + pass + + def refresh_batch_state(attn_metadata): + assert attn_metadata is metadata + raise BatchStateRefreshed + + layer = SimpleNamespace( + indexer=None, + skip_topk=True, + layer_name=MLA_LAYER, + impl=SimpleNamespace( + prepare_for_batch=refresh_batch_state, + record_logical_topk_ready=lambda: None, + ), + ) + monkeypatch.setattr( + deepseek_attention, + "get_attention_context", + lambda layer_name: (metadata, None, torch.empty(1), None), + ) + + with pytest.raises(BatchStateRefreshed): + DeepseekV32Attention._sparse_indexer_and_attn( + layer, + torch.empty(1, dtype=torch.long), + torch.empty(1, 1), + torch.empty(1, 1, 1), + torch.empty(1, 1, 1), + None, + None, + None, + None, + None, + torch.empty(1, 1, 1), + torch.empty(1, 1, 1), + torch.empty(1, 1), + ) + + def make_indexer_metadata( *, num_decodes: int = 0, @@ -246,6 +292,10 @@ def test_deepseek_v32_dispatches_selected_mha( _fp8_query=fp8_query, _use_sparse_mha=lambda _: True, rotary_emb=lambda _positions, q: (q + 1, None), + impl=SimpleNamespace( + record_logical_topk_ready=lambda: None, + prepare_for_batch=lambda _: None, + ), forward_impl=record_forward_impl, ) q_nope = torch.randn(2, 1, 2) diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index a8ecd1276e..dd162629b8 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/tests/model_executor/test_flashinfer_autotune_warmup.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from inspect import signature from types import SimpleNamespace from typing import Any from unittest.mock import Mock, call, patch @@ -76,8 +77,16 @@ def test_flashinfer_autotune_token_counts_are_bounded_and_deduplicated(): assert token_counts == (32,) -def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(): +@pytest.mark.parametrize("skip_attn", [False, True]) +def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(skip_attn): + from vllm.v1.worker.gpu.model_runner import GPUModelRunner as V2Runner + from vllm.v1.worker.gpu_model_runner import GPUModelRunner as V1Runner + runner = _make_runner([]) + dummy_run_signature = signature((V2Runner if skip_attn else V1Runner)._dummy_run) + runner._dummy_run.side_effect = lambda **kwargs: dummy_run_signature.bind( + None, **kwargs + ) max_buckets = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 8192) deferred_buckets = (1, 2, 4, 8, 16, 32, 64, 128) @@ -93,7 +102,7 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(): ) as get_buckets, patch("vllm.utils.flashinfer.autotune") as autotune, ): - _run_flashinfer_autotune_dummy_runs(runner) + _run_flashinfer_autotune_dummy_runs(runner, skip_attn=skip_attn) assert get_buckets.call_args_list == [call(8192), call(128)] assert autotune.call_args_list == [ @@ -106,11 +115,13 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(): skip_eplb=True, is_profile=True, randomize_inputs=True, + **({"skip_attn": True} if skip_attn else {}), ), call( num_tokens=128, skip_eplb=True, is_profile=True, randomize_inputs=True, + **({"skip_attn": True} if skip_attn else {}), ), ] diff --git a/tests/test_config.py b/tests/test_config.py index 2524968cce..7b6027313a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -19,10 +19,12 @@ import vllm.config.vllm as vllm_config_module import vllm.envs as envs from vllm.compilation.backends import VllmBackend from vllm.config import ( + AttentionConfig, CacheConfig, CompilationConfig, DeviceConfig, EngramConfig, + HiSparseConfig, KernelConfig, KVTransferConfig, ModelConfig, @@ -279,6 +281,62 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected): assert envs.VLLM_USE_V2_MODEL_RUNNER is expected +def test_hisparse_requires_v2_model_runner(): + config = object.__new__(VllmConfig) + config.attention_config = AttentionConfig(hisparse_config=HiSparseConfig()) + + with patch.object(envs, "VLLM_USE_V2_MODEL_RUNNER", None): + assert config.use_v2_model_runner + with ( + patch.object(envs, "VLLM_USE_V2_MODEL_RUNNER", False), + pytest.raises(ValueError, match="requires Model Runner V2"), + ): + _ = config.use_v2_model_runner + + +def test_hisparse_rejects_decode_context_parallelism(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr(current_platform, "device_count", lambda: 2) + with pytest.raises(ValueError, match="decode context parallelism"): + VllmConfig( + attention_config=AttentionConfig(hisparse_config=HiSparseConfig()), + parallel_config=ParallelConfig( + tensor_parallel_size=2, + decode_context_parallel_size=2, + ), + ) + + +def test_hisparse_rejects_pipeline_parallelism(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr(current_platform, "device_count", lambda: 2) + with pytest.raises(ValueError, match="pipeline parallelism"): + VllmConfig( + attention_config=AttentionConfig(hisparse_config=HiSparseConfig()), + parallel_config=ParallelConfig(pipeline_parallel_size=2), + ) + + +def test_hisparse_rejects_disabled_hybrid_kv_cache_manager(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True) + with pytest.raises(ValueError, match="requires the hybrid KV cache manager"): + VllmConfig( + attention_config=AttentionConfig(hisparse_config=HiSparseConfig()), + scheduler_config=SchedulerConfig( + max_model_len=2048, + is_encoder_decoder=False, + disable_hybrid_kv_cache_manager=True, + ), + ) + + +def test_hisparse_rejects_non_cuda(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: False) + with pytest.raises(ValueError, match="requires NVIDIA CUDA"): + VllmConfig(attention_config=AttentionConfig(hisparse_config=HiSparseConfig())) + + def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): """ROCm keeps the DSA models (DeepSeek V3.2/V4, GLM-5.2) on their compiled MRV1 paths and off breakable cudagraphs by default.""" @@ -305,7 +363,8 @@ def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): # (warning_once args must be hashable for its lru_cache). monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False) config = SimpleNamespace( - model_config=SimpleNamespace(architectures=["DeepseekV32ForCausalLM"]) + model_config=SimpleNamespace(architectures=["DeepseekV32ForCausalLM"]), + attention_config=AttentionConfig(), ) assert VllmConfig.use_v2_model_runner.fget(config) is False finally: @@ -385,6 +444,7 @@ def test_dsa_models_default_to_mrv2_and_breakable_cudagraph( ) config = SimpleNamespace( model_config=model_config, + attention_config=AttentionConfig(), speculative_config=SimpleNamespace(method="mtp") if with_mtp else None, parallel_config=SimpleNamespace(prefill_context_parallel_size=1), compilation_config=CompilationConfig( @@ -897,7 +957,10 @@ def test_models_default_to_v2_model_runner(model_config, expected, monkeypatch): monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False) monkeypatch.setattr(vllm_config_module, "HAS_TRITON", True) monkeypatch.setattr(current_platform, "is_rocm", lambda: False) - config = SimpleNamespace(model_config=model_config) + config = SimpleNamespace( + model_config=model_config, + attention_config=AttentionConfig(), + ) config._get_v2_model_runner_unsupported_features = lambda: [] assert VllmConfig.use_v2_model_runner.fget(config) is expected @@ -2787,6 +2850,7 @@ def test_draft_sample_method_gumbel_is_rejected(): def _watermarked_vllm_config() -> VllmConfig: config = object.__new__(VllmConfig) config.watermark_config = WatermarkConfig(key=42) + config.attention_config = AttentionConfig() config.speculative_config = None return config diff --git a/tests/v1/attention/test_backend_per_kind.py b/tests/v1/attention/test_backend_per_kind.py index 9d180d3aee..ab77d58825 100644 --- a/tests/v1/attention/test_backend_per_kind.py +++ b/tests/v1/attention/test_backend_per_kind.py @@ -1,16 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for per-KV-group attention backend selection (backend_per_kind).""" +"""Tests for attention configuration and per-KV-group backend selection.""" + +from types import SimpleNamespace import pytest -from vllm.config.attention import AttentionConfig +from vllm.config.attention import AttentionConfig, HiSparseConfig from vllm.model_executor.layers.attention.attention import ( _largest_kernel_block_within, ) from vllm.v1.attention.backend import AttentionType, MultipleOf from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.selector import get_attn_spec_kind +from vllm.v1.hisparse.runtime import ResolvedHiSparseConfig from vllm.v1.kv_cache_interface import KVCacheSpecKind @@ -86,3 +89,52 @@ def test_backend_per_kind_rejects_unknown_kind(): def test_backend_per_kind_defaults_empty(): assert AttentionConfig().backend_per_kind == {} + + +def test_hisparse_device_buffer_size_boundaries(): + vllm_config = SimpleNamespace( + attention_config=AttentionConfig(hisparse_config=HiSparseConfig()), + speculative_config=None, + ) + resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128) + assert resolved is not None + assert resolved.device_buffer_size == 256 + + vllm_config.attention_config.hisparse_config = HiSparseConfig( + device_buffer_size=127 + ) + with pytest.raises(ValueError, match="expected at least 128"): + ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128) + + vllm_config.attention_config.hisparse_config = HiSparseConfig( + device_buffer_size=32768 + ) + resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128) + assert resolved is not None + assert resolved.device_buffer_size == 32768 + + vllm_config.attention_config.hisparse_config = HiSparseConfig( + device_buffer_size=32769 + ) + with pytest.raises(ValueError, match="int16 slot-index limit"): + ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128) + + +def test_hisparse_device_buffer_covers_speculative_window(): + vllm_config = SimpleNamespace( + attention_config=AttentionConfig(hisparse_config=HiSparseConfig()), + speculative_config=SimpleNamespace( + num_speculative_tokens=3, + parallel_drafting=False, + ), + ) + + resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128) + assert resolved is not None + assert resolved.device_buffer_size == 5 * 128 + + vllm_config.attention_config.hisparse_config = HiSparseConfig( + device_buffer_size=4 * 128 - 1 + ) + with pytest.raises(ValueError, match="expected at least 512"): + ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index edc6cbaded..ee5737a995 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -200,6 +200,7 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype( layer = SimpleNamespace( kv_cache_dtype=cache_dtype, head_size=576, + indexer=None, non_causal_multi_token_decode=False, sliding_window=None, ) diff --git a/tests/v1/attention/test_mla_context_chunks.py b/tests/v1/attention/test_mla_context_chunks.py index 2ed73b2c98..2d63b165a0 100644 --- a/tests/v1/attention/test_mla_context_chunks.py +++ b/tests/v1/attention/test_mla_context_chunks.py @@ -55,6 +55,7 @@ def build_chunked_context( def test_sparse_context_lengths_do_not_force_gpu_sync(monkeypatch): """Sparse MLA must pin computed context lengths before the H2D copy.""" builder = SimpleNamespace( + _context_lens_cpu=torch.empty(2, dtype=torch.int32, pin_memory=True), chunked_prefill_workspace=torch.empty((2048, 1)), chunked_prefill_workspace_size=1024, kv_cache_spec=SimpleNamespace(block_size=BLOCK_SIZE), @@ -113,6 +114,7 @@ def test_chunks_gather_every_context_row_exactly_once(context_lens, workspace_si query_lens = [4] * len(context_lens) metadata = build_chunked_context(context_lens, query_lens, workspace_size) assert metadata is not None + assert metadata.context_lens.tolist() == context_lens gathered: dict[int, list[tuple[int, int]]] = {} previous_request_start = -1 @@ -253,6 +255,7 @@ def test_dcp_chunks_fit_the_per_rank_row_budget(): dcp_local_block_size=interleave, ) assert metadata is not None + assert metadata.context_lens.tolist() == context_lens virtual_block_size = interleave * dcp_world_size local_cursor: dict[int, int] = {} @@ -330,3 +333,23 @@ def test_dcp_reorg_uses_each_chunks_local_starts(): ) torch.testing.assert_close(reorganized, torch.cat(expected)) + + +def test_chunked_context_metadata_owns_context_lengths(): + """Reusing CPU scratch must not overwrite an earlier batch's context lengths.""" + context_lens = torch.tensor([16, 32], dtype=torch.int32) + metadata = build_mla_chunked_context_metadata( + context_lens_cpu=context_lens, + prefill_query_start_loc_cpu=torch.tensor([0, 4, 8], dtype=torch.int32), + chunked_prefill_workspace=torch.empty((64, 1)), + chunked_prefill_workspace_size=64, + block_size=BLOCK_SIZE, + align_chunk_to_block=True, + device=torch.device("cpu"), + dcp_world_size=1, + dcp_local_block_size=1, + dcp_virtual_block_size=1, + ) + context_lens.zero_() + assert metadata is not None + assert metadata.context_lens.tolist() == [16, 32] diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 425db95a8a..e7239320c2 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -3,7 +3,9 @@ """Unit tests for the sparse MLA backends and utilities.""" import math +from collections import deque from types import MethodType, SimpleNamespace +from unittest.mock import MagicMock import pytest import torch @@ -20,13 +22,19 @@ from tests.v1.attention.utils import ( create_vllm_config, ) from vllm import _custom_ops as ops -from vllm.config import set_current_vllm_config +from vllm.config import HiSparseConfig, SpeculativeConfig, set_current_vllm_config +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( + HiSparseConnectorWorker, +) +from vllm.model_executor.layers.attention import mla_attention from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha from vllm.model_executor.layers.attention.sparse_mla_attention import ( GLOBAL_TOPK_MASK_MAX_BYTES, SparseMLACommonImpl, + SparseMLAPrefillMetadata, _masked_mha_workspace_fits, _topk_mask_shape, + _use_dense_mha_prefill, ) from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.platforms import current_platform @@ -41,7 +49,6 @@ if not current_platform.is_cuda(): allow_module_level=True, ) -import vllm.v1.attention.backends.mla.flashinfer_mla_sparse as flashinfer_sparse_mod from vllm.model_executor.layers.attention.mla_attention import ( _canonicalize_sparse_mla_kv_cache_dtype, ) @@ -49,10 +56,19 @@ from vllm.model_executor.layers.attention.sparse_mla_attention import ( SharedTopkIndicesBuffer, ) from vllm.utils.math_utils import cdiv +from vllm.utils.torch_utils import current_stream +from vllm.v1.attention.backends.mla import index_group as index_group_module +from vllm.v1.attention.backends.mla.flashattn_mla_sparse import ( + FlashAttnMLASparseImpl, +) from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseImpl, + FlashInferMLASparseMetadataBuilder, FlashInferMLASparseTRTLLMBackend, ) +from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import ( + FlashInferMLASparseSM120Impl, +) from vllm.v1.attention.backends.mla.flashmla_sparse import ( FlashMLASparseBackend, FlashMLASparseImpl, @@ -60,12 +76,29 @@ from vllm.v1.attention.backends.mla.flashmla_sparse import ( FlashMLASparseMetadataBuilder, triton_convert_req_index_to_global_index, ) -from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadataBuilder +from vllm.v1.attention.backends.mla.index_group import ( + HiSparseMLAIndexGroup, + SparseMLAIndexGroupBuilder, +) +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadataBuilder, +) +from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend from vllm.v1.attention.backends.utils import ( split_decodes_and_prefills, split_prefill_chunks, ) from vllm.v1.attention.ops import flashmla +from vllm.v1.hisparse import runtime as hisparse_runtime +from vllm.v1.hisparse.runtime import ( + HiSparseCacheHandle, + HiSparseRuntime, + ResolvedHiSparseConfig, + _has_hisparse_ops, + build_hisparse_prefill_staging_plan, + hisparse_prefill_staging_remap, +) +from vllm.v1.hisparse.types import SparseKVRowMirror SPARSE_BACKEND_BATCH_SPECS = { name: BATCH_SPECS[name] @@ -113,10 +146,11 @@ def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch): impl.bmm2_scale = None impl.is_nope_mla = True impl.need_to_return_lse_for_decode = False - monkeypatch.setattr( - flashinfer_sparse_mod, - "triton_convert_req_index_to_global_index", - lambda *args, **kwargs: (topk, torch.ones(1, dtype=torch.int32)), + impl.index_group = MagicMock() + impl.index_group_index = 0 + impl.index_group.convert_logical_to_physical_topk.return_value = ( + topk, + torch.ones(1, dtype=torch.int32), ) import flashinfer.decode @@ -142,6 +176,41 @@ def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch): assert recorded_scale != kv_lora_rank**-0.5 +def test_hisparse_routes_prefill_to_sparse_mqa(): + layer = SimpleNamespace(hisparse_cache=object()) + + assert not mla_attention.MLAAttention._use_sparse_mha(layer, SimpleNamespace()) + + +def test_hisparse_metadata_keeps_short_prefill_indexer_enabled(): + config = SimpleNamespace( + attention_config=SimpleNamespace( + hisparse_config=object(), sparse_mla_force_mqa=False + ) + ) + + assert not _use_dense_mha_prefill(config, prefill_max_seq_len=512, topk_tokens=2048) + + +def test_unified_mla_prepares_batch_without_slot_mapping(monkeypatch): + layer = SimpleNamespace( + prepare_kv_cache_update=MagicMock(), + update_kv_cache=MagicMock(), + ) + metadata = object() + monkeypatch.setattr( + mla_attention, + "get_attention_context", + lambda _: (metadata, layer, torch.empty(0), None), + ) + tensor = torch.empty(0, device=DEVICE_TYPE) + + mla_attention.unified_mla_kv_cache_update(tensor, tensor, "layer", "auto", tensor) + + layer.prepare_kv_cache_update.assert_called_once_with(metadata) + layer.update_kv_cache.assert_not_called() + + def _float_to_e8m0_truncate(f: float) -> float: """Simulate SM100's float -> e8m0 -> bf16 scale conversion. e8m0 format only stores the exponent (power of 2). @@ -618,8 +687,6 @@ def test_sparse_backend_decode_correctness( # The sparse builder clones the layer's dense-MHA prefill backend from # static_forward_context; register a mock layer carrying one. - from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend - prefill_backend = get_mla_prefill_backend(vllm_config)( num_heads=num_heads, scale=scale, @@ -972,6 +1039,8 @@ def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens(): ), block_table=torch.arange(40, dtype=torch.int32, device=device).view(4, 10), block_size=block_size, + num_decode_tokens=num_mqa_tokens, + num_actual_tokens=num_mqa_tokens, ) assert attn_metadata.req_id_per_token.shape[0] == num_batch_tokens @@ -990,11 +1059,30 @@ def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens(): def _stub_kernel(q, kv, indices, lengths, actual_num_heads): captured["indices"] = indices captured["actual_num_heads"] = actual_num_heads - return torch.zeros(q.shape[0], q.shape[1], 512, dtype=q.dtype, device=q.device) + return ( + torch.zeros(q.shape[0], q.shape[1], 512, dtype=q.dtype, device=q.device), + None, + ) - stub_impl = SimpleNamespace(_bf16_flash_mla_kernel=_stub_kernel) + def _convert_topk(indices, metadata, **kwargs): + return triton_convert_req_index_to_global_index( + metadata.req_id_per_token[: indices.shape[0]], + metadata.block_table, + indices, + BLOCK_SIZE=metadata.block_size, + BLOCK_STRIDE_ROWS=kwargs["block_stride_rows"], + NUM_TOPK_TOKENS=indices.shape[1], + return_valid_counts=kwargs["return_valid_counts"], + ) - out = FlashMLASparseImpl._forward_bf16_kv( + stub_impl = SimpleNamespace( + _bf16_flash_mla_kernel=_stub_kernel, + _convert_logical_to_physical_topk=_convert_topk, + index_group=None, + index_group_index=0, + ) + + out, _ = FlashMLASparseImpl._forward_bf16_kv( stub_impl, q, kv_cache, topk_indices, attn_metadata, q.shape[1] ) @@ -1283,8 +1371,6 @@ def test_sparse_backend_prefill_correctness( # The sparse builder clones the layer's dense-MHA prefill backend from # static_forward_context; register a mock layer carrying one. - from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend - prefill_backend = get_mla_prefill_backend(vllm_config)( num_heads=num_heads, scale=scale, @@ -1524,6 +1610,1616 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int): assert torch.all(result[row, num_valid:] == -1) +def test_sparse_mla_index_group_converts_decode_indices_once(monkeypatch): + device = torch.device(DEVICE_TYPE) + logical = torch.full((2, 128), -1, dtype=torch.int32, device=device) + logical[:, :4] = torch.tensor( + [[0, 2, 5, -1], [1, 3, 7, -1]], dtype=torch.int32, device=device + ) + builder = SparseMLAIndexGroupBuilder(logical) + group, leader_index = builder.register_layer(True) + follower_group, follower_index = builder.register_layer(False) + assert follower_group is group + + block_table = torch.tensor([[2, 3], [4, 5]], dtype=torch.int32, device=device) + metadata = SimpleNamespace( + block_table=block_table, + block_size=4, + req_id_per_token=torch.tensor([0, 1], dtype=torch.int32, device=device), + num_decode_tokens=2, + num_actual_tokens=2, + ) + calls = [] + convert = index_group_module.triton_convert_req_index_to_global_index + + def convert_spy(*args, **kwargs): + calls.append((args, kwargs)) + return convert(*args, **kwargs) + + monkeypatch.setattr( + index_group_module, + "triton_convert_req_index_to_global_index", + convert_spy, + ) + group.set_logical_topk_ready(leader_index) + leader_result = group.convert_logical_to_physical_topk( + leader_index, + logical, + metadata, + block_stride_rows=None, + return_valid_counts=True, + ) + follower_result = group.convert_logical_to_physical_topk( + follower_index, + logical, + metadata, + block_stride_rows=None, + return_valid_counts=True, + ) + torch.accelerator.synchronize() + + assert len(calls) == 1 + leader_indices, leader_counts = leader_result + follower_indices, follower_counts = follower_result + torch.testing.assert_close(follower_indices, leader_indices) + torch.testing.assert_close(follower_counts, leader_counts) + + +def test_sparse_mla_index_group_falls_back_for_prefill_sized_batch(monkeypatch): + device = torch.device(DEVICE_TYPE) + logical = torch.full((4, 128), -1, dtype=torch.int32, device=device) + logical[:, 0] = torch.tensor([0, 1, 2, 3], dtype=torch.int32, device=device) + builder = SparseMLAIndexGroupBuilder(logical, max_decode_rows=2) + group, leader_index = builder.register_layer(True) + follower_group, follower_index = builder.register_layer(False) + assert follower_group is group + + metadata = SimpleNamespace( + block_table=torch.tensor( + [[2], [3], [4], [5]], dtype=torch.int32, device=device + ), + block_size=4, + req_id_per_token=torch.arange(4, dtype=torch.int32, device=device), + ) + calls = 0 + convert = index_group_module.triton_convert_req_index_to_global_index + + def convert_spy(*args, **kwargs): + nonlocal calls + calls += 1 + return convert(*args, **kwargs) + + monkeypatch.setattr( + index_group_module, + "triton_convert_req_index_to_global_index", + convert_spy, + ) + leader_result = group.convert_logical_to_physical_topk( + leader_index, + logical, + metadata, + block_stride_rows=None, + return_valid_counts=False, + ) + follower_result = group.convert_logical_to_physical_topk( + follower_index, + logical, + metadata, + block_stride_rows=None, + return_valid_counts=False, + ) + torch.accelerator.synchronize() + + assert calls == 2 + torch.testing.assert_close(follower_result, leader_result) + + +def test_sparse_mla_index_groups_own_distinct_physical_buffers(): + logical = torch.empty((8, 128), dtype=torch.int32, device=DEVICE_TYPE) + builder = SparseMLAIndexGroupBuilder(logical, max_decode_rows=2) + + first, _ = builder.register_layer(True) + follower, _ = builder.register_layer(False) + second, _ = builder.register_layer(True) + + assert follower is first + assert second is not first + assert first.physical_topk_indices[:2].shape == (2, 128) + assert second.physical_topk_indices[:2].shape == (2, 128) + assert first.physical_topk_indices.data_ptr() != ( + second.physical_topk_indices.data_ptr() + ) + + +# HiSparse is host-resident-only and kernel-only: runtime construction +# raises without the compiled CUDA ops. +requires_hisparse_ops = pytest.mark.skipif( + not _has_hisparse_ops(), + reason="HiSparse CUDA ops not compiled", +) + + +def _make_hisparse_runtime( + *, + top_k: int = 4, + device_buffer_size: int = 5, + max_num_reqs: int = 2, + row_width: int = 8, + block_size: int = 64, + max_swap_rows: int | None = None, + index_group: hisparse_runtime.HiSparseIndexGroup | None = None, +) -> HiSparseRuntime: + runtime = HiSparseRuntime( + config=ResolvedHiSparseConfig( + top_k=top_k, + device_buffer_size=device_buffer_size, + ), + max_num_reqs=max_num_reqs, + row_width=row_width, + kv_dtype=torch.float32, + device=DEVICE_TYPE, + max_swap_rows=max_swap_rows, + index_group=index_group, + ) + blocks_per_request = cdiv(runtime.region_stride, block_size) + # Leave one extra physical block for resident-tier tests. HMA group layouts + # alias this backing, but distinct group allocations use distinct block IDs. + num_blocks = max_num_reqs * blocks_per_request + 1 + raw = torch.zeros( + num_blocks * block_size * row_width, + dtype=torch.float32, + device=DEVICE_TYPE, + ).view(torch.int8) + block_table = torch.arange( + max_num_reqs * blocks_per_request, + dtype=torch.int32, + device=DEVICE_TYPE, + ).view(max_num_reqs, blocks_per_request) + runtime.bind_hot_cache( + raw, + byte_offset=0, + block_stride=block_size * row_width * torch.float32.itemsize, + num_blocks=num_blocks, + block_size=block_size, + block_table=block_table, + ) + runtime.request_state_indices = torch.arange( + max_num_reqs, dtype=torch.int32, device=runtime.device + ) + return runtime + + +def _make_hisparse_cache_handle( + *, + top_k: int = 4, + device_buffer_size: int = 5, + max_num_reqs: int = 2, + row_width: int = 8, + block_size: int = 64, + max_swap_rows: int | None = None, +) -> HiSparseCacheHandle: + runtime = _make_hisparse_runtime( + top_k=top_k, + device_buffer_size=device_buffer_size, + max_num_reqs=max_num_reqs, + row_width=row_width, + block_size=block_size, + max_swap_rows=max_swap_rows, + ) + return HiSparseCacheHandle(runtime) + + +@requires_hisparse_ops +def test_hisparse_uses_graph_stable_request_state_mapping(): + device = torch.device(DEVICE_TYPE) + block_size, row_width = 64, 8 + runtime = _make_hisparse_runtime( + top_k=1, + device_buffer_size=2, + max_num_reqs=2, + row_width=row_width, + block_size=block_size, + ) + runtime.bind_source_cache( + torch.arange(2 * block_size * row_width, dtype=torch.float32) + .view(2, block_size, row_width) + .pin_memory() + ) + runtime.request_state_indices = torch.tensor([1], dtype=torch.int32, device=device) + + cache_handle = HiSparseCacheHandle(runtime) + cache_handle.all_context_pages_resident = False + runtime.begin_forward() + cache_handle.swap_in( + req_id_per_token=torch.tensor([0], dtype=torch.int32, device=device), + block_table=torch.tensor([[1]], dtype=torch.int32, device=device), + logical_topk_indices=torch.tensor([[0]], dtype=torch.int32, device=device), + block_size=block_size, + ) + torch.accelerator.synchronize() + + assert (runtime.index_group.device_global_indices[0] == -1).all() + assert (runtime.index_group.device_global_indices[1] == block_size).any() + + +@requires_hisparse_ops +def test_hisparse_maps_speculative_rows_through_request_state(): + """Multiple verification rows for one request share its persistent state.""" + device = torch.device(DEVICE_TYPE) + block_size, row_width = 64, 8 + runtime = _make_hisparse_runtime( + top_k=1, + device_buffer_size=2, + max_num_reqs=1, + row_width=row_width, + block_size=block_size, + max_swap_rows=4, + ) + runtime.bind_source_cache( + torch.zeros(1, block_size, row_width, dtype=torch.float32).pin_memory() + ) + + cache_handle = HiSparseCacheHandle(runtime) + resident_table = torch.tensor([[1]], dtype=torch.int32, device=device) + cache_handle.bind_cache( + runtime.hot.cache.view(torch.int8), + byte_offset=0, + block_stride=block_size * row_width * torch.float32.itemsize, + num_blocks=2, + block_size=block_size, + block_table=resident_table, + slot_mapping=torch.arange(4, dtype=torch.int64, device=device), + ) + cache_handle.all_context_pages_resident = True + runtime.begin_forward() + physical = cache_handle.swap_in( + req_id_per_token=torch.zeros(4, dtype=torch.int32, device=device), + block_table=torch.zeros((1, 1), dtype=torch.int32, device=device), + logical_topk_indices=torch.zeros((4, 1), dtype=torch.int32, device=device), + block_size=block_size, + ) + torch.accelerator.synchronize() + + assert physical.tolist() == [[block_size]] * 4 + + +@requires_hisparse_ops +def test_hisparse_resident_rows_bypass_hot_lru(): + device = torch.device(DEVICE_TYPE) + block_size, row_width = 64, 8 + cache_handle = _make_hisparse_cache_handle( + top_k=4, + device_buffer_size=5, + max_num_reqs=1, + row_width=row_width, + block_size=block_size, + ) + raw = cache_handle.runtime.hot.cache.view(torch.int8) + resident_table = torch.tensor([[1]], dtype=torch.int32, device=device) + resident_slots = torch.tensor([block_size], dtype=torch.int64, device=device) + cache_handle.bind_cache( + raw, + byte_offset=0, + block_stride=block_size * row_width * torch.float32.itemsize, + num_blocks=2, + block_size=block_size, + block_table=resident_table, + slot_mapping=resident_slots, + ) + + host = torch.randn(1, block_size, row_width).pin_memory() + cache_handle.runtime.bind_source_cache(host) + resident_rows = host[0].to(device).add(7) + cache_handle.runtime.hot.cache[1].copy_(resident_rows) + source_table = torch.tensor([[5]], dtype=torch.int32, device=device) + request_ids = torch.zeros(1, dtype=torch.int32, device=device) + topk = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32, device=device) + lru_before = cache_handle.runtime.index_group.lru_slots.clone() + + cache_handle.runtime.begin_forward() + indices, valid_counts = cache_handle.swap_in( + req_id_per_token=request_ids, + block_table=source_table, + logical_topk_indices=topk, + block_size=block_size, + return_valid_counts=True, + ) + torch.accelerator.synchronize() + + expected_indices = topk + block_size + torch.testing.assert_close(indices, expected_indices) + torch.testing.assert_close( + valid_counts, torch.tensor([topk.shape[1]], dtype=torch.int32, device=device) + ) + torch.testing.assert_close(cache_handle.runtime.index_group.lru_slots, lru_before) + torch.testing.assert_close( + cache_handle.runtime.hot.attention_cache.view(-1, row_width)[ + indices.to(torch.long) + ], + resident_rows[topk.to(torch.long)], + ) + + +@requires_hisparse_ops +def test_hisparse_swap_in_preserves_rows_across_eviction(): + device = torch.device(DEVICE_TYPE) + block_size = 64 + row_width = 64 + num_blocks = 64 + top_k = 128 + buf = 256 + num_reqs = 4 + + kv_pool = torch.randn( + (num_blocks + 1, block_size, row_width), dtype=torch.float32 + ).pin_memory() + flat_pool = kv_pool.reshape(-1, row_width) + + runtime = _make_hisparse_runtime( + top_k=top_k, + device_buffer_size=buf, + max_num_reqs=num_reqs, + row_width=row_width, + ) + runtime.bind_source_cache(kv_pool) + cache = HiSparseCacheHandle(runtime) + cache.all_context_pages_resident = False + cache.runtime.begin_forward() + + blocks_per_req = num_blocks // num_reqs + block_table = torch.arange( + 1, num_blocks + 1, dtype=torch.int32, device=device + ).view(num_reqs, blocks_per_req) + req_ids = torch.arange(num_reqs, dtype=torch.int32, device=device) + seq_len = blocks_per_req * block_size + base = torch.arange(top_k, dtype=torch.int32, device=device) + for step in range(3): + cache.runtime.begin_forward() + topk = torch.stack( + [(base + step * top_k + row * 17) % seq_len for row in range(num_reqs)] + ) + topk[:, -1] = -1 + + hot_indices, valid_counts = cache.swap_in( + req_id_per_token=req_ids, + block_table=block_table, + logical_topk_indices=topk.clone(), + block_size=block_size, + return_valid_counts=True, + ) + torch.accelerator.synchronize() + + global_ref = _triton_convert_reference_impl( + req_ids, block_table, topk, block_size, top_k + ) + valid = hot_indices >= 0 + torch.testing.assert_close( + valid_counts, + (global_ref >= 0).sum(dim=1, dtype=torch.int32), + rtol=0, + atol=0, + ) + gathered = cache.runtime.hot.attention_cache.reshape(-1, row_width)[ + hot_indices[valid].to(torch.long) + ].cpu() + expected = flat_pool[global_ref[valid].cpu().to(torch.long)] + torch.testing.assert_close(gathered, expected) + + +@requires_hisparse_ops +def test_hisparse_multi_step_swaps_match_independent(): + """The shared LRU must retain every speculative step until followers run.""" + device = torch.device(DEVICE_TYPE) + torch.manual_seed(0) + block_size = 64 + row_width = 64 + num_blocks = 64 + top_k = 128 + buf = 3 * top_k + num_reqs = 4 + + kv_pool = torch.randn( + (num_blocks, block_size, row_width), dtype=torch.float32 + ).pin_memory() + + def make( + index_group: hisparse_runtime.HiSparseIndexGroup | None = None, + ) -> HiSparseCacheHandle: + runtime = _make_hisparse_runtime( + top_k=top_k, + max_num_reqs=num_reqs, + device_buffer_size=buf, + row_width=row_width, + block_size=block_size, + max_swap_rows=2 * num_reqs, + index_group=index_group, + ) + runtime.bind_source_cache(kv_pool) + cache = HiSparseCacheHandle(runtime) + cache.all_context_pages_resident = False + return cache + + blocks_per_req = num_blocks // num_reqs + block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).view( + num_reqs, blocks_per_req + ) + req_ids = torch.arange(num_reqs, dtype=torch.int32, device=device) + seq_len = blocks_per_req * block_size + producer = make() + shared = make(producer.runtime.index_group) + indep = make() + base = torch.arange(top_k, dtype=torch.int32, device=device) + for iteration in range(2): + producer.runtime.begin_forward() + shared.runtime.begin_forward() + indep.runtime.begin_forward() + topks = [ + torch.stack( + [ + (base + (iteration * 2 + step) * top_k + row * 17) % seq_len + for row in range(num_reqs) + ] + ) + for step in range(2) + ] + kw = dict( + req_id_per_token=req_ids, + block_table=block_table, + block_size=block_size, + ) + producer_indices = [ + producer.swap_in(logical_topk_indices=topk.clone(), **kw) for topk in topks + ] + independent_indices = [ + indep.swap_in(logical_topk_indices=topk.clone(), **kw) for topk in topks + ] + follower_indices = [ + shared.swap_in(logical_topk_indices=topk.clone(), **kw) for topk in topks + ] + torch.accelerator.synchronize() + + for idx_full, idx_indep, idx_shared in zip( + producer_indices, independent_indices, follower_indices + ): + torch.testing.assert_close(idx_shared, idx_full, rtol=0, atol=0) + torch.testing.assert_close(idx_indep, idx_full, rtol=0, atol=0) + torch.testing.assert_close(shared.runtime.hot.cache, indep.runtime.hot.cache) + + shared_rows = shared.runtime.hot.attention_cache.reshape(-1, row_width) + host_rows = kv_pool.reshape(-1, row_width) + for physical_rows, logical_rows in zip(follower_indices, topks): + global_rows = _triton_convert_reference_impl( + req_ids, block_table, logical_rows, block_size, top_k + ) + valid = physical_rows >= 0 + torch.testing.assert_close( + shared_rows[physical_rows[valid].to(torch.long)].cpu(), + host_rows[global_rows[valid].cpu().to(torch.long)], + ) + + +@requires_hisparse_ops +def test_hisparse_multi_step_writes_request_major_output(): + device = torch.device(DEVICE_TYPE) + block_size = 64 + row_width = 8 + num_reqs = 2 + query_len = 3 + top_k = 4 + cache = _make_hisparse_cache_handle( + top_k=top_k, + device_buffer_size=query_len * top_k, + max_num_reqs=num_reqs, + row_width=row_width, + block_size=block_size, + max_swap_rows=num_reqs * query_len, + ) + host = torch.arange( + (num_reqs + 1) * block_size * row_width, dtype=torch.float32 + ).view(num_reqs + 1, block_size, row_width) + cache.runtime.bind_source_cache(host.pin_memory()) + block_table = torch.arange(1, num_reqs + 1, dtype=torch.int32, device=device).view( + -1, 1 + ) + request_ids = torch.arange(num_reqs, dtype=torch.int32, device=device) + logical = torch.tensor( + [ + [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], + [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]], + ], + dtype=torch.int32, + device=device, + ) + physical = torch.full_like(logical, -1) + valid_counts = torch.zeros((num_reqs, query_len), dtype=torch.int32, device=device) + + cache.runtime.begin_forward() + for step in range(query_len): + cache.swap_in( + request_ids, + block_table, + logical[:, step], + block_size=block_size, + return_valid_counts=True, + attention_indices_out=physical[:, step], + valid_counts_out=valid_counts[:, step], + ) + torch.accelerator.synchronize() + + torch.testing.assert_close( + valid_counts, + torch.full_like(valid_counts, top_k), + ) + global_rows = _triton_convert_reference_impl( + torch.arange(num_reqs, device=device, dtype=torch.int32).repeat_interleave( + query_len + ), + block_table, + logical.view(num_reqs * query_len, top_k), + block_size, + top_k, + ) + hot_rows = cache.runtime.hot.attention_cache.reshape(-1, row_width) + host_rows = host.reshape(-1, row_width) + torch.testing.assert_close( + hot_rows[physical.view(num_reqs * query_len, top_k).to(torch.long)].cpu(), + host_rows[global_rows.cpu().to(torch.long)], + ) + + +@requires_hisparse_ops +def test_hisparse_kv_update_writes_resident_and_staging_caches(): + device = torch.device(DEVICE_TYPE) + block_size = 4 + row_width = 8 + cache_handle = _make_hisparse_cache_handle( + max_num_reqs=1, + row_width=row_width, + block_size=block_size, + ) + resident_slots = torch.tensor([4, 5, -1], dtype=torch.int64, device=device) + cache_handle.bind_cache( + cache_handle.runtime.hot.cache.view(torch.int8), + byte_offset=0, + block_stride=block_size * row_width * torch.float32.itemsize, + num_blocks=cache_handle.runtime.hot.cache.shape[0], + block_size=block_size, + block_table=torch.tensor([[1]], dtype=torch.int32, device=device), + slot_mapping=resident_slots, + ) + cache_handle.mirror_staging_cache = torch.empty( + (1, block_size, row_width), dtype=torch.float32, device=device + ) + cache_handle.mirror_staging_slots = torch.arange( + block_size, dtype=torch.int64, device=device + ) + slots = torch.tensor([3, 7, -1], dtype=torch.int64, device=device) + kv_c = torch.randn(8, row_width - 2, device=device) + k_pe = torch.randn(8, 1, 2, device=device) + cache_handle.num_actual_tokens = slots.numel() + cache_handle.decode_batch = False + cache_handle.host_mirror_required = True + source_cache = torch.zeros_like(cache_handle.view.cache) + impl = object.__new__(FlashMLASparseImpl) + layer = SimpleNamespace( + hisparse_cache=cache_handle, + use_pcp=False, + impl=impl, + ) + mla_attention.MLAAttention.update_kv_cache( + layer, + kv_c, + k_pe, + source_cache, + slots, + SimpleNamespace(num_decode_tokens=0), + "auto", + torch.tensor(1.0, device=device), + ) + torch.accelerator.synchronize() + + expected = torch.cat([kv_c[:2], k_pe[:2, 0]], dim=-1).cpu() + torch.testing.assert_close( + cache_handle.view.cache.view(-1, row_width)[resident_slots[:2]], + expected.to(device), + ) + torch.testing.assert_close(source_cache, torch.zeros_like(source_cache)) + staged = cache_handle.mirror_staging_cache.view(-1, row_width) + staged_expected = torch.cat([kv_c[:3], k_pe[:3, 0]], dim=-1) + torch.testing.assert_close(staged[:3], staged_expected) + + +@requires_hisparse_ops +def test_hisparse_row_dma_copies_discontiguous_spans_across_layers(): + device = torch.device(DEVICE_TYPE) + block_size = 4 + row_width = 8 + num_rows = 16 + resident_caches = tuple( + ( + torch.arange(num_rows * row_width, dtype=torch.uint8, device=device) + .add_(layer_index * 17) + .view(num_rows // block_size, block_size, row_width) + ) + for layer_index in range(2) + ) + host_caches = tuple( + torch.zeros((num_rows, row_width), dtype=torch.uint8, pin_memory=True) + for _ in resident_caches + ) + worker = object.__new__(HiSparseConnectorWorker) + worker.is_host_writer = True + worker.kernel_block_size = block_size + worker.resident_caches = resident_caches + worker.host_caches = host_caches + worker.cache_handles = [ + SimpleNamespace( + decode_batch=True, + runtime=SimpleNamespace(resident_source_index=0), + ) + for _ in resident_caches + ] + worker.dma_stream = torch.cuda.Stream(device=device) + worker.host_write_event = torch.Event() + worker._dma_free_descriptors = [] + worker._pending_dma_descriptors = deque() + worker._pending_transfer_events = deque() + worker._enqueued_transfer_ids = [] + worker._dma_submitted = False + worker._set_row_mirrors( + ( + SparseKVRowMirror((1,), 5, 2), + SparseKVRowMirror((8,), 10, 4), + ) + ) + + worker._enqueue_row_dma(range(2)) + current_stream().wait_event(worker.host_write_event) + torch.accelerator.synchronize() + + for resident, host in zip(resident_caches, host_caches): + flat_resident = resident.view(num_rows, row_width).cpu() + torch.testing.assert_close(host[5:7], flat_resident[1:3]) + torch.testing.assert_close(host[10:14], flat_resident[8:12]) + + +@requires_hisparse_ops +def test_hisparse_remaps_strided_hma_rows_for_attention(): + device = torch.device(DEVICE_TYPE) + block_size, row_width = 4, 8 + runtime = _make_hisparse_runtime( + max_num_reqs=1, + block_size=block_size, + row_width=row_width, + ) + page_elements = block_size * row_width + stride_elements = 3 * page_elements + byte_offset = page_elements * torch.float32.itemsize + raw = torch.zeros(2 * stride_elements, dtype=torch.float32, device=device).view( + torch.int8 + ) + runtime.bind_hot_cache( + raw, + byte_offset=byte_offset, + block_stride=stride_elements * torch.float32.itemsize, + num_blocks=2, + block_size=block_size, + block_table=torch.tensor([[0, 1]], dtype=torch.int32, device=device), + ) + + source = ( + torch.arange(36 * row_width, dtype=torch.float32) + .view(9, block_size, row_width) + .pin_memory() + ) + runtime.bind_source_cache(source) + cache = HiSparseCacheHandle(runtime) + cache.all_context_pages_resident = False + runtime.begin_forward() + attention_indices = cache.swap_in( + req_id_per_token=torch.tensor([0], dtype=torch.int32, device=device), + block_table=torch.arange(1, 9, dtype=torch.int32, device=device).view(1, 8), + logical_topk_indices=torch.tensor( + [[0, 1, 2, 3]], dtype=torch.int32, device=device + ), + block_size=block_size, + ) + + assert runtime.hot.attention_cache is not None + assert runtime.hot.attention_cache.is_contiguous() + physical_indices = runtime.index_group.shared_topk.device_topk_rows[0] + expected = ( + physical_indices // block_size * (stride_elements // row_width) + + physical_indices % block_size + ) + torch.testing.assert_close(attention_indices[0], expected) + + +@requires_hisparse_ops +def test_hisparse_gather_zeroes_unaligned_destination(): + """Invalid host rows are safely zeroed even at a byte-unaligned address.""" + device = torch.device(DEVICE_TYPE) + row_bytes = 16 + host_cache = torch.ones(1, row_bytes, dtype=torch.uint8).pin_memory() + storage = torch.full((2 * row_bytes + 1,), 7, dtype=torch.uint8, device=device) + hot_cache = torch.as_strided( + storage, + size=(2, 1, row_bytes), + stride=(row_bytes + 1, row_bytes, 1), + ) + global_indices = torch.tensor([[1]], dtype=torch.int32, device=device) + hot_indices = torch.tensor([[1]], dtype=torch.int32, device=device) + miss_mask = torch.ones_like(global_indices) + + torch.ops._C_cache_ops.hisparse_gather_plan( + host_cache, + hot_cache, + global_indices, + hot_indices, + miss_mask, + None, + ) + torch.accelerator.synchronize() + + torch.testing.assert_close(hot_cache[1, 0], torch.zeros_like(hot_cache[1, 0])) + assert storage[row_bytes].item() == 7 + + +@requires_hisparse_ops +def test_hisparse_gather_rejects_short_attention_block_stride(): + """Reject a stride that would let the gather kernel cross block bounds.""" + device = torch.device(DEVICE_TYPE) + host_cache = torch.ones(1, 16, dtype=torch.uint8).pin_memory() + hot_cache = torch.empty((1, 2, 16), dtype=torch.uint8, device=device) + global_indices = torch.zeros((1, 1), dtype=torch.int32, device=device) + hot_indices = torch.zeros_like(global_indices) + miss_mask = torch.ones_like(global_indices) + attention_indices = torch.empty_like(global_indices) + + with pytest.raises(RuntimeError, match="stride must cover one hot block"): + torch.ops._C_cache_ops.hisparse_gather_plan( + host_cache, + hot_cache, + global_indices, + hot_indices, + miss_mask, + None, + attention_indices, + 1, + ) + + +@requires_hisparse_ops +def test_hisparse_newest_write_and_recycled_slot_invalidation(): + """Newest writes clamp padding and invalidated rows are loaded again.""" + device = torch.device(DEVICE_TYPE) + block_size = 4 + row_width = 8 + num_blocks = 8 + + kv_pool = ( + torch.arange(num_blocks * block_size * row_width, dtype=torch.float32) + .view(num_blocks, block_size, row_width) + .pin_memory() + ) + flat_pool = kv_pool.reshape(-1, row_width) + + cache_handle = _make_hisparse_cache_handle(block_size=block_size) + cache_handle.runtime.eager_host_mirror = True + + block_table = torch.tensor([[2, 0, 4]], dtype=torch.int32, device=device) + resident_slots = torch.tensor([4, -1], dtype=torch.int64, device=device) + cache_handle.bind_cache( + cache_handle.runtime.hot.cache.view(torch.int8), + byte_offset=0, + block_stride=block_size * row_width * torch.float32.itemsize, + num_blocks=cache_handle.runtime.hot.cache.shape[0], + block_size=block_size, + block_table=torch.tensor([[0, 0, 1]], dtype=torch.int32, device=device), + slot_mapping=resident_slots, + ) + req_ids = torch.tensor([0], dtype=torch.int32, device=device) + newest_global = 4 * block_size + padded_global = 1 + padded_row = flat_pool[padded_global].clone() + slot_mapping = torch.tensor( + [newest_global, padded_global], dtype=torch.int64, device=device + ) + request_state_indices = torch.tensor([0, -1], dtype=torch.int32, device=device) + cache_handle.runtime.request_state_indices = request_state_indices + + kv_c = torch.randn(3, row_width - 2, device=device) + k_pe = torch.randn(3, 1, 2, device=device) + cache_handle.runtime.bind_source_cache(kv_pool) + cache_handle.num_actual_tokens = slot_mapping.numel() + ops.concat_and_cache_mla( + kv_c[:1], + k_pe[:1, 0], + cache_handle.view.cache, + resident_slots[:1], + kv_cache_dtype="auto", + scale=torch.tensor(1.0, device=device), + ) + torch.accelerator.synchronize() + expected_row = torch.cat([kv_c[0], k_pe[0, 0]]).cpu() + flat_pool[newest_global].copy_(expected_row) + torch.testing.assert_close(flat_pool[newest_global], expected_row) + torch.testing.assert_close(flat_pool[padded_global], padded_row) + flat_hot = cache_handle.runtime.hot.cache.view(-1, row_width) + torch.testing.assert_close(flat_hot[resident_slots[0]].cpu(), expected_row) + + topk = torch.tensor([[0, -1, -1, -1]], dtype=torch.int32, device=device) + cache_handle.all_context_pages_resident = False + cache_handle.runtime.begin_forward() + hot_indices = cache_handle.swap_in( + req_id_per_token=req_ids, + block_table=block_table, + logical_topk_indices=topk, + block_size=block_size, + ) + torch.accelerator.synchronize() + stale_hot_slot = hot_indices[0, 0].item() + stale_row = flat_hot[stale_hot_slot].clone() + + flat_pool[8] += 1000 + cache_handle.runtime.invalidate_written_slots( + torch.tensor([8], dtype=torch.int64, device=device), req_ids + ) + cache_handle.runtime.begin_forward() + hot_indices = cache_handle.swap_in( + req_id_per_token=req_ids, + block_table=block_table, + logical_topk_indices=topk, + block_size=block_size, + ) + torch.accelerator.synchronize() + idx = hot_indices.cpu().tolist()[0][0] + assert not torch.equal(flat_hot[idx], stale_row) + torch.testing.assert_close(flat_hot[idx].cpu(), flat_pool[8]) + + +@requires_hisparse_ops +def test_hisparse_mixed_batch_bf16_row_split( + default_vllm_config, dist_init, workspace_init +): + """Host-resident mixed batch on the bf16 path is row-split. + + Two long-context speculative-decode requests + one short local-prefill + chunk: every decode step must be served from the bounded hot buffer before + it is reused, while only the prefill rows' blocks are staged host->GPU. + """ + ok, reason = flashmla.is_flashmla_sparse_supported() + if not ok: + pytest.skip(reason) + + device = torch.device(DEVICE_TYPE) + dtype = torch.bfloat16 + torch.manual_seed(0) + + num_heads = 64 + kv_lora_rank = 512 + qk_nope_head_dim = 128 + qk_rope_head_dim = 64 + v_head_dim = 128 + head_size = kv_lora_rank + qk_rope_head_dim + topk_tokens = 128 + block_size = 64 + + # Long decode contexts + a short prefill chunk (router shortcut shape). + batch_spec = BatchSpec(seq_lens=[2048, 2048, 192], query_lens=[2, 2, 64]) + max_seqlen = max(batch_spec.seq_lens) + total_cache_tokens = sum(batch_spec.seq_lens) + total_tokens = batch_spec.compute_num_tokens() + + vllm_config = create_vllm_config( + model_name="deepseek-ai/DeepSeek-V2-Lite-Chat", + tensor_parallel_size=1, + max_model_len=max_seqlen, + num_gpu_blocks=max(2048, cdiv(total_cache_tokens, block_size) + 1), + block_size=block_size, + hf_config_override={ + "index_topk": topk_tokens, + "attn_module_list_cfg": [{"topk_tokens": topk_tokens}], + }, + ) + vllm_config.attention_config.hisparse_config = HiSparseConfig( + device_buffer_size=2 * topk_tokens, + ) + vllm_config.speculative_config = SpeculativeConfig( + method="ngram", num_speculative_tokens=1 + ) + model_config = vllm_config.model_config + model_config.hf_text_config = SimpleNamespace( + index_topk=topk_tokens, + q_lora_rank=None, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + model_type="deepseek_v2", + ) + model_config.dtype = dtype + model_config.get_num_attention_heads = MethodType( + lambda self, parallel_config: num_heads, model_config + ) + model_config.get_num_kv_heads = MethodType( + lambda self, parallel_config: 1, model_config + ) + model_config.get_head_size = MethodType(lambda self: head_size, model_config) + model_config.get_sliding_window = MethodType(lambda self: None, model_config) + # create_hisparse_cache_handle sizes hot buffers per layer. + model_config.get_num_layers = MethodType( + lambda self, parallel_config: 1, model_config + ) + + kv_cache_spec = create_standard_kv_cache_spec(vllm_config) + common_attn_metadata = create_common_attn_metadata( + batch_spec, + vllm_config.cache_config.block_size, + device, + arange_block_indices=True, + ) + # Prepopulate every position of every sequence so the forward needs no + # KV-cache update (the row split itself is what is under test). + kv_c_contexts = [ + torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device) + for s_len in batch_spec.seq_lens + ] + k_pe_contexts = [ + torch.rand(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device) + for s_len in batch_spec.seq_lens + ] + kv_cache = create_and_prepopulate_kv_cache( + kv_c_contexts=kv_c_contexts, + k_pe_contexts=k_pe_contexts, + block_size=vllm_config.cache_config.block_size, + head_size=head_size, + dtype=dtype, + device=device, + num_blocks=vllm_config.cache_config.num_gpu_blocks, + common_attn_metadata=common_attn_metadata, + randomize_blocks=False, + kv_cache_dtype="auto", + ) + + prefill_backend = get_mla_prefill_backend(vllm_config)( + num_heads=num_heads, + scale=1.0 / math.sqrt(head_size), + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + vllm_config=vllm_config, + ) + vllm_config.compilation_config.static_forward_context["placeholder"] = ( + SimpleNamespace(prefill_backend=prefill_backend) + ) + + builder_cls = FlashMLASparseBackend.get_builder_cls() + builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device) + metadata = builder.build( + common_prefix_len=0, common_attn_metadata=common_attn_metadata + ) + assert isinstance(metadata.prefill, SparseMLAPrefillMetadata) + + # Per-token sparse indices bounded by each token's position, with unique + # offsets and -1 padding (same construction as the decode parity test). + positions: list[int] = [] + for i in range(batch_spec.batch_size): + ctx_len = batch_spec.seq_lens[i] - batch_spec.query_lens[i] + positions.extend(ctx_len + q_idx for q_idx in range(batch_spec.query_lens[i])) + sparse_indices = torch.empty( + total_tokens, topk_tokens, dtype=torch.int32, device=device + ) + for tok_idx in range(total_tokens): + max_valid_idx = positions[tok_idx] + offset = tok_idx * 7 # Prime number for varied offsets + num_valid = min(topk_tokens // 2, max_valid_idx + 1) + valid_range = torch.arange(num_valid, device=device, dtype=torch.int32) + sparse_indices[tok_idx] = torch.cat( + [ + (valid_range + offset) % (max_valid_idx + 1), + torch.full( + (topk_tokens - num_valid,), -1, device=device, dtype=torch.int32 + ), + ] + ) + + q = torch.rand(total_tokens, num_heads, head_size, dtype=dtype, device=device) + mock_indexer = SimpleNamespace( + topk_indices_buffer=sparse_indices, topk_tokens=topk_tokens + ) + + impl_cls = FlashMLASparseBackend.get_impl_cls() + with set_current_vllm_config(vllm_config): + impl = impl_cls( + num_heads=num_heads, + head_size=head_size, + scale=1.0 / math.sqrt(head_size), + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="auto", + logits_soft_cap=None, + attn_type="decoder", + kv_sharing_target_layer_name=None, + q_lora_rank=None, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + qk_head_dim=qk_nope_head_dim + qk_rope_head_dim, + v_head_dim=v_head_dim, + kv_b_proj=None, + indexer=mock_indexer, + ) + assert isinstance(impl.index_group, HiSparseMLAIndexGroup) + cache_handle = impl.index_group.cache(impl.index_group_index) + blocks_per_request = cdiv(cache_handle.runtime.region_stride, block_size) + num_hot_blocks = cache_handle.runtime.max_num_reqs * blocks_per_request + hot_cache = torch.zeros( + num_hot_blocks * block_size * head_size, + dtype=dtype, + device=device, + ).view(torch.int8) + cache_handle.runtime.bind_hot_cache( + hot_cache, + byte_offset=0, + block_stride=block_size * head_size * dtype.itemsize, + num_blocks=num_hot_blocks, + block_size=block_size, + block_table=torch.arange(num_hot_blocks, dtype=torch.int32, device=device).view( + cache_handle.runtime.max_num_reqs, blocks_per_request + ), + ) + cache_handle.runtime.request_state_indices = torch.arange( + cache_handle.runtime.max_num_reqs, dtype=torch.int32, device=device + ) + impl.prepare_for_batch(metadata) + + # Device-resident reference: the whole batch converted against the full + # block table and run as one kernel call over the GPU cache. + ref_topk, ref_topk_length = triton_convert_req_index_to_global_index( + metadata.req_id_per_token, + metadata.block_table, + sparse_indices, + BLOCK_SIZE=metadata.block_size, + NUM_TOPK_TOKENS=topk_tokens, + return_valid_counts=True, + ) + reference, _ = impl._bf16_flash_mla_kernel(q, kv_cache, ref_topk, ref_topk_length) + + # Host-resident pool with identical contents. + kv_pool = kv_cache.squeeze(1).cpu().pin_memory() + cache_handle.runtime.bind_source_cache(kv_pool) + cache_handle.source_block_table = metadata.block_table + + staging_calls = [] + original_gather = cache_handle.runtime.gather_prefill_cache + + def spy_gather(self, kv, plan, resident_cache=None): + staged = original_gather(kv, plan, resident_cache) + staging_calls.append((plan, staged.shape)) + return staged + + cache_handle.runtime.gather_prefill_cache = MethodType( + spy_gather, cache_handle.runtime + ) + + backend_output, _ = impl._forward_bf16_kv( + q, kv_pool, sparse_indices, metadata, q.shape[1] + ) + torch.accelerator.synchronize() + + # Only the prefill rows' blocks were staged: the decode rows' 2048-token + # contexts (32 blocks each) must stay off the staging gather. + assert len(staging_calls) == 1 + plan, staged_shape = staging_calls[0] + assert plan is metadata.prefill.host_staging_plan + prefill_blocks = cdiv(batch_spec.seq_lens[-1], block_size) + assert staged_shape[0] <= prefill_blocks + 1 # +1: block-0 tail padding + + torch.testing.assert_close(backend_output, reference, rtol=0.01, atol=0.01) + + +def test_hisparse_prefill_staging_remap(): + """Compacted staging references the same host rows as direct indexing.""" + block_size = 4 + block_table = torch.tensor( + [[5, 2, -1, -1], [2, 7, 3, -1], [9, -1, -1, -1]], dtype=torch.int32 + ) + + new_bt, row_ids = hisparse_prefill_staging_remap(block_table, block_size) + + assert new_bt.shape == block_table.shape + assert row_ids.dtype == torch.int32 and new_bt.dtype == torch.int32 + n_unique = len({0, 2, 3, 5, 7, 9}) # -1 clamps to block 0 + flat_rows = row_ids.flatten() + assert (flat_rows >= 0).sum() == n_unique * block_size + for i in range(block_table.shape[0]): + for j in range(block_table.shape[1]): + orig = max(int(block_table[i, j]), 0) + staged = int(new_bt[i, j]) + for k in range(block_size): + assert int(flat_rows[staged * block_size + k]) == orig * block_size + k + + +def test_hisparse_prefill_staging_plan_masks_unused_blocks(): + block_table = torch.tensor([[5, 2, 7], [9, 3, 4]], dtype=torch.int32) + plan = build_hisparse_prefill_staging_plan( + block_table, + seq_lens=torch.tensor([5, 8], dtype=torch.int32), + block_size=4, + staging_block_capacity=4, + ) + + assert (plan.block_table[:, 2] == 0).all() + valid_rows = plan.row_ids[plan.row_ids >= 0] + assert set((valid_rows[::4] // 4).tolist()) == {0, 2, 3, 5, 9} + for row, used in enumerate((2, 2)): + for column in range(used): + staged_block = int(plan.block_table[row, column]) + host_block = int(block_table[row, column]) + assert int(valid_rows[staged_block * 4]) == host_block * 4 + + +def test_hisparse_prefill_staging_plan_resolves_resident_sources(): + """Per-page resident hits become device sources; misses stay host DMAs.""" + block_size = 4 + resident_block_size = 2 + block_table = torch.tensor([[5, 2, 0], [9, 3, 0]], dtype=torch.int32) + seq_lens = torch.tensor([5, 8], dtype=torch.int32) + plan = build_hisparse_prefill_staging_plan( + block_table, seq_lens, block_size, staging_block_capacity=4 + ) + # Two resident pages per host block; 0 entries are null (not resident). + resident_table = torch.tensor( + [[11, 12, 0, 13, 0, 0], [21, 0, 22, 23, 0, 0]], dtype=torch.int32 + ) + + plan.ensure_gpu_sources(resident_table, resident_block_size) + + assert plan.gpu_row_ids is not None + unique_hosts = (plan.row_ids[0].view(-1, block_size)[:, 0] // block_size).tolist() + gpu_rows = plan.gpu_row_ids[0].view(-1, block_size) + reps = {5: (0, 0), 2: (0, 1), 9: (1, 0), 3: (1, 1)} + for u, host_id in enumerate(unique_hosts): + for t in range(block_size): + if host_id < 0: + assert int(gpu_rows[u, t]) == -1 + continue + if host_id == 0: + assert int(gpu_rows[u, t]) == -1 + continue + row, col = reps[host_id] + page = t // resident_block_size + res_block = int(resident_table[row, col * 2 + page]) + expected = ( + res_block * resident_block_size + t % resident_block_size + if res_block > 0 + else -1 + ) + assert int(gpu_rows[u, t]) == expected + valid_rows = plan.row_ids >= 0 + torch.testing.assert_close( + plan.miss_mask, + ((plan.gpu_row_ids < 0) & valid_rows).int(), + check_dtype=False, + ) + + plan.ensure_gpu_sources(torch.zeros_like(resident_table), resident_block_size) + + assert plan.gpu_row_ids is not None + assert (plan.gpu_row_ids == -1).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_hisparse_gather_prefill_cache_prefers_resident_rows(): + """Staged rows come from the resident cache when a shadow page exists.""" + if not _has_hisparse_ops(): + pytest.skip("hisparse CUDA ops unavailable") + device = torch.device("cuda") + block_size, resident_block_size, row_width = 4, 2, 16 + block_table = torch.tensor([[5, 2, 0], [9, 3, 0]], dtype=torch.int32, device=device) + seq_lens = torch.tensor([5, 8], dtype=torch.int32, device=device) + plan = build_hisparse_prefill_staging_plan( + block_table, seq_lens, block_size, staging_block_capacity=4 + ) + resident_table = torch.tensor( + [[11, 12, 0, 13, 0, 0], [21, 0, 22, 23, 0, 0]], + dtype=torch.int32, + device=device, + ) + plan.ensure_gpu_sources(resident_table, resident_block_size) + + num_host_blocks, num_res_blocks = 10, 24 + host_cache = ( + torch.arange(num_host_blocks * block_size * row_width, dtype=torch.float32) + .view(num_host_blocks, block_size, row_width) + .pin_memory() + ) + resident_cache = ( + ( + -torch.arange( + num_res_blocks * resident_block_size * row_width, dtype=torch.float32 + ) + - 1.0 + ) + .view(num_res_blocks, resident_block_size, row_width) + .to(device) + ) + + staged = HiSparseRuntime.gather_prefill_cache( + None, host_cache, plan, resident_cache=resident_cache + ) + + staged_flat = staged.view(-1, row_width).cpu() + host_flat = host_cache.view(-1, row_width) + resident_flat = resident_cache.reshape(-1, row_width).cpu() + gpu_rows = plan.gpu_row_ids[0].cpu() + host_rows = plan.row_ids[0].cpu() + for i in range(staged_flat.shape[0]): + if int(host_rows[i]) < 0: + assert int(plan.miss_mask[0, i]) == 0 + continue + expected = ( + resident_flat[int(gpu_rows[i])] + if int(gpu_rows[i]) >= 0 + else host_flat[int(host_rows[i])] + ) + torch.testing.assert_close(staged_flat[i], expected) + + +def test_hisparse_fp8_decode_resolves_steps_then_runs_batched_attention(monkeypatch): + """FP8 decode must use active dimensions when common metadata is padded.""" + device = torch.device("cpu") + num_decodes = 2 + query_len = 3 + num_tokens = num_decodes * query_len + q = torch.randn(num_tokens, 2, 4, device=device) + topk = ( + torch.arange(num_tokens, dtype=torch.int32, device=device) + .view(-1, 1) + .expand(-1, 4) + ) + steps: list[int] = [] + kernel_shapes: list[torch.Size] = [] + + def swap_in(req_ids, *, logical_topk_indices, **kwargs): # noqa: ARG001 + steps.append(int(logical_topk_indices[0, 0])) + output = kwargs["attention_indices_out"] + if output is not None: + output.copy_(logical_topk_indices + 10) + return topk[:num_decodes] + + def run_kernel(self, *, q, **kwargs): # noqa: ARG001 + kernel_shapes.append(q.shape) + return q[..., :1], None + + runtime = SimpleNamespace( + hot=SimpleNamespace(attention_cache=torch.empty(1, device=device)) + ) + leader_cache = SimpleNamespace( + runtime=runtime, + source_block_table=torch.empty( + num_decodes, 1, dtype=torch.int32, device=device + ), + swap_in=MagicMock(side_effect=swap_in), + ) + follower_cache = SimpleNamespace( + runtime=runtime, + source_block_table=torch.empty( + num_decodes, 1, dtype=torch.int32, device=device + ), + swap_in=MagicMock(side_effect=swap_in), + ) + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.caches = [leader_cache, follower_cache] + index_group.physical_topk_indices = torch.empty( + (num_tokens + 1, 4), dtype=torch.int32, device=device + ) + index_group.valid_topk_counts = torch.empty( + num_tokens + 1, dtype=torch.int32, device=device + ) + index_group.request_ids = torch.arange( + num_decodes, dtype=torch.int32, device=device + ) + impl = SimpleNamespace( + kv_lora_rank=1, + index_group=index_group, + index_group_index=0, + ) + impl._fp8_flash_mla_kernel = MethodType(run_kernel, impl) + metadata = SimpleNamespace( + num_decodes=num_tokens, + num_decode_tokens=num_tokens, + decode_max_query_len=1, + query_start_loc=torch.arange( + 0, + num_tokens + 1, + query_len, + dtype=torch.int32, + device=device, + ), + block_table=torch.empty(num_decodes, 1, dtype=torch.int32, device=device), + block_size=64, + ) + + output = FlashMLASparseImpl._host_backed_fp8_decode( + impl, + q, + topk, + metadata, + SimpleNamespace(), + num_decodes, + query_len, + ) + + assert steps == [0, 1, 2] + assert kernel_shapes == [(num_decodes, query_len, 2, 4)] + assert output.shape == (num_tokens, 2, 1) + torch.testing.assert_close( + index_group.physical_topk_indices[:num_tokens], topk + 10 + ) + + follower_indices = index_group.convert_decode_logical_to_physical_topk( + 1, + topk, + metadata, + return_valid_counts=False, + num_decodes=num_decodes, + decode_query_len=query_len, + ) + assert steps == [0, 1, 2, 0, 1, 2] + torch.testing.assert_close(follower_indices, topk + 10) + + +def test_hisparse_varlen_decode_uses_canonical_request_rows(): + """Graph metadata's query upper bound must preserve packed token rows.""" + num_decodes = 2 + topk = torch.arange(18, dtype=torch.int32).view(9, 2) + + def swap_in(_request_ids, *, logical_topk_indices, **_kwargs): + return logical_topk_indices + 10 + + cache = SimpleNamespace( + source_block_table=torch.zeros((num_decodes, 1), dtype=torch.int32), + swap_in=MagicMock(side_effect=swap_in), + ) + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.caches = [cache] + index_group.physical_topk_indices = torch.empty((10, 2), dtype=torch.int32) + index_group.valid_topk_counts = torch.empty(10, dtype=torch.int32) + index_group.request_ids = torch.arange(num_decodes, dtype=torch.int32) + metadata = SimpleNamespace( + num_decodes=num_decodes, + num_decode_tokens=topk.shape[0], + decode_max_query_len=15, + query_start_loc=torch.tensor([0, 8, 9], dtype=torch.int32), + block_size=64, + ) + + physical_topk = index_group.convert_decode_logical_to_physical_topk( + 0, + topk, + metadata, + return_valid_counts=False, + ) + + torch.testing.assert_close(physical_topk, topk + 10) + request_rows = cache.swap_in.call_args_list[0].args[0] + torch.testing.assert_close(request_rows, index_group.request_ids) + + +def test_flashinfer_hisparse_decode_runs_batched_attention(): + device = torch.device("cpu") + num_tokens = 6 + q = torch.randn(num_tokens, 2, 4, device=device) + topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=device) + physical_topk = topk + 10 + valid_counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=device) + kernel_shapes: list[torch.Size] = [] + + def convert_decode(self, *args, **kwargs): # noqa: ARG001 + return physical_topk, valid_counts + + def prepare_kernel(self, *args, **kwargs): # noqa: ARG001 + pass + + def run_kernel(self, q, cache, indices, counts): # noqa: ARG001 + kernel_shapes.append(q.shape) + return q[..., :1], None + + cache_handle = SimpleNamespace( + runtime=SimpleNamespace( + hot=SimpleNamespace(attention_cache=torch.empty(1, device=device)) + ), + ) + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.caches = [cache_handle] + index_group.convert_decode_logical_to_physical_topk = MethodType( + convert_decode, index_group + ) + impl = object.__new__(FlashInferMLASparseImpl) + impl.topk_indices_buffer = topk + impl.index_group = index_group + impl.index_group_index = 0 + impl._prepare_mqa_kernel = MethodType(prepare_kernel, impl) + impl._run_mqa_kernel = MethodType(run_kernel, impl) + metadata = SimpleNamespace(num_decode_tokens=num_tokens) + + output, lse = FlashInferMLASparseImpl.forward_mqa( + impl, + q, + torch.empty(1, device=device), + metadata, + SimpleNamespace(), + ) + + assert kernel_shapes == [q.shape] + assert output.shape == (num_tokens, 2, 1) + assert lse is None + + +def test_flashattn_hisparse_decode_uses_index_group(): + num_tokens = 4 + q_nope = torch.empty(num_tokens, 2, 3, device=DEVICE_TYPE) + q_rope = torch.empty(num_tokens, 2, 1, device=DEVICE_TYPE) + topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=DEVICE_TYPE) + physical = topk.clone() + counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE) + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.convert_decode_logical_to_physical_topk = MagicMock( + return_value=(physical, counts) + ) + index_group.physical_kv_cache = MagicMock( + return_value=torch.empty(1, device=DEVICE_TYPE) + ) + impl = object.__new__(FlashAttnMLASparseImpl) + impl.topk_indices_buffer = topk + impl.index_group = index_group + impl.index_group_index = 0 + impl._run_mqa_kernel = MagicMock(return_value=q_nope[..., :1]) + metadata = SimpleNamespace(num_decode_tokens=num_tokens, block_size=64) + + output, lse = FlashAttnMLASparseImpl.forward_mqa( + impl, + (q_nope, q_rope), + torch.empty(1, device=DEVICE_TYPE), + metadata, + SimpleNamespace(), + ) + + assert output.shape == (num_tokens, 2, 1) + assert lse is None + index_group.convert_decode_logical_to_physical_topk.assert_called_once() + impl._run_mqa_kernel.assert_called_once() + + +def test_flashinfer_sm120_hisparse_decode_uses_index_group(): + num_tokens = 4 + q = torch.empty(num_tokens, 2, 4, device=DEVICE_TYPE) + topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=DEVICE_TYPE) + physical = topk.clone() + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.convert_decode_logical_to_physical_topk = MagicMock( + return_value=physical + ) + index_group.physical_kv_cache = MagicMock( + return_value=torch.empty(1, device=DEVICE_TYPE) + ) + impl = object.__new__(FlashInferMLASparseSM120Impl) + impl.topk_indices_buffer = topk + impl.index_group = index_group + impl.index_group_index = 0 + impl._run_mqa_kernel = MagicMock(return_value=q[..., :1]) + metadata = SimpleNamespace( + num_decode_tokens=num_tokens, + topk_tokens=topk.shape[1], + ) + + output, lse = FlashInferMLASparseSM120Impl.forward_mqa( + impl, + q, + torch.empty(1, device=DEVICE_TYPE), + metadata, + SimpleNamespace(), + ) + + assert output.shape == (num_tokens, 2, 1) + assert lse is None + index_group.convert_decode_logical_to_physical_topk.assert_called_once() + impl._run_mqa_kernel.assert_called_once() + + +def test_hisparse_resident_prefill_uses_attention_block_stride(): + expected = torch.tensor([[19]], dtype=torch.int32) + cache_handle = SimpleNamespace( + all_context_pages_resident=True, + view=SimpleNamespace(block_size=64, attention_block_stride=832), + block_table=torch.tensor([[3]], dtype=torch.int32), + ) + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.caches = [cache_handle] + index_group.physical_topk_indices = torch.empty((1, 1), dtype=torch.int32) + index_group._convert_once = MagicMock(return_value=expected) + topk = torch.zeros((2, 1), dtype=torch.int32) + metadata = SimpleNamespace( + req_id_per_token=torch.zeros(2, dtype=torch.int32), + ) + + result = index_group.convert_logical_to_physical_topk( + 0, + topk, + metadata, + block_stride_rows=None, + return_valid_counts=False, + ) + + assert result is expected + assert index_group._convert_once.call_args.kwargs["block_stride_rows"] == 832 + + +def test_hisparse_shared_sparse_builder_routes_multi_token_chunks_to_prefill(): + """Default sparse backends must route multi-token chunks to prefill.""" + builder = object.__new__(FlashInferMLASparseMetadataBuilder) + builder.vllm_config = SimpleNamespace( + attention_config=SimpleNamespace(hisparse_config=object()), + speculative_config=SimpleNamespace( + num_speculative_tokens=4, + parallel_drafting=False, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + ) + + builder._init_reorder_batch_threshold( + 1024, + supports_spec_as_decode=True, + supports_dcp_with_varlen=True, + ) + + assert builder.reorder_batch_threshold == 1 + + +def test_hisparse_flashmla_reorders_full_speculative_window_as_decode(): + builder = object.__new__(FlashMLASparseMetadataBuilder) + builder.vllm_config = SimpleNamespace( + attention_config=SimpleNamespace(hisparse_config=object()), + speculative_config=SimpleNamespace( + num_speculative_tokens=3, + parallel_drafting=False, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + ) + + builder._init_reorder_batch_threshold( + 256, + supports_spec_as_decode=True, + supports_dcp_with_varlen=True, + ) + + assert builder.reorder_batch_threshold == 4 + + +def test_hisparse_mixed_single_token_batch_is_not_decode_only(): + cache = SimpleNamespace( + runtime=SimpleNamespace( + begin_forward=lambda: None, + eager_host_mirror=True, + ), + ) + metadata = SimpleNamespace( + num_decode_tokens=1, + num_actual_tokens=2, + max_query_len=1, + num_reqs=2, + req_id_per_token=torch.arange(2), + ) + + HiSparseCacheHandle._prepare_for_batch(cache, metadata) + + assert not cache.decode_batch + + def test_flashmla_cache_dtype_aliases_use_ds_layout(): # kv-cache dtype aliases are canonicalized to fp8_ds_mla before the layer # stores kv_cache_dtype, so they cannot bypass the gate. @@ -1534,41 +3230,6 @@ def test_flashmla_cache_dtype_aliases_use_ds_layout(): ) -def test_flashmla_fp8_metadata_reuses_common_batch_split(): - builder = SimpleNamespace( - device=torch.device(DEVICE_TYPE), - vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=8)), - ) - common_metadata = SimpleNamespace( - num_actual_tokens=1, - seq_lens_cpu_upper_bound=torch.tensor([1]), - query_start_loc_cpu=torch.tensor([0, 1]), - block_table_tensor=torch.zeros(1, 1, dtype=torch.int32, device=DEVICE_TYPE), - ) - metadata = FlashMLASparseMetadata( - num_reqs=1, - max_query_len=1, - max_seq_len=1, - num_actual_tokens=1, - query_start_loc=torch.tensor([0, 1], device=DEVICE_TYPE), - slot_mapping=torch.tensor([0], device=DEVICE_TYPE), - block_table=torch.zeros(1, 1, dtype=torch.int32, device=DEVICE_TYPE), - req_id_per_token=torch.zeros(1, dtype=torch.int32, device=DEVICE_TYPE), - num_decodes=0, - num_prefills=1, - num_decode_tokens=0, - ) - - fp8_metadata = FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode( - builder, common_metadata, metadata - ) - - assert fp8_metadata.num_decodes == 0 - assert fp8_metadata.num_prefills == 1 - assert fp8_metadata.num_decode_tokens == 0 - assert fp8_metadata.num_prefill_tokens == 1 - - def test_flashmla_common_metadata_requires_uniform_decodes(): common_metadata = SimpleNamespace( max_query_len=3, @@ -1687,6 +3348,7 @@ def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: b fp8_extra_metadata=fp8_metadata, fp8_use_mixed_batch=use_mixed_batch, num_actual_tokens=num_batch_tokens, + num_decode_tokens=num_decode_tokens, req_id_per_token=torch.empty( num_batch_tokens, dtype=torch.int32, device=DEVICE_TYPE ), @@ -1698,9 +3360,19 @@ def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: b topk_indices_buffer=topk_indices, num_heads=2, kv_lora_rank=1, + index_group=None, + index_group_index=0, dcp_world_size=1, need_to_return_lse_for_decode=False, _fp8_flash_mla_kernel=run_kernel, + _convert_logical_to_physical_topk=( + lambda indices, metadata, **kwargs: convert_indices( + metadata.req_id_per_token[: indices.shape[0]], + metadata.block_table, + indices, + return_valid_counts=kwargs["return_valid_counts"], + ) + ), ) impl._forward_fp8_kv_mixed_batch = MethodType( FlashMLASparseImpl._forward_fp8_kv_mixed_batch, impl @@ -1856,6 +3528,8 @@ def test_fp8_mixed_batch_dcp_neutralizes_empty_rows(monkeypatch): cp_kv_cache_interleave_size=1, ) impl = SimpleNamespace( + index_group=None, + index_group_index=0, dcp_world_size=2, dcp_rank=0, need_to_return_lse_for_decode=True, @@ -1876,6 +3550,117 @@ def test_fp8_mixed_batch_dcp_neutralizes_empty_rows(monkeypatch): assert not lse.isnan().any() +def test_hisparse_prefill_reuses_builder_staging_plan(): + """Every layer must reuse the batch plan instead of synchronizing to dedupe.""" + plan = SimpleNamespace( + block_table=torch.tensor([[0]], dtype=torch.int32), + ensure_gpu_sources=MagicMock(), + ) + staged = torch.empty((1, 1, 8)) + calls = [] + + def gather(kv_cache, staging_plan, resident_cache=None): + calls.append((kv_cache, staging_plan, resident_cache)) + return staged + + resident_cache = torch.empty((1, 1, 8)) + resident_block_table = torch.tensor([[1]], dtype=torch.int32) + cache = SimpleNamespace( + runtime=SimpleNamespace( + gather_prefill_cache=gather, + ), + view=SimpleNamespace(cache=resident_cache, block_size=1), + block_table=resident_block_table, + ) + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.caches = [cache] + source = torch.empty((1, 1, 8)) + metadata = SimpleNamespace( + num_decodes=0, + num_decode_tokens=0, + seq_lens=torch.tensor([1], dtype=torch.int32), + prefill=SimpleNamespace(host_staging_plan=plan), + req_id_per_token=torch.tensor([0], dtype=torch.int32), + ) + + result, block_table, request_ids = index_group.stage_prefill_rows( + 0, source, metadata + ) + + assert result is staged + assert block_table is plan.block_table + torch.testing.assert_close(request_ids, metadata.req_id_per_token) + plan.ensure_gpu_sources.assert_called_once() + args = plan.ensure_gpu_sources.call_args.args + torch.testing.assert_close(args[0], resident_block_table) + assert args[1] == 1 + assert calls == [(source, plan, resident_cache)] + + +def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch): + compute_stream = object() + prefill_stream = MagicMock() + gather = MagicMock() + monkeypatch.setattr(index_group_module, "current_stream", lambda: compute_stream) + monkeypatch.setattr( + index_group_module.ops, "cp_gather_and_upconvert_fp8_kv_cache", gather + ) + + resident_cache = torch.empty((2, 4, 656), dtype=torch.uint8) + ready = MagicMock() + cache = SimpleNamespace( + view=SimpleNamespace(cache=resident_cache, block_size=4), + block_table=torch.tensor([[2, 3]], dtype=torch.int32), + ) + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.caches = [cache] + index_group.prefill_stream = prefill_stream + index_group.prefill_ready_events = [ready] + source = torch.empty((3, 4, 656), dtype=torch.uint8) + dst = torch.empty((8, 576), dtype=torch.bfloat16) + block_table = torch.tensor([[0, 1]], dtype=torch.int32) + workspace_starts = torch.tensor([0], dtype=torch.int32) + plan = SimpleNamespace( + block_table=torch.tensor([[0, 1]], dtype=torch.int32), + row_ids=torch.arange(8, dtype=torch.int32).view(1, -1), + gpu_row_ids=torch.full((1, 8), -1, dtype=torch.int32), + ensure_gpu_sources=MagicMock(), + ) + metadata = SimpleNamespace( + prefill=SimpleNamespace(host_staging_plan=plan), + num_decodes=0, + ) + + result = index_group.gather_fp8_prefill( + 0, + source, + dst, + block_table, + workspace_starts, + 1, + metadata, + 0, + ) + + prefill_stream.wait_stream.assert_called_once_with(compute_stream) + ready.record.assert_called_once_with(prefill_stream) + assert result is ready + gather.assert_called_once() + args = gather.call_args + assert args.args[0] is resident_cache + assert args.args[1] is dst + torch.testing.assert_close(args.args[2], plan.block_table) + assert args.args[3] is workspace_starts + assert args.args[4] == 1 + assert args.kwargs["host_cache"].data_ptr() == source.data_ptr() + assert args.kwargs["host_row_ids"] is plan.row_ids + assert args.kwargs["device_row_ids"] is plan.gpu_row_ids + plan.ensure_gpu_sources.assert_called_once() + ensure_args = plan.ensure_gpu_sources.call_args.args + torch.testing.assert_close(ensure_args[0], cache.block_table) + assert ensure_args[1] == 4 + + def test_sparse_impl_observes_repointed_indexer_buffer(): """The MTP proposer repoints the draft's indexer at the target model's buffer after the backend impl is built, so the impl must resolve the buffer per read. diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 97b846219f..8012c82f78 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -41,7 +41,9 @@ def test_connector_without_divergent_hit_support_uses_common_lookup(): manager = MagicMock() manager.get_computed_blocks.return_value = (common_blocks, 0, 0) scheduler = SimpleNamespace( - connector=SimpleNamespace(supports_divergent_local_hybrid_hits=False), + connector=SimpleNamespace( + supports_divergent_local_hybrid_hits=False, + ), kv_cache_manager=manager, ) @@ -61,7 +63,9 @@ def test_capable_connector_uses_divergent_partial_hit_lookup(): True, ) scheduler = SimpleNamespace( - connector=SimpleNamespace(supports_divergent_local_hybrid_hits=True), + connector=SimpleNamespace( + supports_divergent_local_hybrid_hits=True, + ), kv_cache_manager=manager, ) @@ -922,71 +926,6 @@ def test_connector_finish_registers_partial_tail_before_cleanup(): scheduler.kv_cache_manager.remove_skipped_blocks.assert_called_once() -def test_block_pool_touch_pins_released_cow_target(): - """The connector can rescue an offered CoW target after its step-scoped - retention is released by using the bound BlockPool's touch method.""" - hash_block_size = 2 - block_size = 2 * hash_block_size - kv_cache_config = KVCacheConfig( - num_blocks=24, - kv_cache_tensors=[], - kv_cache_groups=[ - KVCacheGroupSpec( - ["full"], - FullAttentionSpec( - block_size=hash_block_size, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - ), - ), - KVCacheGroupSpec( - ["mamba"], - MambaSpec( - block_size=block_size, - shapes=(1, 1), - dtypes=(torch.float32,), - mamba_cache_mode="align", - ), - ), - ], - ) - manager = make_kv_cache_manager( - kv_cache_config=kv_cache_config, - max_model_len=8192, - enable_caching=True, - hash_block_size=hash_block_size, - ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) - computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) - assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None - req0.num_computed_tokens = 6 - req0.append_output_token_ids([3]) - assert manager.allocate_slots(req0, 1) is not None - - # Retention released before the drain (defer_block_free=False ordering). - _copies, retained = manager.take_kv_cache_block_copies() - manager.block_pool.free_blocks(retained) - - offloads = drain_boundary_state_offloads(manager) - ((_group_id, block_id, boundary_tokens),) = offloads["0"] - assert boundary_tokens == 6 - cow_block = manager.block_pool.blocks[block_id] - assert cow_block.ref_cnt == 0 - assert block_id in _free_block_ids(manager) - - manager.block_pool.touch([cow_block]) - assert cow_block.ref_cnt == 1 - assert block_id not in _free_block_ids(manager) - - # The connector-pinned block is out of the free queue: draining every free block - # neither trips the allocator's ref_cnt assert nor hands it out. - new_blocks = manager.block_pool.get_new_blocks( - manager.block_pool.get_num_free_blocks() - ) - assert block_id not in {b.block_id for b in new_blocks} - - def test_boundary_state_offload_dropped_when_request_freed_before_drain(): """A hand-off recorded in the same scheduling pass as the request's death must not be drained: its release hook has already run, so draining would diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 0a21e896b0..7ff63d019c 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -13,7 +13,7 @@ from vllm.v1.request import RequestStatus from vllm.v1.structured_output import StructuredOutputGrammar from vllm.v1.utils import ConstantList -from .utils import create_requests, create_scheduler +from .utils import create_requests, create_scheduler, mock_kv pytestmark = pytest.mark.cpu_test @@ -145,6 +145,32 @@ def test_abort(): assert req.num_output_tokens == abort_order_copy.index(i) +def test_connector_metadata_precedes_async_placeholder_advance(monkeypatch): + """Mirroring must see earlier unresolved outputs, excluding the current step.""" + scheduler = create_scheduler( + async_scheduling=True, + use_kv_connector=mock_kv(matched_tokens=0, is_async=False), + ) + (request,) = create_requests(num_requests=1, num_tokens=4, max_tokens=4) + scheduler.add_request(request) + + assert scheduler.connector is not None + build_connector_meta = scheduler.connector.build_connector_meta + observed_placeholders = [] + + def record_placeholders(scheduler_output): + observed_placeholders.append(request.num_output_placeholders) + return build_connector_meta(scheduler_output) + + monkeypatch.setattr( + scheduler.connector, "build_connector_meta", record_placeholders + ) + scheduler.schedule() + scheduler.schedule() + + assert observed_placeholders == [0, 1] + + def test_preempt(): scheduler = create_scheduler(async_scheduling=True) requests = create_requests(num_requests=10, max_tokens=20) diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 2bdb5e74a7..2fcf70c6ef 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -75,6 +75,7 @@ def _mock_vllm_config(layout: str | None): config.cache_config = CacheConfig() config.cache_config.num_gpu_blocks_override = None config.cache_config.kv_cache_layout = layout + config.attention_config.hisparse_config = None return config diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index f0bd296b71..e7eee03e7b 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -22,6 +22,7 @@ from vllm.config import ( SchedulerConfig, VllmConfig, ) +from vllm.config.attention import HiSparseConfig from vllm.config.kv_events import KVEventsConfig from vllm.lora.request import LoRARequest from vllm.multimodal.inputs import ( @@ -52,10 +53,16 @@ from vllm.v1.core.kv_cache_utils import ( make_block_hash_with_group_id, tensor_data, ) +from vllm.v1.hisparse.layout import ( + create_hisparse_layout, + get_hisparse_gpu_memory_usage, +) from vllm.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, FullAttentionSpec, HiddenStateCacheSpec, + HiSparseHotSpec, + HiSparseResidentSpec, KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, @@ -68,18 +75,139 @@ from vllm.v1.kv_cache_interface import ( SinkFullAttentionSpec, SlidingWindowMLASpec, SlidingWindowSpec, + SparseCacheRole, UniformTypeKVCacheSpecs, get_kv_cache_spec_kind, get_kv_cache_spec_sliding_window, is_full_attention_spec, iter_layer_specs, ) +from vllm.v1.kv_cache_layout import KVCacheLayout from vllm.v1.metrics.stats import CachingMetrics, PrefixCacheStats from vllm.v1.request import Request pytestmark = pytest.mark.cpu_test +@pytest.mark.parametrize("gpu_block_size", [32, 64]) +def test_hisparse_hma_uses_resolved_gpu_block_size(monkeypatch, gpu_block_size): + specs = { + "model.layers.0.self_attn": MLAAttentionSpec( + block_size=gpu_block_size, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + is_index_group_leader=True, + ), + "model.layers.0.self_attn.indexer": MLAAttentionSpec( + block_size=gpu_block_size, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + cache_role=SparseCacheRole.INDEXER, + ), + } + group_spec = UniformTypeKVCacheSpecs.from_specs(specs) + assert group_spec is not None + group = KVCacheGroupSpec(list(specs), group_spec) + config = SimpleNamespace( + attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()), + model_config=SimpleNamespace( + hf_config=SimpleNamespace(index_topk=128), + max_model_len=gpu_block_size, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + cache_config=SimpleNamespace( + num_gpu_blocks_override=7, + prefix_cache_retention_interval=None, + get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC, + ), + ) + indexer_spec = specs["model.layers.0.self_attn.indexer"] + assert get_hisparse_gpu_memory_usage(config, [group]) == ( + indexer_spec.max_memory_usage_bytes(config) + ) + + monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30) + cache_config = kv_cache_utils.get_kv_cache_config_from_groups( + config, [group], available_memory=2**30 + ) + assert cache_config.num_blocks == 7 + assert cache_config.hisparse_host_num_blocks is not None + assert cache_config.hisparse_host_num_blocks > 7 + + host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups + assert host_group.host_resident + assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups]) + host_layers = set(host_group.layer_names) + for tensor in cache_config.kv_cache_tensors: + assert all( + (name in host_layers) == tensor.host_resident for name in tensor.layers + ) + host_tensor = next(t for t in cache_config.kv_cache_tensors if t.host_resident) + host_spec = host_group.kv_cache_spec.kv_cache_specs[host_tensor.layers[0]] + assert host_tensor.block_stride == host_spec.page_size_bytes + assert host_tensor.layer_stride == ( + host_spec.page_size_bytes * cache_config.hisparse_host_num_blocks + ) + assert host_group.kv_cache_spec.block_size == gpu_block_size + assert indexer_group.kv_cache_spec.block_size == gpu_block_size + host_specs = host_group.kv_cache_spec.kv_cache_specs + gpu_indexer_specs = indexer_group.kv_cache_spec.kv_cache_specs + assert set(host_specs) == {"model.layers.0.self_attn"} + assert set(gpu_indexer_specs) == {"model.layers.0.self_attn.indexer"} + assert indexer_group.kv_cache_spec.prefix_cacheable + assert host_group.enable_kv_transfer + auxiliary_specs = [group.kv_cache_spec for group in auxiliary_groups] + assert any(isinstance(spec, HiSparseResidentSpec) for spec in auxiliary_specs) + assert any(isinstance(spec, HiSparseHotSpec) for spec in auxiliary_specs) + assert all( + spec.block_size == gpu_block_size + for spec in auxiliary_specs + if isinstance(spec, (HiSparseResidentSpec, HiSparseHotSpec)) + ) + scheduler_block_size, hash_block_size = kv_cache_utils.resolve_kv_cache_block_sizes( + cache_config, + SimpleNamespace( + cache_config=SimpleNamespace( + block_size=16, + enable_prefix_caching=True, + prefix_match_unit=None, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + kv_transfer_config=object(), + ), + ) + assert scheduler_block_size == hash_block_size == gpu_block_size + + +def test_hisparse_rejects_deepseek_v4(): + full_specs = { + "model.layers.0.attn": MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=512, + dtype=torch.uint8, + model_version="deepseek_v4", + ) + } + full_uniform = UniformTypeKVCacheSpecs.from_specs(full_specs) + assert full_uniform is not None + group = KVCacheGroupSpec(list(full_specs), full_uniform) + config = SimpleNamespace( + attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()), + model_config=SimpleNamespace(hf_config=SimpleNamespace(index_topk=512)), + cache_config=SimpleNamespace(num_gpu_blocks_override=7), + ) + + with pytest.raises(ValueError, match="does not support DeepSeek V4"): + create_hisparse_layout( + config, + [group], + host_budget=2**30, + ) + + @pytest.fixture(autouse=True) def _auto_init_hash_fn(request): hash_fn: Callable @@ -3848,6 +3976,22 @@ def test_iter_layer_specs_returns_group_members(): assert list(iter_layer_specs(wrapped)) == [full, mla] +def test_wrapped_mamba_group_requires_block_zeroing(): + mamba = MambaSpec( + block_size=4, + shapes=((4, 1),), + dtypes=(torch.float32,), + ) + wrapped = UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs={"mamba": mamba}) + config = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec(["mamba"], wrapped)], + ) + + assert config.needs_kv_cache_zeroing + + def _spec_decode_grouping_config(method="dspark", model_type=None): """Grouping config with an EAGLE-family speculative method enabled.""" return SimpleNamespace( diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 1760ca14d0..d91260e30a 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -7,6 +7,7 @@ from collections.abc import Callable from dataclasses import replace from math import lcm from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch @@ -14,11 +15,21 @@ import torch import vllm.v1.core.kv_cache_manager as kv_cache_manager import vllm.v1.core.kv_cache_utils as kv_cache_utils from vllm.distributed.kv_events import ( + MEDIUM_CPU, MEDIUM_GPU, AllBlocksCleared, BlockRemoved, BlockStored, ) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( + HiSparseConnector, + HiSparseConnectorScheduler, + HiSparseConnectorWorkerMetadata, +) +from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector +from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import ( + OffloadingConnector, +) from vllm.lora.request import LoRARequest from vllm.multimodal.inputs import ( MultiModalFeatureSpec, @@ -40,16 +51,23 @@ from vllm.v1.core.kv_cache_utils import ( init_none_hash, make_block_hash_with_group_id, ) +from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.hisparse.coordinator import get_hisparse_coordinator from vllm.v1.kv_cache_interface import ( FullAttentionSpec, + HiSparseHotSpec, + HiSparseResidentSpec, KVCacheConfig, + KVCacheGroupRole, KVCacheGroupSpec, KVCacheSpecKind, MambaSpec, MLAAttentionSpec, SlidingWindowSpec, ) +from vllm.v1.outputs import KVConnectorOutput +from vllm.v1.request import RequestStatus pytestmark = pytest.mark.cpu_test @@ -118,7 +136,10 @@ def make_kv_cache_manager(kv_cache_config: KVCacheConfig, **kwargs) -> KVCacheMa kv_cache_config, prefix_cache_retention_interval=kwargs.pop("retention_interval"), ) - return KVCacheManager(kv_cache_config, **kwargs) + manager = KVCacheManager(kv_cache_config, **kwargs) + if kv_cache_config.hisparse_host_num_blocks is not None: + get_hisparse_coordinator(manager) + return manager def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig: @@ -139,6 +160,999 @@ def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig: ) +HISPARSE_BLOCK_SIZE = 16 + + +def make_hisparse_kv_cache_config( + num_blocks: int, + host_num_blocks: int, + *, + transfer_device_cache: bool = False, +) -> KVCacheConfig: + source_spec = FullAttentionSpec( + block_size=HISPARSE_BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ) + groups = [ + KVCacheGroupSpec( + ["source"], + source_spec, + role=KVCacheGroupRole.HISPARSE_SOURCE, + host_resident=True, + ), + KVCacheGroupSpec( + ["indexer"], + source_spec, + enable_kv_transfer=transfer_device_cache, + role=KVCacheGroupRole.HISPARSE_INDEXER, + ), + ] + groups.append( + KVCacheGroupSpec( + ["resident"], + HiSparseResidentSpec( + block_size=HISPARSE_BLOCK_SIZE, + page_size=HISPARSE_BLOCK_SIZE * 4, + ), + enable_kv_transfer=transfer_device_cache, + ) + ) + groups.append( + KVCacheGroupSpec( + ["hot"], + HiSparseHotSpec( + block_size=HISPARSE_BLOCK_SIZE, + page_size=HISPARSE_BLOCK_SIZE * 4, + blocks_per_request=2, + ), + enable_kv_transfer=False, + ) + ) + return KVCacheConfig( + num_blocks=num_blocks, + hisparse_host_num_blocks=host_num_blocks, + kv_cache_tensors=[], + kv_cache_groups=groups, + ) + + +def make_hisparse_kv_cache_manager( + num_blocks: int, + host_num_blocks: int, + *, + max_model_len: int = 128, + max_in_flight_tokens: int | None = None, + enable_caching: bool = False, + **config_kwargs, +) -> KVCacheManager: + config = make_hisparse_kv_cache_config(num_blocks, host_num_blocks, **config_kwargs) + return make_kv_cache_manager( + config, + max_model_len=max_model_len, + max_in_flight_tokens=max_in_flight_tokens, + enable_caching=enable_caching, + hash_block_size=HISPARSE_BLOCK_SIZE, + ) + + +def test_hisparse_builds_dma_row_mirrors_across_pages(): + manager = make_hisparse_kv_cache_manager(32, 16) + request = make_request( + "request", + list(range(2 * HISPARSE_BLOCK_SIZE)), + HISPARSE_BLOCK_SIZE, + sha256, + ) + assert manager.allocate_slots(request, num_new_tokens=32) is not None + coordinator = get_hisparse_coordinator(manager) + resident_blocks = coordinator.resident_managers[0].req_to_blocks[request.request_id] + host_blocks = coordinator.host_manager.req_to_blocks[request.request_id] + + mirrors = coordinator.build_row_mirrors([(request.request_id, 14, 4)]) + + assert [mirror.source_starts for mirror in mirrors] == [ + (resident_blocks[0].block_id * HISPARSE_BLOCK_SIZE + 14,), + (resident_blocks[1].block_id * HISPARSE_BLOCK_SIZE,), + ] + assert [mirror.destination_start for mirror in mirrors] == [ + host_blocks[0].block_id * HISPARSE_BLOCK_SIZE + 14, + host_blocks[1].block_id * HISPARSE_BLOCK_SIZE, + ] + assert [mirror.num_rows for mirror in mirrors] == [2, 2] + + +def test_hisparse_async_speculation_mirrors_uncertain_position_range(): + """Unresolved drafts must not leave gaps in the eager host mirror.""" + coordinator = MagicMock(host_group_id=0) + coordinator.take_block_table_updates.return_value = {} + coordinator.build_offload_command.return_value = None + coordinator.build_row_mirrors.return_value = () + scheduler = HiSparseConnectorScheduler( + async_speculative=True, + draft_kv_lookahead=4, + ) + scheduler.bind_coordinator(coordinator) + request = SimpleNamespace(request_id="request", num_output_placeholders=3) + connector = object.__new__(HiSparseConnector) + connector.connector_scheduler = scheduler + connector.update_state_after_alloc(request, None, 0) + scheduler_output = SimpleNamespace( + block_table_updates=None, + kv_cache_block_copies=None, + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace( + new_block_ids=[], + req_ids=["request"], + num_computed_tokens=[103], + ), + num_scheduled_tokens={"request": 4}, + ) + + scheduler.build_connector_meta(scheduler_output) + + coordinator.build_row_mirrors.assert_called_once_with((("request", 100, 11),)) + + # Acceptance arrives between steps; the connector must read the live count. + request.num_output_placeholders = 1 + scheduler.build_connector_meta(scheduler_output) + coordinator.build_row_mirrors.assert_called_with((("request", 102, 9),)) + + connector.request_finished_all_groups(request, ()) + assert not scheduler.requests + + +def test_hisparse_reports_when_context_is_fully_resident(): + """A page the pool reuses under a host-reading request stops being resident.""" + manager = make_hisparse_kv_cache_manager(32, 16) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None + coordinator = get_hisparse_coordinator(manager) + scheduled = ((request.request_id, len(tokens) - 1, 1),) + assert coordinator.all_context_pages_resident(scheduled) + + for hot_manager in coordinator.hot_managers: + hot_manager.require_hot(request.request_id) + assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None + _publish_hisparse_pages(manager) + pool = manager.block_pool + pool.get_new_blocks(pool.get_num_free_blocks()) + + assert not coordinator.all_context_pages_resident(scheduled) + assert coordinator.take_block_table_updates().keys() == {request.request_id} + + +def test_hisparse_host_prefix_can_be_completed_by_indexer_offload(): + """Keep indexer-only imports host-backed in the resident block table.""" + manager = make_hisparse_kv_cache_manager( + 32, + 16, + enable_caching=True, + ) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers + spill_counts = {spill.transfer_id: 1 for spill in spills} + get_hisparse_coordinator(manager).update_spills(spill_counts, spill_counts) + _, indexer_blocks, _, _ = manager.get_blocks(original.request_id).blocks + evicted_indexer_id = indexer_blocks[2].block_id + manager.free(original) + manager.block_pool.evict_blocks({evicted_indexer_id}) + + resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) + blocks, num_local, _ = manager.get_computed_blocks(resumed) + max_completion = HISPARSE_BLOCK_SIZE + assert num_local == 2 * HISPARSE_BLOCK_SIZE + assert [len(group_blocks) for group_blocks in blocks.blocks] == [3, 2, 0, 0] + + allocated = manager.allocate_slots( + resumed, + num_new_tokens=1, + num_new_computed_tokens=num_local, + new_computed_blocks=blocks, + num_external_computed_tokens=max_completion, + ) + + assert allocated is not None + source, indexer, resident, hot = manager.get_blocks(resumed.request_id).blocks + assert len(source) == len(indexer) == len(resident) == 4 + assert len(hot) == 2 + # The local prefix is adopted from shadow pages (GPU-resident), while the + # externally imported page stays host-backed until its tail allocation. + assert not any(block.is_null for block in resident[:2]) + assert resident[2].is_null + assert not resident[3].is_null + + +def test_hisparse_indexer_offload_is_capped_by_missing_host_prefix(): + manager = make_hisparse_kv_cache_manager( + 32, + 16, + enable_caching=True, + ) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers + spill_counts = {spill.transfer_id: 1 for spill in spills} + get_hisparse_coordinator(manager).update_spills(spill_counts, spill_counts) + host_blocks, _, _, _ = manager.get_blocks(original.request_id).blocks + evicted_host_id = host_blocks[0].block_id + manager.free(original) + get_hisparse_coordinator(manager).host_manager.block_pool.evict_blocks( + {evicted_host_id} + ) + + resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) + _, num_local, _ = manager.get_computed_blocks(resumed) + connector = object.__new__(OffloadingConnector) + connector._kv_cache_config = manager.kv_cache_config + connector.bind_kv_cache_manager(manager) + assert num_local == 0 + assert connector._max_loadable_tokens(resumed, num_local) == 0 + + +@pytest.mark.parametrize("nested", [False, True]) +@pytest.mark.parametrize( + "available_tokens,imported_tokens,host_blocks", + [(0, 0, 2), (16, 16, 2), (64, 32, 3)], +) +def test_connector_completes_partial_prefix_without_importing_missing_host_kv( + nested, available_tokens, imported_tokens, host_blocks +): + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + tokens = list(range(64)) + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(original, num_new_tokens=64) is not None + _publish_hisparse_pages(manager) + host, indexer, _, _ = manager.get_blocks(original.request_id).blocks + missing_indexer = indexer[1].block_id + missing_host = host[0].block_id + manager.free(original) + manager.block_pool.evict_blocks({missing_indexer}) + + defer = True + + def lookup(request, num_computed_tokens, max_num_new_tokens=None): + nonlocal defer + if defer: + defer = False + return None, False + available = available_tokens + if max_num_new_tokens is not None: + available = min(available, max_num_new_tokens) + return available, available > 0 + + connector = object.__new__(OffloadingConnector) + connector._kv_cache_config = manager.kv_cache_config + connector.connector_scheduler = SimpleNamespace(get_num_new_matched_tokens=lookup) + if nested: + for _ in range(2): + outer = object.__new__(MultiConnector) + outer._connectors = (connector,) + outer._requests_to_connector = {} + connector = outer + connector.bind_kv_cache_manager(manager) + scheduler = SimpleNamespace(connector=connector, kv_cache_manager=manager) + resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) + + blocks, local, _, diverged = Scheduler._get_local_prefix_cache_hit( + scheduler, resumed + ) + assert local == 16 and not diverged + assert connector.get_num_new_matched_tokens(resumed, local) == (None, False) + + blocks, local, _, diverged = Scheduler._get_local_prefix_cache_hit( + scheduler, resumed + ) + external, asynchronous = connector.get_num_new_matched_tokens(resumed, local) + assert external == imported_tokens + assert asynchronous == (external > 0) + assert [len(group) for group in blocks.blocks] == [3, 1, 0, 0] + assert ( + manager.allocate_slots( + resumed, + int(external == 0), + delay_cache_blocks=True, + num_new_computed_tokens=local, + new_computed_blocks=blocks, + num_external_computed_tokens=external, + ) + is not None + ) + allocated = manager.get_blocks(resumed.request_id).blocks + assert len(allocated[0]) == host_blocks + assert allocated[2][1].is_null == (external > 0) + manager.free(resumed) + + # A surviving indexer offload must not make absent host KV look computed. + get_hisparse_coordinator(manager).host_manager.block_pool.evict_blocks( + {missing_host} + ) + retry = make_request("retry", tokens, HISPARSE_BLOCK_SIZE, sha256) + _, local, _, _ = Scheduler._get_local_prefix_cache_hit(scheduler, retry) + assert local == 0 + assert connector.get_num_new_matched_tokens(retry, local) == (0, False) + + +def allocate_external_prefix( + manager: KVCacheManager, request: Request, num_tokens: int +) -> KVCacheBlocks | None: + return manager.allocate_slots( + request, + num_new_tokens=0, + num_external_computed_tokens=num_tokens, + delay_cache_blocks=True, + full_sequence_must_fit=True, + ) + + +@pytest.mark.parametrize("enable_caching", [False, True]) +def test_hisparse_keeps_resident_pages_until_hot_buffer_is_allocated(enable_caching): + """Later admissions must not remove a scheduled request's only readable tier.""" + manager = make_hisparse_kv_cache_manager( + 32, 32, max_model_len=256, enable_caching=enable_caching + ) + first = make_request("first", list(range(160)), HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(first, num_new_tokens=128) is not None + _publish_hisparse_pages(manager) + first.num_computed_tokens = 128 + pool = manager.block_pool + held = pool.get_new_blocks(pool.get_num_free_blocks() - 3) + + assert manager.allocate_slots(first, num_new_tokens=1) is not None + second = make_request("second", [1000] * 16, HISPARSE_BLOCK_SIZE, sha256) + manager.allocate_slots(second, num_new_tokens=16) + coordinator = get_hisparse_coordinator(manager) + scheduled = (("first", 128, 1),) + assert coordinator.all_context_pages_resident(scheduled) + assert all(not m.has_hot("first") for m in coordinator.hot_managers) + + # Once another allocation finishes, A can acquire its hot buffer and + # safely release its clean pages for other requests to reuse. + pool.free_blocks(held) + assert manager.allocate_slots(first, num_new_tokens=1) is not None + assert all(m.has_hot("first") for m in coordinator.hot_managers) + pool.get_new_blocks(pool.get_num_free_blocks()) + assert not coordinator.all_context_pages_resident(scheduled) + assert "first" in coordinator.take_block_table_updates() + + +def test_hisparse_full_pool_keeps_pages_pinned_until_preemption(): + """Without room for a hot buffer, admission must defer instead of losing KV.""" + manager = make_hisparse_kv_cache_manager(18, 18, max_model_len=160) + first = make_request("first", list(range(128)), HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(first, num_new_tokens=128) is not None + _publish_hisparse_pages(manager) + assert manager.block_pool.get_num_free_blocks() == 1 + first.num_computed_tokens = 128 + assert manager.allocate_slots(first, num_new_tokens=1) is None + assert get_hisparse_coordinator(manager).all_context_pages_resident( + (("first", 127, 1),) + ) + + manager.free(first) + second = make_request("second", [1000] * 16, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(second, num_new_tokens=16) is not None + + +def test_hisparse_materializes_prefix_without_allocating_hot_blocks(): + """A host prefix becomes visible only when every page is durable. + + Completing a later page first must not expose a prefix with a hole. + """ + manager = make_hisparse_kv_cache_manager( + 32, + 16, + enable_caching=True, + ) + tokens = list(range(2 * HISPARSE_BLOCK_SIZE)) + request = make_request("resident", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None + blocks = manager.get_block_ids(request.request_id) + assert len(blocks[2]) == 2 + assert blocks[3] == [] + + spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers + assert len(spills) == 2 + assert all(spill.after_forward for spill in spills) + resident_blocks = manager.get_blocks(request.request_id).blocks[2] + assert [block.ref_cnt for block in resident_blocks] == [2, 2] + spill_counts = {spill.transfer_id: 1 for spill in spills} + duplicate = make_request( + "duplicate", + list(range(3 * HISPARSE_BLOCK_SIZE)), + HISPARSE_BLOCK_SIZE, + sha256, + ) + _, num_computed, _ = manager.get_computed_blocks(duplicate) + assert num_computed == 0 + + get_hisparse_coordinator(manager).update_spills(spill_counts, {}) + assert [block.ref_cnt for block in resident_blocks] == [2, 2] + _, num_computed, _ = manager.get_computed_blocks(duplicate) + assert num_computed == 0 + + get_hisparse_coordinator(manager).update_spills({}, {spills[1].transfer_id: 1}) + assert [block.ref_cnt for block in resident_blocks] == [2, 1] + _, num_computed, _ = manager.get_computed_blocks(duplicate) + assert num_computed == 0 + + get_hisparse_coordinator(manager).update_spills({}, {spills[0].transfer_id: 1}) + _, num_computed, _ = manager.get_computed_blocks(duplicate) + assert num_computed == 2 * HISPARSE_BLOCK_SIZE + blocks = manager.get_block_ids(request.request_id) + assert len(blocks[2]) == 2 + assert blocks[3] == [] + + +def test_hisparse_materialization_respects_per_step_spill_budget(): + """Prefix publication must not bypass the configured spill batch limit.""" + manager = make_hisparse_kv_cache_manager( + 32, + 16, + enable_caching=True, + ) + coordinator = get_hisparse_coordinator(manager) + coordinator.max_spill_pages = 1 + tokens = list(range(2 * HISPARSE_BLOCK_SIZE)) + request = make_request("bounded", tokens, HISPARSE_BLOCK_SIZE, sha256) + + assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None + first = coordinator.build_offload_command().page_transfers + assert len(first) == 1 + + coordinator.plan_prefix_materialization(request.request_id, len(tokens)) + second = coordinator.build_offload_command().page_transfers + assert len(second) == 1 + assert first[0].transfer_id != second[0].transfer_id + + +def test_hisparse_deferred_free_retains_host_blocks_until_fence(): + """Detaching a request must not make in-flight host pages reusable.""" + manager = make_hisparse_kv_cache_manager(16, 16) + host_pool = get_hisparse_coordinator(manager).get_host_block_pool() + assert host_pool is not None + host_free = host_pool.get_num_free_blocks() + device_free = manager.block_pool.get_num_free_blocks() + request = make_request("free", list(range(16)), HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, 16, delay_cache_blocks=True) is not None + blocks = manager.pop_blocks_for_free(request) + assert host_pool.get_num_free_blocks() < host_free + assert blocks and any(host_pool.blocks[b.block_id] is b for b in blocks) + manager.block_pool.free_blocks(reversed(blocks)) + assert host_pool.get_num_free_blocks() == host_free + assert manager.block_pool.get_num_free_blocks() == device_free + + +def test_hisparse_host_cow_copy_is_drained_without_a_gpu_pool(): + """Host-only copy-on-write work must reach the worker copy queue.""" + manager = make_hisparse_kv_cache_manager(16, 16) + source_manager = manager.coordinator.single_type_managers[0] + source_block = source_manager.block_pool.get_new_blocks(1)[0] + source_manager.req_to_blocks["cow"] = [source_block] + source_manager._partial_hit_reqs["cow"] = (0, source_block) + + new_blocks = source_manager.allocate_new_blocks( + "cow", HISPARSE_BLOCK_SIZE, HISPARSE_BLOCK_SIZE + ) + new_block_ids = manager.take_new_block_ids() + coordinator = get_hisparse_coordinator(manager) + assert manager.take_kv_cache_block_copies() == ([], []) + connector = HiSparseConnectorScheduler(async_speculative=False) + connector.bind_coordinator(coordinator) + output = SchedulerOutput.make_empty() + copies = connector.build_connector_meta(output).host_block_copies + assert output.has_sync_kv_loads + + assert new_blocks and new_block_ids == [] + assert len(copies) == 1 + assert copies[0].src_block_id == source_block.block_id + assert copies[0].dst_block_id == new_blocks[0].block_id + source_manager.free("cow") + assert source_block.ref_cnt > 0 and new_blocks[0].ref_cnt > 0 + assert coordinator.has_pending_work() + connector.update_connector_output( + KVConnectorOutput( + kv_connector_worker_meta=HiSparseConnectorWorkerMetadata( + {}, {}, completed_host_copy_dst_ids=(copies[0].dst_block_id,) + ) + ) + ) + assert source_block.ref_cnt == new_blocks[0].ref_cnt == 0 + assert not coordinator.has_pending_work() + + +def test_hisparse_inflight_host_import_reserves_remaining_gpu_pages(): + """An in-flight host import must reserve its unwritten resident pages.""" + manager = make_hisparse_kv_cache_manager( + 11, + 16, + transfer_device_cache=True, + ) + request = make_request( + "partial-import", + list(range(6 * HISPARSE_BLOCK_SIZE)), + HISPARSE_BLOCK_SIZE, + sha256, + ) + imported_tokens = 4 * HISPARSE_BLOCK_SIZE + + assert allocate_external_prefix(manager, request, imported_tokens) is not None + required = manager.coordinator.get_num_blocks_to_allocate( + request_id=request.request_id, + num_tokens=request.num_tokens, + new_computed_blocks=manager.empty_kv_cache_blocks.blocks, + num_encoder_tokens=0, + total_computed_tokens=imported_tokens, + num_local_computed_tokens=0, + num_tokens_main_model=request.num_tokens, + apply_admission_cap=True, + ) + + assert required == 4 + + +@pytest.mark.parametrize( + "full_sequence_must_fit,host_num_blocks,admitted", + [(True, 2, False), (True, 3, True), (False, 2, False), (False, 3, True)], +) +def test_hisparse_async_admission_requires_only_import_destinations( + full_sequence_must_fit, host_num_blocks, admitted, tmp_path +): + """Imports need host destinations, but do not reserve future prefill pages.""" + from .utils import create_scheduler, mock_kv + + (tmp_path / "config.json").write_text( + '{"architectures": ["OPTForCausalLM"], "model_type": "opt"}' + ) + scheduler = create_scheduler( + model=str(tmp_path), + skip_tokenizer_init=True, + max_model_len=128, + use_kv_connector=mock_kv(matched_tokens=HISPARSE_BLOCK_SIZE, is_async=True), + ) + manager = make_hisparse_kv_cache_manager(32, host_num_blocks) + scheduler.kv_cache_manager = manager + scheduler.kv_cache_config = manager.kv_cache_config + scheduler.scheduler_reserve_full_isl = full_sequence_must_fit + inflight = make_request( + "inflight", list(range(3 * HISPARSE_BLOCK_SIZE)), HISPARSE_BLOCK_SIZE, sha256 + ) + scheduler.add_request(inflight) + scheduler.schedule() + assert inflight in scheduler._inflight_prefills + + request = make_request( + "waiting", list(range(4 * HISPARSE_BLOCK_SIZE)), HISPARSE_BLOCK_SIZE, sha256 + ) + scheduler.add_request(request) + scheduler.schedule() + + assert (request in scheduler._inflight_prefills) == admitted + assert bool(manager.get_blocks(request.request_id).blocks[0]) == admitted + + +@pytest.mark.parametrize( + "evictable,local_tokens,free_blocks,admitted", + [ + (False, HISPARSE_BLOCK_SIZE, 1, True), + (True, HISPARSE_BLOCK_SIZE, 1, False), + (False, HISPARSE_BLOCK_SIZE - 1, 1, False), + (False, HISPARSE_BLOCK_SIZE - 1, 2, True), + ], +) +def test_hisparse_import_capacity_includes_hits_and_cow( + evictable, local_tokens, free_blocks, admitted +): + """Touching an evictable hit and copying a partial hit consume capacity.""" + manager = make_hisparse_kv_cache_manager(32, free_blocks + 2 - int(evictable)) + coordinator = get_hisparse_coordinator(manager) + host_pool = coordinator.get_host_block_pool() + assert host_pool is not None + hit = host_pool.get_new_blocks(1)[0] + if evictable: + host_pool.free_blocks([hit]) + assert host_pool.get_num_free_blocks() == free_blocks + assert ( + coordinator.host_manager.get_num_blocks_to_allocate( + "waiting", 64, [hit], 32, local_tokens, 64 + ) + <= host_pool.get_num_free_blocks() + ) == admitted + + +@pytest.mark.parametrize("enable_caching", [False, True]) +def test_hisparse_host_exhaustion_keeps_gpu_pages_readable(enable_caching): + """Pages without a host destination must survive GPU cache reclamation.""" + manager = make_hisparse_kv_cache_manager(32, 2, enable_caching=enable_caching) + request = make_request("resident", list(range(64)), HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, 64) is not None + coordinator = get_hisparse_coordinator(manager) + host = coordinator.host_manager + assert host is not None + assert [b.is_null for b in host.req_to_blocks[request.request_id]] == [ + False, + True, + True, + True, + ] + _publish_hisparse_pages(manager) + for hot in coordinator.hot_managers: + hot.require_hot(request.request_id) + hot.allocate_new_blocks(request.request_id, 64, 64) + coordinator.update_residency(request.request_id) + for resident in coordinator.resident_managers: + for block in resident.req_to_blocks[request.request_id][1:]: + assert not block.is_null and block.ref_cnt > 0 + assert coordinator.request_states[request.request_id].valid_pages == {0} + pool = coordinator.get_host_block_pool() + assert pool is not None + manager.free(request) + assert pool.get_num_free_blocks() == 1 + + +def test_hisparse_host_cow_takes_priority_over_new_pages(): + """The last host block must hold the private tail, not a fresh page.""" + manager = make_hisparse_kv_cache_manager(32, 3) + coordinator = get_hisparse_coordinator(manager) + host = coordinator.host_manager + assert host is not None + source = host.block_pool.get_new_blocks(1)[0] + host.req_to_blocks["cow"] = [source] + host._partial_hit_reqs["cow"] = (0, source) + + host.allocate_new_blocks("cow", 32, 32) + + private, missing = host.req_to_blocks["cow"] + assert not private.is_null and private is not source + assert missing.is_null + assert host.take_host_cow_copies() == [(source, private)] + + +@pytest.mark.parametrize("failed", [False, True]) +def test_host_receive_completion_without_spill_metadata(failed): + """Only successful external receives make host pages readable.""" + manager = make_hisparse_kv_cache_manager(32, 16) + request = make_request("import", list(range(32)), HISPARSE_BLOCK_SIZE, sha256) + assert allocate_external_prefix(manager, request, 32) is not None + request.num_computed_tokens = 32 + request.status = RequestStatus.WAITING_FOR_REMOTE_KVS + scheduler = HiSparseConnectorScheduler(async_speculative=False) + scheduler.bind_coordinator(get_hisparse_coordinator(manager)) + scheduler.requests[request.request_id] = request + connector = object.__new__(HiSparseConnector) + connector.connector_scheduler = scheduler + composite = object.__new__(MultiConnector) + composite._connectors = (connector,) + composite.update_connector_output( + KVConnectorOutput( + finished_recving={request.request_id}, + failed_recving={request.request_id} if failed else set(), + ) + ) + state = get_hisparse_coordinator(manager).request_states.get(request.request_id) + if failed: + assert state is None + else: + assert state is not None and state.valid_pages == {0, 1} + + +def test_hisparse_recomputation_does_not_complete_failed_host_import(): + """Allocating recomputation is not evidence that imported host KV is valid.""" + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + request = make_request( + "partial-failure", list(range(48)), HISPARSE_BLOCK_SIZE, sha256 + ) + assert allocate_external_prefix(manager, request, 32) is not None + # A partial load failure leaves only the first page computed. The next + # allocation reaches the original import boundary before the worker runs. + request.num_computed_tokens = 16 + manager.cache_blocks(request, 16) + assert manager.allocate_slots(request, 16) is not None + state = get_hisparse_coordinator(manager).request_states.get(request.request_id) + assert state is None or 1 not in state.valid_pages + + +def test_hisparse_host_import_ignores_unsealed_tail(): + """A partial imported page must not become readable from stale host data.""" + manager = make_hisparse_kv_cache_manager(16, 16) + coordinator = get_hisparse_coordinator(manager) + + coordinator.complete_host_import("partial", HISPARSE_BLOCK_SIZE + 1) + + state = coordinator.request_states["partial"] + assert state.valid_pages == {0} + assert state.ready_prefix_pages == 1 + + +def test_hisparse_resident_request_can_grow_without_hot_capacity(): + """Running resident history must not be mistaken for a new host-prefix hit.""" + manager = make_hisparse_kv_cache_manager(8, 16) + request = make_request("resident", list(range(48)), HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, num_new_tokens=32) is not None + request.num_computed_tokens = 32 + + # Three free blocks fit the next indexer/resident pages, but not a hot region. + assert manager.block_pool.get_num_free_blocks() == 3 + assert manager.allocate_slots(request, num_new_tokens=1) is not None + assert manager.get_block_ids(request.request_id)[3] == [] + + +def test_hisparse_capacity_query_does_not_require_hot_blocks(): + """A read-only capacity query must not mutate hot-block requirements.""" + manager = make_hisparse_kv_cache_manager(16, 16) + coordinator = get_hisparse_coordinator(manager) + host_pool = coordinator.get_host_block_pool() + assert host_pool is not None + computed: tuple[list[KVCacheBlock], ...] = ( + [host_pool.blocks[1]], + [], + [], + [], + ) + + manager.coordinator.get_num_blocks_to_allocate( + request_id="query-only", + num_tokens=HISPARSE_BLOCK_SIZE, + new_computed_blocks=computed, + num_encoder_tokens=0, + total_computed_tokens=HISPARSE_BLOCK_SIZE, + num_local_computed_tokens=HISPARSE_BLOCK_SIZE, + num_tokens_main_model=HISPARSE_BLOCK_SIZE, + ) + + assert all( + "query-only" not in hot_manager.hot_required + for hot_manager in coordinator.hot_managers + ) + + +def _publish_hisparse_pages(manager: KVCacheManager) -> None: + """Ack all planned spills so host pages publish to the prefix cache.""" + command = get_hisparse_coordinator(manager).build_offload_command() + counts = {spill.transfer_id: 1 for spill in command.page_transfers} + get_hisparse_coordinator(manager).update_spills(counts, counts) + + +def test_hisparse_events_report_host_and_device_placement(): + """Host publication and eviction must not advertise GPU-resident KV.""" + manager = make_kv_cache_manager( + make_hisparse_kv_cache_config(32, 16), + max_model_len=128, + enable_caching=True, + hash_block_size=HISPARSE_BLOCK_SIZE, + enable_kv_cache_events=True, + ) + request = make_request("request", list(range(64)), HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, 64) is not None + _publish_hisparse_pages(manager) + + stored = [ + event for event in manager.take_events() if isinstance(event, BlockStored) + ] + assert {(event.group_idx, event.medium) for event in stored} == { + (0, MEDIUM_CPU), + (1, MEDIUM_GPU), + } + + host_block = manager.get_blocks(request.request_id).blocks[0][0] + get_hisparse_coordinator(manager).host_manager.block_pool.evict_blocks( + {host_block.block_id} + ) + [removed] = manager.take_events() + assert isinstance(removed, BlockRemoved) + assert removed.medium == MEDIUM_CPU + assert removed.group_idx == 0 + + +def test_mamba_boundary_handoffs_do_not_pin_obsolete_blocks(): + """Draining unused connector offers must not retain old recurrent states.""" + manager = make_kv_cache_manager( + KVCacheConfig( + num_blocks=16, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=16, + shapes=((1,),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ) + ], + ), + max_model_len=128, + enable_caching=True, + hash_block_size=16, + ) + request = make_request("request", list(range(64)), 16, sha256) + old_blocks = [] + for end in (16, 32, 48): + assert manager.allocate_slots(request, 16) is not None + request.num_computed_tokens = end + offers = manager.take_boundary_state_offloads() + [(_, block_id, boundary)] = offers[request.request_id] + assert boundary == end + old_blocks.append(manager.block_pool.blocks[block_id]) + manager.remove_skipped_blocks(request.request_id, end, 64) + manager.new_step_starts() + + assert all(block.ref_cnt == 0 for block in old_blocks[:-1]) + + +def test_hisparse_prefix_hit_adopts_gpu_shadow_pages(): + """A host prefix hit must come back GPU-resident while shadows survive.""" + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + _publish_hisparse_pages(manager) + original_resident_ids = [ + block.block_id for block in manager.get_blocks("original").blocks[2] + ] + manager.free(original) + + resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) + computed, num_computed, _ = manager.get_computed_blocks(resumed) + assert num_computed == 3 * HISPARSE_BLOCK_SIZE + assert manager.allocate_slots( + resumed, + num_new_tokens=len(tokens) - num_computed, + num_new_computed_tokens=num_computed, + new_computed_blocks=computed, + ) + resident_blocks = manager.get_blocks("resumed").blocks[2] + assert [block.block_id for block in resident_blocks[:3]] == ( + original_resident_ids[:3] + ) + assert not any(block.is_null for block in resident_blocks[:3]) + assert get_hisparse_coordinator(manager).all_context_pages_resident( + ((resumed.request_id, num_computed, len(tokens) - num_computed),) + ) + + +def test_hisparse_host_backed_request_accepts_local_prefix_hit(): + """Local group-completion hits need no external host import allocation.""" + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + _publish_hisparse_pages(manager) + manager.free(original) + + resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) + computed, num_computed, _ = manager.get_computed_blocks(resumed) + assert num_computed == 3 * HISPARSE_BLOCK_SIZE + assert manager.allocate_slots( + resumed, + num_new_tokens=len(tokens) - num_computed, + num_new_computed_tokens=num_computed, + new_computed_blocks=computed, + num_external_computed_tokens=0, + ) + + +def test_hisparse_finished_request_leaves_free_copies(): + """Clean pages of a finished request are free-queue blocks, not pins.""" + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + _publish_hisparse_pages(manager) + manager.free(original) + + coordinator = get_hisparse_coordinator(manager) + assert coordinator.copies + pool = manager.block_pool + assert pool.get_num_free_blocks() == pool.num_gpu_blocks - 1 + pool.get_new_blocks(pool.get_num_free_blocks()) + assert not coordinator.copies + assert not coordinator.spills_to_send + + +def test_hisparse_reset_prefix_cache_drops_copies(): + """Reset must forget GPU copies along with the host hashes they mirror.""" + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None + _publish_hisparse_pages(manager) + manager.free(request) + + coordinator = get_hisparse_coordinator(manager) + assert coordinator.copies + assert manager.reset_prefix_cache() + assert not coordinator.copies + + +def test_hisparse_reused_copy_is_not_adopted(): + """A host prefix hit must not adopt a GPU copy the pool has handed out.""" + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + _publish_hisparse_pages(manager) + indexer_blocks = list(manager.get_blocks("original").blocks[1]) + manager.free(original) + pool = manager.block_pool + pool.touch(indexer_blocks) + reused = pool.get_new_blocks(pool.get_num_free_blocks()) + pool.free_blocks(reused) + pool.free_blocks(indexer_blocks) + + resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) + computed, num_computed, _ = manager.get_computed_blocks(resumed) + assert num_computed == 3 * HISPARSE_BLOCK_SIZE + assert manager.allocate_slots( + resumed, + num_new_tokens=len(tokens) - num_computed, + num_new_computed_tokens=num_computed, + new_computed_blocks=computed, + ) + resident_blocks = manager.get_blocks("resumed").blocks[2] + assert all(block.is_null for block in resident_blocks[:3]) + + +def test_hisparse_external_import_uses_hard_gpu_footprint(): + """An external prefix larger than resident capacity must remain admissible. + + This guards the admission boundary where indexer + one writable resident + page + the fixed hot allocation fit, but indexer + full resident history do + not. Requiring one more resident page would silently strand long P/D inputs. + """ + num_prompt_blocks = 4 + num_prompt_tokens = num_prompt_blocks * HISPARSE_BLOCK_SIZE + manager = make_hisparse_kv_cache_manager( + 8, + 9, + transfer_device_cache=True, + ) + request = make_request( + "host-import", + list(range(num_prompt_tokens)), + HISPARSE_BLOCK_SIZE, + sha256, + ) + + allocated = allocate_external_prefix(manager, request, num_prompt_tokens) + + assert allocated is not None + source, indexer, resident, hot = manager.get_blocks(request.request_id).blocks + assert len(source) == num_prompt_blocks + assert len(indexer) == num_prompt_blocks + assert len(resident) == num_prompt_blocks + assert all(block.is_null for block in resident) + assert len(hot) == 2 + assert manager.block_pool.get_num_free_blocks() == 1 + + +def test_hisparse_external_import_survives_capacity_retry(): + manager = make_hisparse_kv_cache_manager( + 9, + 7, + transfer_device_cache=True, + ) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + first = make_request("first", tokens, HISPARSE_BLOCK_SIZE, sha256) + second = make_request("second", tokens, HISPARSE_BLOCK_SIZE, sha256) + + assert allocate_external_prefix(manager, first, len(tokens)) is not None + + assert allocate_external_prefix(manager, second, len(tokens)) is None + + manager.free(first) + assert allocate_external_prefix(manager, second, len(tokens)) is not None + _, _, resident, hot = manager.get_blocks(second.request_id).blocks + assert all(block.is_null for block in resident) + assert len(hot) == 2 + + def make_kv_cache_config_hybrid_model( block_size: int, num_blocks: int, @@ -3874,6 +4888,42 @@ def test_can_fit_full_sequence_swa_cap_admits_long_prompt(): ) +@pytest.mark.parametrize(("reserved_blocks", "fits"), [(4, True), (10, False)]) +def test_full_sequence_admission_reserves_device_blocks_for_immediate_chunk( + reserved_blocks, fits +): + """Reservations constrain the current allocation, not the entire prompt.""" + block_size = 16 + config = KVCacheConfig( + num_blocks=11, # Ten usable blocks after the null block. + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ) + manager = make_kv_cache_manager( + config, max_model_len=128, enable_caching=True, hash_block_size=block_size + ) + request = make_request("chunk", list(range(128)), block_size, sha256) + + blocks = manager.allocate_slots( + request, + block_size, + full_sequence_must_fit=True, + reserved_blocks=reserved_blocks, + ) + + assert (blocks is not None) is fits + + def test_can_fit_full_sequence_full_attention_still_gates_oversized(): """The cap only loosens the SWA group; a prompt that exceeds the full-attention pool capacity must still be rejected.""" @@ -3926,6 +4976,30 @@ def test_can_fit_full_sequence_full_attention_still_gates_oversized(): assert manager.allocate_slots(req, block_size, full_sequence_must_fit=True) is None +def test_can_fit_full_sequence_hisparse_caps_resident_pages(): + manager = make_hisparse_kv_cache_manager( + num_blocks=11, + host_num_blocks=9, + max_model_len=8 * HISPARSE_BLOCK_SIZE, + max_in_flight_tokens=2 * HISPARSE_BLOCK_SIZE, + ) + request = make_request( + "long-hisparse", + list(range(8 * HISPARSE_BLOCK_SIZE)), + HISPARSE_BLOCK_SIZE, + sha256, + ) + + assert ( + manager.allocate_slots( + request, + num_new_tokens=HISPARSE_BLOCK_SIZE, + full_sequence_must_fit=True, + ) + is not None + ) + + def test_cache_hit_local_and_external(): # Regression test for #33775: when a request hits the local prefix cache # in one KV cache group and needs external (connector) blocks in another, @@ -4771,3 +5845,24 @@ def test_swa_shared_prefix_reuse_under_zero_retention(): assert last_req_hit(retention=0, pin=False) == 0 # retention=0 with the pin keeps the junction window -> reuse restored. assert last_req_hit(retention=0, pin=True) == 4 * block_size + + +def test_device_eviction_keeps_host_prefix(): + """Device eviction must not target a same-numbered host block.""" + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + request = make_request("producer", list(range(64)), HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, 64) is not None + _publish_hisparse_pages(manager) + host, indexer, _, _ = manager.get_blocks(request.request_id).blocks + device_block = indexer[0] + host_block = host[0] + assert device_block.block_id == host_block.block_id + manager.free(request) + assert host_block.block_hash is not None and device_block.block_hash is not None + manager.evict_blocks({device_block.block_id}) + assert device_block.block_hash is None, ( + "Device eviction left the indexer prefix cached" + ) + assert host_block.block_hash is not None, ( + "Device eviction invalidated unrelated host KV" + ) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 1c56eaa6a5..3a972e682e 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -19,7 +19,12 @@ from vllm.config import ( SpeculativeConfig, VllmConfig, ) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( + HiSparseConnector, + HiSparseConnectorScheduler, +) from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats +from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalKwargsItem, @@ -34,6 +39,7 @@ from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs from vllm.v1.engine import FinishReason +from vllm.v1.hisparse.coordinator import get_hisparse_coordinator from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, @@ -148,6 +154,30 @@ def test_get_num_unfinished_requests(): assert scheduler.get_num_unfinished_requests() == len(requests) - i - 1 +def _bind_hisparse_connector(scheduler): + connector = object.__new__(HiSparseConnector) + connector.connector_scheduler = HiSparseConnectorScheduler(async_speculative=False) + from tests.v1.core.test_prefix_caching import make_hisparse_kv_cache_manager + + coordinator = get_hisparse_coordinator(make_hisparse_kv_cache_manager(16, 16)) + connector.connector_scheduler.bind_coordinator(coordinator) + composite = object.__new__(MultiConnector) + composite._connectors = (connector,) + scheduler.connector = composite + return coordinator + + +def test_pending_hisparse_spill_keeps_scheduler_alive(): + """A final host spill must complete after the last request finishes.""" + scheduler = create_scheduler() + coordinator = _bind_hisparse_connector(scheduler) + pending = coordinator.pending_spills + pending[0] = Mock() + assert scheduler.has_requests() + pending.clear() + assert not scheduler.has_requests() + + @pytest.mark.parametrize( "enable_prefix_caching, prompt_logprobs", [ @@ -1181,6 +1211,45 @@ def test_preempt_during_execution(): assert requests[1].output_token_ids[0] == 42 +def test_pending_hisparse_reclamation_defers_preemption(monkeypatch): + """Pending reclamation must stall rather than preempt its owning request.""" + scheduler = create_scheduler(enable_prefix_caching=False) + request = create_requests(num_requests=1, num_tokens=16, block_size=16)[0] + scheduler.add_request(request) + scheduler_output = scheduler.schedule() + scheduler.update_from_output( + scheduler_output, + ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[0]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ), + ) + + monkeypatch.setattr( + scheduler.kv_cache_manager, "allocate_slots", Mock(return_value=None) + ) + coordinator = _bind_hisparse_connector(scheduler) + monkeypatch.setattr( + scheduler.connector, "build_connector_meta", Mock(return_value=None) + ) + monkeypatch.setattr( + coordinator, + "has_pending_reclamation", + Mock(return_value=True), + ) + + deferred_output = scheduler.schedule() + + assert not deferred_output.num_scheduled_tokens + assert scheduler.running == [request] + assert request.status == RequestStatus.RUNNING + assert request.num_preemptions == 0 + + def test_prefix_cache_query_not_inflated_by_connector_defer(): """The GPU prefix-cache query is recorded at admission, so a request the connector defers several times is counted once, not once per retry.""" diff --git a/tests/v1/e2e/general/test_hisparse.py b/tests/v1/e2e/general/test_hisparse.py new file mode 100644 index 0000000000..dd4c588657 --- /dev/null +++ b/tests/v1/e2e/general/test_hisparse.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import prometheus_client +import pytest +import torch + +from tests.conftest import VllmRunner +from tests.utils import fork_new_process_for_each_test +from vllm.config import AttentionConfig, HiSparseConfig, KVTransferConfig +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( + HiSparseConnector, +) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( + HiSparseConnectorWorker, +) +from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import ( + MultiConnector, +) +from vllm.platforms import current_platform +from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.v1.hisparse.coordinator import get_hisparse_coordinator + +MODEL = "deepseek-ai/DeepSeek-V3.2" + +HF_OVERRIDES = { + "num_hidden_layers": 8, + "hidden_size": 256, + "intermediate_size": 512, + "num_attention_heads": 8, + "num_key_value_heads": 1, + "n_routed_experts": 8, + "num_experts_per_tok": 2, + "index_topk": 128, +} + + +def _shrink_config(config): + overrides = HF_OVERRIDES.copy() + if any("MTP" in architecture for architecture in config.architectures): + overrides.pop("num_hidden_layers") + config.update(overrides) + return config + + +def _num_hisparse_spills(runner: VllmRunner) -> int: + client = runner.llm.llm_engine.engine_core + coordinator = get_hisparse_coordinator( + client.engine_core.scheduler.kv_cache_manager + ) + return coordinator.next_spill_id + + +def _offload_load_bytes() -> float: + return sum( + sample.value + for metric in prometheus_client.REGISTRY.collect() + if metric.name == "vllm:kv_offload_load_bytes" + for sample in metric.samples + if sample.name == "vllm:kv_offload_load_bytes_total" + ) + + +def _get_hisparse_worker(runner: VllmRunner) -> HiSparseConnectorWorker: + engine_core = runner.llm.llm_engine.engine_core.engine_core + model_runner = engine_core.model_executor.driver_worker.worker.model_runner + connector = model_runner.kv_connector.kv_connector + if isinstance(connector, MultiConnector): + connector = next( + child + for child in connector.sub_connectors + if isinstance(child, HiSparseConnector) + ) + assert isinstance(connector, HiSparseConnector) + assert connector.connector_worker is not None + return connector.connector_worker + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="HiSparse requires NVIDIA CUDA" +) +@pytest.mark.parametrize( + "with_offloading", [False, True], ids=["standalone", "offload"] +) +@fork_new_process_for_each_test +def test_hisparse_spill_and_prefix_restore( + monkeypatch: pytest.MonkeyPatch, + vllm_runner: type[VllmRunner], + with_offloading: bool, +): + """Spilled prefixes restore and FULL-graph decode writes reach host KV. + + A regression omitted attention metadata before FULL graph replay, so decode + completed normally while its newly written KV rows were never copied to host. + """ + capability = current_platform.get_device_capability() + if capability is None or capability.major < 9: + pytest.skip("Sparse MLA requires Hopper or newer") + + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip") + + target = [1000 + i % 64 for i in range(257)] + pressure = [ + [2000 + request_idx * 128 + i % 64 for i in range(257)] + for request_idx in range(4) + ] + hisparse_connector = { + "kv_connector": "HiSparseConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": {"host_pool_gib": 1}, + } + kv_transfer_config = KVTransferConfig(**hisparse_connector) + if with_offloading: + kv_transfer_config = KVTransferConfig( + kv_connector="MultiConnector", + kv_role="kv_both", + kv_connector_extra_config={ + "connectors": [ + hisparse_connector, + { + "kv_connector": "OffloadingConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30}, + }, + ] + }, + ) + + with vllm_runner( + MODEL, + load_format="dummy", + hf_overrides=_shrink_config, + attention_config=AttentionConfig( + hisparse_config=HiSparseConfig( + device_buffer_size=512, + ) + ), + kv_transfer_config=kv_transfer_config, + block_size=64, + max_model_len=320, + max_num_batched_tokens=1024, + max_num_seqs=4, + num_gpu_blocks_override=128, + disable_log_stats=False, + enable_chunked_prefill=True, + enable_prefix_caching=True, + speculative_config={"method": "mtp", "num_speculative_tokens": 3}, + compilation_config={ + "cudagraph_mode": "FULL_AND_PIECEWISE", + "cudagraph_capture_sizes": [1, 4], + }, + ) as runner: + worker = _get_hisparse_worker(runner) + engine_core = runner.llm.llm_engine.engine_core.engine_core + model_runner = engine_core.model_executor.driver_worker.worker.model_runner + full_graph_calls = 0 + full_replay_pending = False + original_run_fullgraph = model_runner.cudagraph_manager.run_fullgraph + + def record_fullgraph(batch_desc): + nonlocal full_graph_calls, full_replay_pending + full_graph_calls += 1 + full_replay_pending = True + return original_run_fullgraph(batch_desc) + + monkeypatch.setattr( + model_runner.cudagraph_manager, "run_fullgraph", record_fullgraph + ) + + original_finish_forward = worker.finish_forward + verified_full_replay_rows = 0 + copied_nonzero_kv = False + + def finish_forward(): + nonlocal full_replay_pending, verified_full_replay_rows + nonlocal copied_nonzero_kv + try: + original_finish_forward() + if not full_replay_pending: + return + # Synchronize only the test's inspection, not finish_forward itself. + with gpu_sync_allowed(): + torch.accelerator.synchronize() + assert worker._row_mirrors + for layer_index, cache in enumerate(worker.cache_handles): + source_index = cache.runtime.resident_source_index + for mirror in worker._row_mirrors: + source_row = mirror.source_starts[source_index] + destination_row = mirror.destination_start + num_rows = mirror.num_rows + source_block, source_offset = divmod( + source_row, worker.kernel_block_size + ) + with gpu_sync_allowed(): + gpu_rows = worker.resident_caches[layer_index][ + source_block, source_offset : source_offset + num_rows + ].cpu() + host_rows = worker.host_caches[layer_index][ + destination_row : destination_row + num_rows + ] + assert torch.equal(host_rows, gpu_rows) + copied_nonzero_kv |= bool(torch.count_nonzero(gpu_rows).item()) + verified_full_replay_rows += num_rows + finally: + full_replay_pending = False + + monkeypatch.setattr(worker, "finish_forward", finish_forward) + + expected = runner.generate_greedy([target], max_tokens=8) + + assert full_graph_calls > 0 + assert verified_full_replay_rows > 0 + assert copied_nonzero_kv + + runner.generate_greedy(pressure, max_tokens=8) + + assert _num_hisparse_spills(runner) > 0 + load_bytes = _offload_load_bytes() + actual = runner.generate_greedy([target], max_tokens=8) + runner.generate_greedy([[42]], max_tokens=1) + + assert actual == expected + if with_offloading: + assert _offload_load_bytes() > load_bytes diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index cd5356031e..10464d950d 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -21,6 +21,7 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, HiddenStateCacheSpec, KVCacheConfig, + KVCacheGroupRole, KVCacheGroupSpec, KVCacheSpec, KVCacheTensor, @@ -336,6 +337,36 @@ def test_worker_kv_bytes_preserves_tensor_layout(packed: bool): assert offloading_config.cache.blocks_per_chunk == 2 +def test_hisparse_offloads_only_indexer_group(): + source = KVCacheGroupSpec( + ["source"], + _full_attention_spec(), + host_resident=True, + ) + indexer = KVCacheGroupSpec( + ["indexer"], + _full_attention_spec(), + role=KVCacheGroupRole.HISPARSE_INDEXER, + ) + kv_cache_config = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[source, indexer], + hisparse_host_num_blocks=4, + ) + config = _make_vllm_config() + + offloading_config = build_offloading_config(config, kv_cache_config) + scheduler_config = SchedulerOffloadConfig.from_spec( + MockOffloadingSpec(offloading_config), config, kv_cache_config + ) + + assert [ + (group.group_id, group.layer_names) for group in offloading_config.groups + ] == [(1, ("indexer",))] + assert [group.group_idx for group in scheduler_config.kv_group_configs] == [1] + + def test_zero_blocks_skips_tensor_layout_validation(): kv_cache_config = _make_sizing_kv_cache_config(packed=False) kv_cache_config.num_blocks = 0 @@ -792,7 +823,6 @@ def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]): False, id="uniform-uncertifiable-inner", ), - pytest.param([], False, id="attention-free"), ], ) def test_canonical_layout_gate(kv_cache_groups, certified): diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 47fb7103d7..26664e3ac5 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -98,6 +98,14 @@ def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state): ) manager = CPUOffloadingManager(num_chunks=100) spec = SimpleNamespace( + config=SimpleNamespace( + groups=tuple( + SimpleNamespace( + group_id=i, tokens_per_block=group.kv_cache_spec.block_size + ) + for i, group in enumerate(groups) + ) + ), tokens_per_block=(256, 64, 8), tokens_per_hash=8, blocks_per_chunk=4, @@ -303,7 +311,24 @@ def test_aligned_boundary_store_uses_exact_source_with_partial_tail(): assert src_spec.block_indices == [0, 0] -def test_aligned_boundary_store_flushes_before_cow_destination_reuse(): +def test_aligned_boundary_store_maps_sparse_group_id_to_dense_transfer_slot(): + scheduler = _make_partial_tail_scheduler() + indexer_config = scheduler.config.kv_group_configs[1] + scheduler.config = scheduler.config._replace(kv_group_configs=(indexer_config,)) + _make_partial_tail_request(scheduler) + scheduler.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + + jobs = scheduler._build_aligned_boundary_store_jobs({"req": [(1, 98, 16)]}) + + [job] = jobs.values() + assert job.src_spec.group_sizes == [1] + assert job.src_spec.block_indices == [0] + + +@pytest.mark.parametrize("cow_reuse", [False, True]) +def test_aligned_boundary_store_flushes_before_block_reuse(cow_reuse): scheduler = _make_partial_tail_scheduler() _make_partial_tail_request(scheduler) scheduler.manager.prepare_store.side_effect = lambda keys, req_context: ( @@ -320,7 +345,11 @@ def test_aligned_boundary_store_flushes_before_cow_destination_reuse(): [job_id] = meta.store_jobs output = SchedulerOutput.make_empty() - output.kv_cache_block_copies = [KVCacheBlockCopy(98, 99)] + if cow_reuse: + output.kv_cache_block_copies = [KVCacheBlockCopy(98, 99)] + else: + output.scheduled_cached_reqs.req_ids = ["req"] + output.scheduled_cached_reqs.new_block_ids = [([], [99])] meta = scheduler.build_connector_meta(output) assert meta.jobs_to_flush == {job_id} @@ -384,6 +413,20 @@ def test_partial_lookup_returns_exact_boundary_and_group_load_keys(): assert req_status.partial_tail_boundary is None +def test_lookup_cap_stops_at_authoritative_prefix_boundary(): + scheduler = _make_partial_tail_scheduler() + request = _make_partial_tail_request(scheduler) + request.skip_reading_prefix_cache = False + scheduler.manager.lookup.return_value = LookupResult.HIT + + tokens, load_async = scheduler.get_num_new_matched_tokens( + request, 0, max_num_new_tokens=20 + ) + + assert (tokens, load_async) == (20, True) + assert scheduler._req_status["req"].partial_tail_boundary == 20 + + def test_recurrent_group_unhashed_block_does_not_truncate_load_boundary(): """A recurrent group may hold unhashed blocks inside the computed prefix. diff --git a/tests/v1/kv_connector/unit/test_hisparse_connector.py b/tests/v1/kv_connector/unit/test_hisparse_connector.py new file mode 100644 index 0000000000..6c4e79efbe --- /dev/null +++ b/tests/v1/kv_connector/unit/test_hisparse_connector.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole +from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import ( + ExampleConnector, +) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( + HiSparseConnector, + HiSparseConnectorMetadata, + HiSparseConnectorScheduler, +) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( + HiSparseConnectorWorker, +) +from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector +from vllm.v1.hisparse.coordinator import get_hisparse_coordinator +from vllm.v1.hisparse.runtime import HiSparseCacheHandle +from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector + + +@pytest.mark.parametrize("nested", [False, True]) +def test_cache_manager_binding_preserves_hisparse_and_legacy_pool_hooks(nested): + """Composite binding must reach HiSparse and existing pool-only connectors.""" + hisparse = object.__new__(HiSparseConnector) + hisparse._role = KVConnectorRole.SCHEDULER + hisparse.connector_scheduler = HiSparseConnectorScheduler(async_speculative=False) + legacy = object.__new__(ExampleConnector) + legacy.bind_gpu_block_pool = MagicMock() + connector = object.__new__(MultiConnector) + connector._connectors = [hisparse, legacy] + if nested: + parent = object.__new__(MultiConnector) + parent._connectors = [connector] + connector = parent + from tests.v1.core.test_prefix_caching import make_hisparse_kv_cache_manager + + manager = make_hisparse_kv_cache_manager(16, 16) + + connector.bind_kv_cache_manager(manager) + + assert hisparse.connector_scheduler.coordinator is get_hisparse_coordinator(manager) + legacy.bind_gpu_block_pool.assert_called_once_with(manager.block_pool) + + +def test_hisparse_requires_block_outermost_device_layout(): + assert HiSparseConnector.get_required_kvcache_layout(MagicMock()) == "BLHNC" + + +def test_no_forward_enqueues_deferred_hisparse_transfers(): + """A zero-token step must still enqueue deferred post-forward transfers.""" + connector = object.__new__(ActiveKVConnector) + connector._disabled = False + connector.pre_forward = MagicMock() + connector.finish_forward = MagicMock() + connector.post_forward = MagicMock(return_value=None) + + scheduler_output = SimpleNamespace(finished_req_ids=set()) + connector.no_forward(scheduler_output) + + connector.pre_forward.assert_called_once_with(scheduler_output) + connector.finish_forward.assert_called_once_with() + + +def test_full_graph_step_prepares_host_mirror_outside_model(): + """Graph replay must restore host-mirror state cleared at step start.""" + runtime = SimpleNamespace( + is_group_leader=True, + eager_host_mirror=True, + begin_forward=MagicMock(), + invalidate_written_slots=MagicMock(), + ) + handle = HiSparseCacheHandle(runtime) + handle.mirror_slot_mapping = torch.tensor([4, 5]) + worker = object.__new__(HiSparseConnectorWorker) + worker.cache_layer_names = ["layer"] + worker.cache_handles = [handle] + worker._group_leaders = (("layer", handle),) + worker._per_layer_mirrored = set() + worker._submitted_mirror_layers = set() + worker.is_host_writer = True + worker._enqueue_row_dma = MagicMock() + worker.start_step = MagicMock( + side_effect=lambda *_args, **_kwargs: worker._clear_forward_mirror_state() + ) + + connector = object.__new__(HiSparseConnector) + connector.connector_worker = worker + connector._get_connector_metadata = MagicMock( + return_value=HiSparseConnectorMetadata(None, (), (), {}, True) + ) + req_id_per_token = torch.tensor([0, 1], dtype=torch.int32) + attn_metadata = SimpleNamespace( + num_actual_tokens=2, + num_decode_tokens=2, + num_reqs=2, + max_query_len=1, + req_id_per_token=req_id_per_token, + ) + + connector.start_load_kv( + SimpleNamespace(), + request_state_indices=None, + request_ids=[], + attn_metadata={"layer": attn_metadata}, + ) + + worker._enqueue_host_mirror() + + worker._enqueue_row_dma.assert_called_once_with((0,), ready_event=None) + runtime.invalidate_written_slots.assert_called_once() diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 52df7531e7..f613ba720e 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -3297,7 +3297,9 @@ def _make_bare_worker( # path; everything flows through the coordinator). from vllm.v1.kv_cache_interface import ( FullAttentionSpec, + KVCacheConfig, KVCacheGroupSpec, + KVCacheTensor, ) worker.disk_offload_buffer_budget_bytes = None @@ -3311,6 +3313,18 @@ def _make_bare_worker( block_size=block_size, num_kv_heads=8, head_size=64, dtype=None ) group = KVCacheGroupSpec(["layer0", "__cross_layer__"], spec) + worker._kv_cache_config = KVCacheConfig( + num_blocks=num_gpu_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=0, + layers=["layer0", "__cross_layer__"], + layer_stride=0, + block_stride=0, + ) + ], + kv_cache_groups=[group], + ) worker._kv_cache_groups = [group] worker.pcp_size = 1 worker.dcp_size = 1 @@ -4059,6 +4073,77 @@ def test_register_kv_caches_shared_storage(layout: KVCacheLayout): worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes) +def test_register_kv_caches_uses_transfer_group_memory_domain(): + """Derived GPU caches must not change host-source transfer addresses.""" + from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupRole, + KVCacheGroupSpec, + KVCacheTensor, + ) + + host_num_blocks = 4 + gpu_num_blocks = 2 + page_size = 32 + source_spec = FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float16 + ) + indexer_spec = FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float16 + ) + source_group = KVCacheGroupSpec( + ["source"], + source_spec, + host_resident=True, + ) + indexer_group = KVCacheGroupSpec( + ["indexer"], + indexer_spec, + enable_kv_transfer=False, + role=KVCacheGroupRole.HISPARSE_INDEXER, + ) + config = KVCacheConfig( + num_blocks=gpu_num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=host_num_blocks * page_size, + layers=["source"], + layer_stride=host_num_blocks * page_size, + block_stride=page_size, + host_resident=True, + ), + KVCacheTensor( + size=gpu_num_blocks * page_size, + layers=["indexer"], + layer_stride=gpu_num_blocks * page_size, + block_stride=page_size, + ), + ], + kv_cache_groups=[source_group, indexer_group], + hisparse_host_num_blocks=host_num_blocks, + ) + worker = _make_bare_worker(num_gpu_blocks=gpu_num_blocks) + worker._kv_cache_config = config + worker._kv_cache_groups = list(config.transfer_groups) + worker.token_dbs = [ + ChunkedTokenDatabase(KeyMetadata("test-model", 0, 0, 0, 0), block_size=16) + ] + source = torch.zeros(host_num_blocks, page_size, dtype=torch.uint8) + indexer = torch.zeros(gpu_num_blocks, page_size, dtype=torch.uint8) + + _register_with_mocked_threads( + worker, + {"source": source, "indexer": indexer}, + ) + + db = worker.token_dbs[0] + assert db.kv_caches_base_addr == [source.data_ptr()] + assert db.block_len == [page_size] + worker.store.register_buffer.assert_called_once_with( + source.data_ptr(), source.nbytes + ) + + def test_register_kv_caches_separate_head_groups(): # LHBNC gives each head group its own region; the K/V split doubles heads. layout = KVCacheLayout.LHBNC diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index bac0b03738..737d03ac69 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -304,10 +304,10 @@ def test_multi_example_connector_consistency(): events = get_connector_events() storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"]) storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"]) - # Initial events bind the block pool, query completion counts, and exchange + # Initial events bind the cache manager, query completion counts, and exchange # handshake metadata before the request is enqueued. assert storage1_scheduler_events[:7] == [ - "bind_gpu_block_pool", + "bind_kv_cache_manager", "get_finished_count", "set_xfer_handshake_metadata_pp_aware", "on_new_request", @@ -315,8 +315,8 @@ def test_multi_example_connector_consistency(): "update_state_after_alloc num_blocks=[7] 0", "build_connector_meta", ] - # First three events are from initialization. During generate(), layer hooks - # run before the deferred load is started after the forward pass. + # First three events are from initialization. Layer hooks run before the + # deferred load starts after the forward pass. expected_worker_prefix = [ "register_kv_caches", "set_host_xfer_buffer_ops", @@ -328,12 +328,12 @@ def test_multi_example_connector_consistency(): ] for connector_name in ("storage1-WORKER", "storage2-WORKER"): worker_events = events[connector_name] - assert worker_events[:7] == expected_worker_prefix + assert worker_events[: len(expected_worker_prefix)] == expected_worker_prefix assert worker_events.index("start_load_kv") > worker_events.index( "save_kv_layer" ) assert storage2_scheduler_events[:7] == [ - "bind_gpu_block_pool", + "bind_kv_cache_manager", "get_finished_count", "set_xfer_handshake_metadata_pp_aware", "on_new_request", diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 282458d1a9..3fcbda4c62 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -29,7 +29,10 @@ from vllm.distributed.kv_transfer.kv_connector.utils import ( get_current_attn_backend, ) from vllm.distributed.kv_transfer.kv_connector.v1 import nixl -from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorRole, + KVConnectorTransferResults, +) from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import ( MultiKVConnectorStats, @@ -518,6 +521,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): slot_size_bytes = 4096 self.slot_size_per_layer = [slot_size_bytes] self.block_len_per_layer = [slot_size_bytes * self.block_size] + self.num_blocks = self.kv_cache_config.num_blocks self.num_regions = 1 self.block_stride_per_layer = list(self.block_len_per_layer) self.region_num_blocks = [self.num_blocks] @@ -635,6 +639,12 @@ class TestNixlHandshake: if pcp_rank > 0: assert connector.get_finished(set()) == (set(), set()) + worker.get_transfer_results = MagicMock( + return_value=KVConnectorTransferResults(finished_sending={"sent"}) + ) + results = connector.get_transfer_results(set()) + assert results.finished_sending == ({"sent"} if pcp_rank == 0 else set()) + @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, @@ -1195,6 +1205,7 @@ class TestNixlHandshake: ssm_sizes=(0, 0), attn_backend_name=worker.backend_name, physical_blocks_per_logical_kv_block=1, + region_num_blocks=None, ) assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [ @@ -1435,9 +1446,9 @@ def test_kv_connector_stats(default_vllm_config, dist_init): metadata = NixlConnectorMetadata() metadata.add_new_req_to_recv( request_id=request_id, - local_block_ids=([1, 2, 3],), + local_block_ids=([0],), kv_transfer_params={ - "remote_block_ids": ([4, 5, 6],), + "remote_block_ids": ([0],), "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID, "remote_request_id": f"prefill-{request_id}", "remote_host": "localhost", @@ -1482,6 +1493,10 @@ def test_kv_connector_stats(default_vllm_config, dist_init): assert stats_after_reset is None +@patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, +) def test_reqs_to_send_deadline_rebased_to_worker_clock(default_vllm_config, dist_init): """reqs_to_send deadlines are stamped with the scheduler process's perf_counter, whose epoch differs across processes and (by boot-time @@ -1906,9 +1921,10 @@ def test_mixed_memory_local_descriptors_split_by_memory_type(): assert handle == 22 assert worker._dram_src_handles_by_block_size[worker.block_size] == 11 assert [block[2] for block in blocks] == [0, 0, 3, 3] - assert [ + memory_types = [ call.args[1] for call in worker.nixl_wrapper.get_xfer_descs.call_args_list - ] == ["DRAM", "VRAM"] + ] + assert memory_types == ["DRAM", "VRAM"] def test_mixed_memory_read_notifies_after_both_transfers_finish(): @@ -1919,6 +1935,7 @@ def test_mixed_memory_read_notifies_after_both_transfers_finish(): worker._recving_metadata = {"request": MagicMock()} worker._recving_transfers = defaultdict(list) worker._pending_recv_notifs = {} + worker._failed_inflight_recvs = set() worker.xfer_stats = MagicMock() worker.nixl_wrapper = MagicMock() worker.nixl_wrapper.make_prepped_xfer.side_effect = [101, 102] @@ -1959,6 +1976,7 @@ def test_mixed_memory_read_failure_does_not_notify_producer(): worker._recving_metadata = {"request": MagicMock()} worker._recving_transfers = {"request": [101, 102]} worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]} + worker._failed_inflight_recvs = set() worker._is_hma_required = True worker._failed_recv_reqs = MagicMock() worker._log_failure = MagicMock() # type: ignore[method-assign] @@ -2871,9 +2889,11 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init): "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FailingNixlWrapper, ) -def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init): - """Test that transfer setup failures mark blocks invalid - and return via get_finished.""" +@pytest.mark.parametrize("is_hma", [False, True]) +def test_transfer_setup_failure_returns_finished( + default_vllm_config, dist_init, is_hma +): + """Setup failures report the request; only non-HMA reports block IDs.""" vllm_config = create_vllm_config() connector = NixlConnector( @@ -2883,6 +2903,7 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init) vllm_config, connector.engine_id, hand_shake_latency=0 ) connector.connector_worker.nixl_wrapper.fail_transfer_setup = True + connector.connector_worker._is_hma_required = is_hma request_id = "test_transfer_fail" metadata = NixlConnectorMetadata() @@ -2912,13 +2933,106 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init) time.sleep(0.1) connector.start_load_kv(dummy_ctx) - # check that blocks were marked invalid + results = connector.get_transfer_results(finished_req_ids=set()) + assert request_id in results.finished_recving + assert results.failed_recving == {request_id} invalid_blocks = connector.get_block_ids_with_load_errors() - assert invalid_blocks == {7, 8, 9} + assert invalid_blocks == (set() if is_hma else {7, 8, 9}) - # ensure request appears in get_finished + +class _ScriptedXferWrapper(FakeNixlWrapper): + """Scripts per-handle xfer states; forbids releasing in-flight handles.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # handle -> successive states; the last one repeats. + self.xfer_states: dict[int, list[str]] = {} + self.released: list[int] = [] + + def check_xfer_state(self, handle: int) -> str: + states = self.xfer_states[handle] + return states.pop(0) if len(states) > 1 else states[0] + + def release_xfer_handle(self, handle: int) -> None: + # A posted-but-unfinished transfer cannot be aborted: releasing it + # leaves the RDMA READ armed. Production code must never do this. + assert self.xfer_states.get(handle, ["DONE"])[0] != "PROC", ( + f"released in-flight handle {handle}" + ) + self.released.append(handle) + + +def _make_split_read_connector(vllm_config, request_id, states): + """Seed a request with scripted in-flight xfer handles + recv metadata.""" + connector = NixlConnector( + vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + ) + connector.connector_worker = FakeNixlConnectorWorker( + vllm_config, connector.engine_id, hand_shake_latency=0 + ) + worker = connector.connector_worker + wrapper = _ScriptedXferWrapper("agent") + worker.nixl_wrapper = wrapper + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id=request_id, + local_block_ids=([7, 8, 9],), + kv_transfer_params={ + "remote_block_ids": ([10, 11, 12],), + "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID, + "remote_request_id": f"prefill-{request_id}", + "remote_host": "localhost", + "remote_port": 1234, + "remote_tp_size": 1, + }, + ) + worker._recving_metadata[request_id] = metadata.reqs_to_recv[request_id] + wrapper.xfer_states = dict(states) + worker._recving_transfers[request_id] = list(states) + return connector, worker, wrapper + + +@patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, +) +def test_split_read_failure_defers_report_until_last_handle( + default_vllm_config, dist_init +): + """One half of a split (mixed DRAM/VRAM) read failing must not report the + request — nor invalidate its blocks — while the sibling xfer is still in + flight: a posted READ cannot be aborted and would DMA into blocks the + scheduler could free and reuse. The report happens exactly once, when the + last handle is terminal.""" + request_id = "split_read_partial_failure" + err_handle, live_handle = 11, 22 + connector, worker, wrapper = _make_split_read_connector( + create_vllm_config(), + request_id, + {err_handle: ["ERR"], live_handle: ["PROC", "PROC", "DONE"]}, + ) + + # Poll 1: one half fails; the sibling is in flight -> nothing reported. _, done_recving = connector.get_finished(finished_req_ids=set()) - assert request_id in done_recving + assert done_recving == set() + assert connector.get_block_ids_with_load_errors() == set() + assert request_id in worker._recving_metadata + assert wrapper.released == [err_handle] + + # Poll 2: sibling still in flight. + _, done_recving = connector.get_finished(finished_req_ids=set()) + assert done_recving == set() + + # Poll 3: sibling terminal -> reported exactly once, blocks invalidated. + _, done_recving = connector.get_finished(finished_req_ids=set()) + assert done_recving == {request_id} + assert connector.get_block_ids_with_load_errors() == {7, 8, 9} + assert request_id not in worker._recving_metadata + assert wrapper.released == [err_handle, live_handle] + + # Poll 4: nothing left; no double report. + _, done_recving = connector.get_finished(finished_req_ids=set()) + assert done_recving == set() @patch( @@ -2938,7 +3052,7 @@ def test_failed_request_skips_kv_postprocessing( default_vllm_config, dist_init, failure_mode ): """Test that failed requests skip KV sync and post-processing in - get_finished(). + get_transfer_results(). This is the core safety behavior: when a KV transfer fails at any stage, the request must still appear in done_recving (so the scheduler can apply @@ -3023,11 +3137,10 @@ def test_failed_request_skips_kv_postprocessing( patch.object(worker, "sync_recved_kv_to_device") as mock_sync, patch.object(worker, "post_process_device_kv_on_receive") as mock_postprocess, ): - _, done_recving = connector.get_finished(finished_req_ids=set()) + results = connector.get_transfer_results(finished_req_ids=set()) - # The failed request must appear in done_recving so the scheduler - # can handle it (e.g., trigger recompute via kv_load_failure_policy). - assert request_id in done_recving + assert request_id in results.finished_recving + assert results.failed_recving == {request_id} # Critical: KV sync and post-processing must NOT have been called # since no valid KV data was received for the failed request. @@ -3049,19 +3162,7 @@ def test_failed_request_skips_kv_postprocessing( def test_handles_failing_in_separate_polls_do_not_kill_the_engine( default_vllm_config, dist_init ): - """A request whose handles fail in different polls must not crash the engine. - - One transfer handle is created per remote rank, and when a peer goes away - they do not all fail in the same poll: one errors while another is still - PROC. The first failure reports the request via _failed_recv_reqs and - get_finished() pops its metadata; when the remaining handles fail in a - later poll, the request must be cleaned up without being reported again. - - Reporting twice kills the EngineCore: the scheduler's assert in - _update_from_kv_xfer_finished only expects a finished recv for a request - still waiting for KVs, having moved this one out of - WAITING_FOR_REMOTE_KVS to recompute locally after the first report. - """ + """Failures across polls report once, after the last writer is terminal.""" vllm_config = create_vllm_config() connector = NixlConnector( vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) @@ -3107,20 +3208,21 @@ def test_handles_failing_in_separate_polls_do_not_kill_the_engine( ): _, done_recving = connector.get_finished(finished_req_ids=set()) - assert request_id in done_recving - assert request_id not in worker._recving_metadata + assert request_id not in done_recving + assert request_id in worker._recving_metadata assert worker._recving_transfers[request_id] == [second] + assert connector.get_block_ids_with_load_errors() == set() - # Poll 2: the second handle fails too. The request was already reported, - # so it must not be reported a second time: the scheduler has since moved - # it out of WAITING_FOR_REMOTE_KVS to recompute locally, and asserts on a - # finished recv for a request in that state. Its remaining handles are - # still released and the transfer entry removed. + # Only the last terminal handle permits the scheduler to reuse the blocks. with patch.object(worker.nixl_wrapper, "check_xfer_state", return_value="ERR"): _, done_recving = connector.get_finished(finished_req_ids=set()) - assert request_id not in done_recving + assert done_recving == {request_id} assert request_id not in worker._recving_transfers + assert request_id not in worker._recving_metadata + assert connector.get_block_ids_with_load_errors() == {1, 2, 3} + _, done_recving = connector.get_finished(finished_req_ids=set()) + assert done_recving == set() def _set_test_speculative_config( diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 08506d82f2..259e9536c1 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -5,16 +5,31 @@ import gc from unittest.mock import MagicMock, patch +import msgspec import pytest import torch from tests.v1.attention.utils import MockMambaBuilder from vllm import LLM, SamplingParams -from vllm.config import KVTransferConfig +from vllm.config import KVTransferConfig, set_current_vllm_config +from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlAgentMetadata, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, +) from vllm.v1.core.single_type_kv_cache_manager import ( FullAttentionManager, SlidingWindowManager, ) +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + MLAAttentionSpec, +) from .utils import ( create_request, @@ -272,26 +287,24 @@ def test_logical_to_kernel_block_ids_with_hma(): @pytest.mark.cpu_test @pytest.mark.parametrize( - "is_rocm,has_mamba,use_host_buffer,done_recving,failed_recving,expected_syncs", + "is_rocm,has_mamba,use_host_buffer,done_recving,expected_syncs", [ - (True, True, False, {"req"}, set(), 1), - (False, True, False, {"req"}, set(), 0), - (True, False, False, {"req"}, set(), 0), - (True, True, True, {"req"}, set(), 0), - (True, True, False, set(), set(), 0), - (True, True, False, {"req"}, {"req"}, 0), + (True, True, False, {"req"}, 1), + (False, False, False, {"req"}, 0), + (False, True, False, {"req"}, 0), + (True, True, True, {"req"}, 0), + (True, True, False, set(), 0), ], ) -def test_sync_device_after_mamba_recv_gates( +def test_sync_device_after_direct_recv_gates( monkeypatch, is_rocm, has_mamba, use_host_buffer, done_recving, - failed_recving, expected_syncs, ): - """Only direct-GPU Mamba receives on ROCm need a device fence.""" + """Only backends requiring visibility fences synchronize direct receives.""" from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( NixlConnectorWorker, @@ -309,7 +322,7 @@ def test_sync_device_after_mamba_recv_gates( lambda: sync_calls.append(True), ) - worker._sync_device_after_mamba_recv(done_recving, failed_recving) + worker._sync_device_after_direct_recv(done_recving) assert len(sync_calls) == expected_syncs @@ -399,6 +412,8 @@ def test_read_blocks_for_req_expands_remote_ids( ) remote_engine_id = "remote-engine" + worker.region_group_ids = list(range(len(resolved_types))) + worker.dst_region_group_ids = {remote_engine_id: list(worker.region_group_ids)} worker.transfer_topo = MagicMock() # tp_ratio not exercised (all_source_ranks is empty so no reads run), @@ -888,6 +903,45 @@ def _make_mock_worker_for_desc_ids( return worker +@pytest.mark.cpu_test +def test_shared_nixl_region_broadcasts_every_group_block_table(): + """A shared region must receive every group's block table exactly once.""" + worker = _make_mock_worker_for_desc_ids( + num_regions=3, + has_mamba=False, + group_spec_types=(FullAttentionSpec,) * 3, + block_len_per_layer=[100] * 3, + ) + worker.region_group_ids = [0, bw._SHARED_REGION_GROUP_ID, 2] + worker._uses_region_group_mapping = True + + result = worker._compute_desc_ids( + block_ids=([1], [2], [3]), + dst_num_blocks=10, + block_size_ratio=None, + physical_blocks_per_logical=1, + region_num_blocks=[10, 10, 10], + ) + + assert result.tolist() == [1, 11, 12, 13, 23] + + +@pytest.mark.cpu_test +def test_nixl_region_sort_preserves_pipeline_layer_order(): + """Numeric layer order prevents PP peers from pairing different regions.""" + names = [ + "model.layers.10.self_attn", + "model.layers.2.self_attn", + "model.layers.1.self_attn", + ] + + assert sorted(names, key=bw._region_sort_key) == [ + "model.layers.1.self_attn", + "model.layers.2.self_attn", + "model.layers.10.self_attn", + ] + + @pytest.mark.cpu_test def test_get_block_descs_ids_hybrid_ssm(): """Test _compute_desc_ids uses per-group strides for hybrid @@ -1097,10 +1151,7 @@ def test_post_process_zeroes_untransferred_tail(): worker.device_kv_caches = {"attn.0": attn_cache, "mamba.0": mamba_cache} fa_group = MagicMock(layer_names=["attn.0"]) ssm_group = MagicMock(layer_names=["mamba.0"]) - worker.kv_cache_config = MagicMock( - kv_cache_groups=[fa_group, ssm_group], - transfer_groups=[fa_group, ssm_group], - ) + worker.kv_cache_config = MagicMock(transfer_groups=[fa_group, ssm_group]) # The cached property filters mamba layers out of the permuted caches. attn_caches = NixlConnectorWorker._attention_kv_caches.func(worker) assert len(attn_caches) == 1 and attn_caches[0] is attn_cache @@ -1188,6 +1239,11 @@ def _make_fake_kv_cache_manager(): manager = object.__new__(KVCacheManager) manager.coordinator = MagicMock() + manager.kv_cache_config = MagicMock() + manager.kv_cache_config.kv_cache_groups = ( + MagicMock(host_resident=False), + MagicMock(host_resident=False), + ) manager.coordinator.single_type_managers = ( _FakeSingleTypeManager(True, 16, [10, 11, 12, 13, 14, 15]), # attention _FakeSingleTypeManager(False, 16, [20, 21, 22, 23, 24, 25]), # mamba @@ -1723,16 +1779,6 @@ def test_exchange_clipped_blocks_ssm_positional_states(): assert clipped == ([1, 2, 3], [0, 5, 6, 7]) -@pytest.mark.cpu_test -def test_exchange_clipped_blocks_excludes_nontransfer_groups(): - """Scheduler block tables must match the worker's registered regions.""" - sched = make_nixl_scheduler() - sched.kv_cache_config = make_kv_cache_config(block_size=16, swa_enabled=True) - sched.kv_cache_config.kv_cache_groups[1].enable_kv_transfer = False - - assert sched.get_exchange_clipped_blocks(([1, 2], [9])) == ([1, 2],) - - # ── Hybrid MLA+SSM (KimiLinear-shaped KDA+MLA) tests ───────────────────── @@ -1786,6 +1832,95 @@ def _make_hybrid_mla_kv_cache_config(num_blocks: int = 4): ) +@pytest.mark.cpu_test +@pytest.mark.parametrize("kernel_block_size", [16, 8]) +def test_nixl_keeps_device_block_count_with_hisparse_host_pool(kernel_block_size): + host_num_blocks = 4 + gpu_num_blocks = 9 + spec = MLAAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=8, + dtype=torch.float16, + ) + kv_cache_config = KVCacheConfig( + num_blocks=gpu_num_blocks, + hisparse_host_num_blocks=host_num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=host_num_blocks * spec.page_size_bytes, + layers=["mla.host"], + layer_stride=host_num_blocks * spec.page_size_bytes, + block_stride=spec.page_size_bytes, + host_resident=True, + ), + KVCacheTensor( + size=gpu_num_blocks * spec.page_size_bytes, + layers=["mla.device"], + layer_stride=gpu_num_blocks * spec.page_size_bytes, + block_stride=spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["mla.host"], + spec, + host_resident=True, + ), + KVCacheGroupSpec(["mla.device"], spec), + ], + ) + vllm_config = create_vllm_config(block_size=16) + vllm_config.kv_transfer_config.kv_buffer_device = "cuda" + fake_backend = MagicMock() + fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size] + fake_backend.get_name.return_value = "FLASHMLA" + fake_backend.full_cls_name.return_value = "fake.FLASHMLA" + fake_backend.get_kv_cache_shape.return_value = (1, kernel_block_size, 1, 1) + fake_platform = MagicMock() + fake_platform.device_type = "cuda" + fake_platform.get_nixl_memory_type.return_value = "VRAM" + physical_per_logical = spec.block_size // kernel_block_size + physical_page_bytes = spec.page_size_bytes // physical_per_logical + host_cache = torch.zeros( + host_num_blocks * physical_per_logical, physical_page_bytes, dtype=torch.uint8 + ) + device_cache = torch.zeros( + gpu_num_blocks * physical_per_logical, physical_page_bytes, dtype=torch.uint8 + ) + + with ( + patch.object(bw, "NixlWrapper"), + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0), + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1), + patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]), + patch.object(bw, "current_platform", fake_platform), + set_current_vllm_config(vllm_config), + ): + worker = NixlConnectorWorker(vllm_config, "test-engine", kv_cache_config) + worker.use_mla = True + worker.nixl_wrapper.get_agent_metadata.return_value = b"metadata" + worker.register_kv_caches({"mla.host": host_cache, "mla.device": device_cache}) + + assert worker.num_blocks == gpu_num_blocks * physical_per_logical + assert dict(zip(worker.region_group_ids, worker.region_num_blocks)) == { + 0: host_num_blocks * physical_per_logical, + 1: gpu_num_blocks * physical_per_logical, + } + caches = {"mla.host": host_cache, "mla.device": device_cache} + expected_addrs = [ + caches[name].data_ptr() + block * stride + for name, stride, count in zip( + worker.region_names, + worker.block_stride_per_layer, + worker.region_num_blocks, + strict=True, + ) + for block in range(count) + ] + assert worker.src_blocks_data[:, 0].tolist() == expected_addrs + + @pytest.mark.cpu_test def test_register_kv_caches_hybrid_mla_dual_purpose_regions(): """Hybrid MLA+KDA registration: HMA tensors shared by both layer types @@ -1803,6 +1938,9 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions(): kv_cache_config = _make_hybrid_mla_kv_cache_config() unified_page = kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes vllm_config = create_vllm_config(block_size=12) + # The engine rewrites this global value to the smallest scheduler-group + # block size. NIXL must retain the transferable groups' 12-token geometry. + vllm_config.cache_config.block_size = 4 # kv_buffer_device defaults to the *real* platform's device type, which on # a CPU-only test host would make this a host-buffer worker: host xfer # buffers are per-layer, so the HMA shared tensors would not be @@ -1833,7 +1971,13 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions(): worker.use_mla = True # opt-125m test config is not MLA; force the flag worker.nixl_wrapper.get_agent_metadata.return_value = b"fake-agent-metadata" - tensors = [torch.zeros(4 * unified_page, dtype=torch.uint8) for _ in range(2)] + kernel_page = unified_page // 3 + block_stride = 2 * kernel_page + backing = torch.zeros(12, block_stride, dtype=torch.uint8) + tensors = [ + backing[:, :kernel_page], + backing[:, kernel_page:], + ] # KDA layer first per tensor: exercises the dual-purpose flag merge. worker.register_kv_caches( { @@ -1849,6 +1993,8 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions(): # 12-token logical blocks over the 4-token MLA kernel block. assert worker._physical_blocks_per_logical_kv_block == 3 assert worker.block_size == 4 and worker.num_blocks == 12 + assert worker.region_group_ids == [bw._SHARED_REGION_GROUP_ID] * 2 + assert worker.block_stride_per_layer == [block_stride, block_stride] # Both shared tensors are dual-purpose: their FA view is MLA even though # a KDA layer registered them first. assert worker._region_is_mla == [True, True] @@ -1861,8 +2007,20 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions(): # Mamba descs: 2 regions x (3 conv sub-projections + 1 ssm) x 4 blocks. assert worker.src_blocks_data.shape == (24 + 32, 3) fa_descs = worker.src_blocks_data[:24] - assert fa_descs[1][0] - fa_descs[0][0] == unified_page // 3 + expected_addrs = [ + tensor.data_ptr() + block * block_stride + for tensor in tensors + for block in range(12) + ] + assert fa_descs[:, 0].tolist() == expected_addrs assert all(size == unified_page // 3 for size in fa_descs[:, 1]) + worker.nixl_wrapper.register_memory.assert_called_once() + metadata = msgspec.msgpack.decode( + worker.xfer_handshake_metadata.agent_metadata_bytes, + type=NixlAgentMetadata, + ) + assert metadata.block_strides == [block_stride, block_stride] + assert metadata.region_num_blocks == [12, 12] @pytest.mark.cpu_test @@ -1909,6 +2067,8 @@ def test_push_write_hybrid_mla_replicates_attention(): worker.dst_xfer_side_handles = {engine_id: {0: 100, 1: 101}} worker.src_xfer_handles_by_tp_ratio = {(-2, 4): [200, 201]} worker.src_xfer_handles_by_block_size = {4: 300} + worker.region_group_ids = [0, 1] + worker.dst_region_group_ids = {engine_id: [0, 1]} worker._sending_transfers = defaultdict(list) worker._sending_transfers_lock = threading.Lock() worker.kv_cache_config = _make_hybrid_mla_kv_cache_config() diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index 32908f78ad..607e4822b9 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -186,11 +186,16 @@ def test_overlaid_transfer_groups_share_region_geometry(): worker.host_buffer_kv_cache_layout = "NHD" worker._physical_blocks_per_logical_kv_block = 1 worker._logical_num_blocks = num_blocks + worker.region_mem_types = [] worker.region_group_ids = [] worker.region_mem_types = [] worker._mixed_mem_types = False worker.region_names = [] worker.region_num_blocks = [] + worker._mixed_mem_types = False + worker._desc_is_dram_by_block_size = {} + worker._desc_pos_by_block_size = {} + worker._dram_src_handles_by_block_size = {} worker._region_is_mla = [] worker.block_len_per_layer = [] worker.block_stride_per_layer = [] @@ -304,8 +309,6 @@ def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blo # would not be deduplicated. Pin it to the faked device type. vllm_config.kv_transfer_config.kv_buffer_device = "cuda" - from unittest.mock import MagicMock - fake_backend = MagicMock() fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size] fake_backend.get_name.return_value = "FLASHMLA" @@ -552,6 +555,16 @@ def _make_remote_meta( ) +def _register_remote_agents(worker, metadata, tp_size): + """Mirror the async handshake callback that publishes prepared agents.""" + worker._remote_agents[metadata.engine_id] = { + (0, rank): worker.add_remote_agent( + metadata, remote_tp_rank=rank, remote_tp_size=tp_size + ) + for rank in range(tp_size) + } + + def _owned_byte_ranges(worker, group_logical_ids): """Byte ranges owned by a request: for each HMA region tensor, every logical block id of every group maps to one unified page.""" @@ -609,8 +622,7 @@ def test_hetero_ppl_multi_read_writes_stay_within_request_blocks(): remote_num_logical=12, remote_ssm_sizes=(24, 32), ) - for rank in (0, 1): - worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=2) + _register_remote_agents(worker, meta_r, 2) # Request B: 17 matched tokens. Local: 2 logical blocks (24 tok # capacity); remote: 16 prefilled tokens -> 2 remote logical blocks. @@ -626,7 +638,7 @@ def test_hetero_ppl_multi_read_writes_stay_within_request_blocks(): "remote_block_ids": remote_ids, "remote_engine_id": "remote-engine", "remote_request_id": "prefill-req-b", - "remote_host": "localhost", + "remote_host": "remote-host", "remote_port": 1234, "tp_size": 2, }, @@ -709,8 +721,7 @@ def _run_hetero_case( remote_num_logical=max(2 * n_remote + 4, 8), remote_ssm_sizes=(48 // tp_size, 64 // tp_size), ) - for rank in range(tp_size): - worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=tp_size) + _register_remote_agents(worker, meta_r, tp_size) # Sparse ids so neighbors exist between the request's blocks. local_attn = [2 * i + 1 for i in range(n_local)] @@ -726,7 +737,7 @@ def _run_hetero_case( "remote_block_ids": remote_ids, "remote_engine_id": "remote-engine", "remote_request_id": "prefill-req-b", - "remote_host": "localhost", + "remote_host": "remote-host", "remote_port": 1234, "tp_size": tp_size, }, diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 5b7760e7e3..ce233bbb60 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -15,7 +15,7 @@ requiring a real NIXL agent or network: * The worker matches D registrations against P finished blocks (both scenario directions) and forwards non-PUSH_REG NIXL notifs to the main thread's ``_get_new_notifs``. -* ``get_finished`` enqueues evictions for the writer. +* ``get_transfer_results`` enqueues evictions for the writer. """ from __future__ import annotations @@ -32,6 +32,9 @@ from unittest.mock import MagicMock, patch import msgspec import pytest +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorTransferResults, +) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( PUSH_REG_NOTIF_PREFIX, NixlAgentMetadata, @@ -347,6 +350,8 @@ class _StubWriterWorker(NixlPushConnectorWorker): # Base worker fields touched by start_load_kv / _get_new_notifs. w._recving_metadata = {} + w._pending_recv_notifs = {} + w._failed_inflight_recvs = set() w._recving_transfers = defaultdict(list) w._reqs_to_process = set() w._reqs_to_send = {} @@ -742,21 +747,19 @@ class TestPushWriterNotifs: assert notified == set() assert request_id in w._recving_transfers - def test_get_finished_evicts_completed_state(self): - """``get_finished`` should enqueue evictions and wake the writer.""" + def test_get_transfer_results_evicts_completed_state(self): + """Transfer completion should enqueue evictions and wake the writer.""" w = _StubWriterWorker.fresh() - # Stub the base ``get_finished`` to return one done_sending entry. - # Patch via the MRO's parent class. with patch.object( NixlPushConnectorWorker.__mro__[1], - "get_finished", - return_value=({"req-done"}, set()), + "get_transfer_results", + return_value=KVConnectorTransferResults(finished_sending={"req-done"}), ): - done_sending, done_recving = w.get_finished() + results = w.get_transfer_results() - assert "req-done" in done_sending - assert done_recving == set() + assert results.finished_sending == {"req-done"} + assert results.finished_recving == set() # Eviction enqueued for the writer. evicted = [] while True: @@ -928,19 +931,21 @@ class TestPushWriterNegative: assert len(w._pending_d_registrations) == 1 assert w.start_push_calls == [] - def test_get_finished_enqueues_eviction_for_each_done_request(self): - """``get_finished`` must enqueue an eviction for every request + def test_get_transfer_results_enqueues_each_completed_request(self): + """``get_transfer_results`` must enqueue every completed request in ``done_sending`` so the writer can drop stale matching state. Unlike the happy-path test, this verifies the *cardinality*: N completed requests -> N evictions, in order.""" w = _StubWriterWorker.fresh() with patch.object( NixlPushConnectorWorker.__mro__[1], - "get_finished", - return_value=({"req-1", "req-2", "req-3"}, set()), + "get_transfer_results", + return_value=KVConnectorTransferResults( + finished_sending={"req-1", "req-2", "req-3"} + ), ): - done_sending, _ = w.get_finished() - assert done_sending == {"req-1", "req-2", "req-3"} + results = w.get_transfer_results() + assert results.finished_sending == {"req-1", "req-2", "req-3"} evicted: list[str] = [] while True: @@ -950,23 +955,47 @@ class TestPushWriterNegative: break assert sorted(evicted) == ["req-1", "req-2", "req-3"] - def test_get_finished_with_no_completions_does_not_enqueue_eviction(self): + def test_empty_transfer_results_do_not_enqueue_eviction(self): """If there's nothing newly done, no eviction should be enqueued. The wake event IS still set because ``get_finished`` always wakes the writer to drain notifs.""" w = _StubWriterWorker.fresh() with patch.object( NixlPushConnectorWorker.__mro__[1], - "get_finished", - return_value=(set(), set()), + "get_transfer_results", + return_value=KVConnectorTransferResults(), ): - done_sending, done_recving = w.get_finished() - assert done_sending == set() - assert done_recving == set() + results = w.get_transfer_results() + assert results.finished_sending == set() + assert results.finished_recving == set() assert w._evict_finished_inbox.qsize() == 0 # Wake set so the writer drains NIXL notifs even when idle. assert w._push_writer_wake.is_set() + def test_failed_push_transfer_finishes_sending(self): + w = _StubWriterWorker.fresh() + w.nixl_wrapper = MagicMock() + w.nixl_wrapper.check_xfer_state.return_value = "ERR" + w.xfer_stats = MagicMock() + w._log_failure = MagicMock() + w._sending_transfers["req-failed"] = [17] + w._reqs_to_send["req-failed"] = time.perf_counter() + 30 + w._reqs_to_process.add("req-failed") + + with patch.object( + NixlPushConnectorWorker.__mro__[1], + "get_transfer_results", + return_value=KVConnectorTransferResults(), + ): + results = w.get_transfer_results() + + assert results.finished_sending == {"req-failed"} + assert results.finished_recving == set() + assert "req-failed" not in w._sending_transfers + assert "req-failed" not in w._reqs_to_send + assert "req-failed" not in w._reqs_to_process + w.nixl_wrapper.release_xfer_handle.assert_called_once_with(17) + def test_get_new_notifs_unknown_request_is_logged_and_skipped(self, caplog): """A completion notif for a request the worker doesn't know about should be logged but not crash.""" @@ -1160,6 +1189,7 @@ class TestPushWriterMlaReplication: ) } w._logical_to_kernel_block_ids = lambda block_ids, ratio: block_ids + w.dst_region_group_ids[engine_id] = [0] w.dst_xfer_side_handles = {engine_id: {r: 1000 + r for r in d_ranks}} w.src_xfer_handles_by_block_size = {16: 2000} w._remote_agents = {engine_id: {(0, r): f"agent-{r}" for r in d_ranks}} @@ -1240,6 +1270,8 @@ class TestPushWriterMlaReplication: ) } w._logical_to_kernel_block_ids = lambda block_ids, ratio: block_ids + w.region_group_ids = [0, 1, 2, 3] + w.dst_region_group_ids[engine_id] = [0, 1, 2, 3] w.dst_xfer_side_handles = {engine_id: {0: 1000, 1: 1001}} w.src_xfer_handles_by_tp_ratio = {(-2, 16): [2000, 2001]} diff --git a/tests/v1/kv_connector/unit/test_output_aggregator.py b/tests/v1/kv_connector/unit/test_output_aggregator.py index b083ccef98..e757a65842 100644 --- a/tests/v1/kv_connector/unit/test_output_aggregator.py +++ b/tests/v1/kv_connector/unit/test_output_aggregator.py @@ -15,12 +15,14 @@ class DummyModelRunnerOutput(ModelRunnerOutput): finished_sending: set[str] | None = None, finished_recving: set[str] | None = None, invalid_block_ids: set[int] | None = None, + failed_recving: set[str] | None = None, expected_finished_count: int = 0, ): self.kv_connector_output = KVConnectorOutput( finished_sending=finished_sending, finished_recving=finished_recving, invalid_block_ids=invalid_block_ids or set(), + failed_recving=failed_recving or set(), expected_finished_count=expected_finished_count, ) @@ -73,7 +75,9 @@ def test_aggregate_workers_output(): output1 = DummyModelRunnerOutput(invalid_block_ids={3, 4}) output2 = DummyModelRunnerOutput( - finished_recving={"req2"}, invalid_block_ids={4, 5} + finished_recving={"req2"}, + invalid_block_ids={4, 5}, + failed_recving={"req3"}, ) aggregated = aggregator.aggregate([output1, output2]) @@ -83,6 +87,12 @@ def test_aggregate_workers_output(): assert aggregated.finished_sending is None assert aggregated.finished_recving == {"req2"} assert aggregated.invalid_block_ids == {3, 4, 5} + assert not aggregated.failed_recving + + output1 = DummyModelRunnerOutput(finished_recving={"req3"}) + output2 = DummyModelRunnerOutput(finished_recving={"req3"}) + aggregated = aggregator.aggregate([output1, output2]) + assert aggregated.kv_connector_output.failed_recving == {"req3"} def test_aggregate_workers_output_with_expected_finished_count(): diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 80cab94b34..7b70065d7f 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -539,6 +539,7 @@ def make_nixl_scheduler( sched._heartbeat_interval = kv_lease_duration // 6 # Fields touched by build_connector_meta / request_finished: sched._reqs_need_recv = {} + sched._hisparse_host_blocks_to_recv = {} sched._reqs_need_send = {} sched._reqs_in_batch = set() sched._reqs_not_processed = set() diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index 251cef3251..16539677c4 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -62,7 +62,7 @@ def _make_offloading_config( normalized_extra_config["cpu_bytes_to_use"] = cpu_bytes_to_use if groups is None: - groups = (OffloadingGroupConfig(16, ("layer",)),) + groups = (OffloadingGroupConfig(16, ("layer",), 0),) return OffloadingConfig( groups=groups, @@ -493,8 +493,8 @@ def test_offloading_spec_has_replicated_layout_default(): def test_offloading_spec_uses_normalized_chunk_geometry(): groups = ( - OffloadingGroupConfig(12, ("full_layer",)), - OffloadingGroupConfig(16, ("mla_layer",)), + OffloadingGroupConfig(12, ("full_layer",), 0), + OffloadingGroupConfig(16, ("mla_layer",), 1), ) spec = _create_spec( groups=groups, diff --git a/tests/v1/kv_offload/test_file_mapper.py b/tests/v1/kv_offload/test_file_mapper.py index 26da4079b4..8cd1ae9b10 100644 --- a/tests/v1/kv_offload/test_file_mapper.py +++ b/tests/v1/kv_offload/test_file_mapper.py @@ -26,8 +26,11 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper: OffloadingGroupConfig( tokens_per_block=tokens_per_block, layer_names=(layer_name,), + group_id=group_id, + ) + for group_id, (tokens_per_block, layer_name) in enumerate( + kwargs.get("groups", ()) ) - for tokens_per_block, layer_name in kwargs.get("groups", ()) ), worker_kv_bytes_per_block=0, enable_kv_cache_events=False, diff --git a/tests/v1/simple_kv_offload/test_worker.py b/tests/v1/simple_kv_offload/test_worker.py index cbc7f20613..a62dacc310 100644 --- a/tests/v1/simple_kv_offload/test_worker.py +++ b/tests/v1/simple_kv_offload/test_worker.py @@ -400,6 +400,8 @@ def test_register_mixed_page_sizes_in_one_cache_group(monkeypatch): vllm_config.cache_config = CacheConfig() vllm_config.cache_config.kv_cache_layout = layout.name vllm_config.cache_config.num_gpu_blocks_override = None + vllm_config.attention_config.hisparse_config = None + vllm_config.kv_transfer_config = None pages = [spec.page_size_bytes for spec in specs.values()] num_blocks = 4 @@ -458,6 +460,8 @@ def test_register_mixed_page_sizes_odd_block_counts(monkeypatch, rank_blocks): vllm_config.cache_config = CacheConfig() vllm_config.cache_config.kv_cache_layout = layout.name vllm_config.cache_config.num_gpu_blocks_override = None + vllm_config.attention_config.hisparse_config = None + vllm_config.kv_transfer_config = None pages = [spec.page_size_bytes for spec in specs.values()] kv_cache_config = get_kv_cache_config_from_groups( @@ -507,6 +511,8 @@ def test_mixed_page_byte_placement_is_dcp_invariant(): vllm_config.cache_config.kv_cache_layout = KVCacheLayout.LBNHC.name vllm_config.cache_config.num_gpu_blocks_override = None vllm_config.parallel_config.decode_context_parallel_size = dcp_size + vllm_config.attention_config.hisparse_config = None + vllm_config.kv_transfer_config = None page_bytes = sum(spec.page_size_bytes for spec in specs.values()) config = get_kv_cache_config_from_groups(vllm_config, [group], page_bytes * 4) diff --git a/tests/v1/test_kv_cache_spec_registry.py b/tests/v1/test_kv_cache_spec_registry.py index 58c23ce285..6da14b4e42 100644 --- a/tests/v1/test_kv_cache_spec_registry.py +++ b/tests/v1/test_kv_cache_spec_registry.py @@ -16,6 +16,7 @@ from vllm.v1.core.single_type_kv_cache_manager import ( ChunkedLocalAttentionManager, CrossAttentionManager, FullAttentionManager, + HiSparseSourceManager, MambaManager, SingleTypeKVCacheManager, SinkFullAttentionManager, @@ -27,6 +28,7 @@ from vllm.v1.kv_cache_interface import ( CrossAttentionSpec, FullAttentionSpec, HiddenStateCacheSpec, + KVCacheGroupRole, KVCacheSpec, KVCacheSpecKind, MambaSpec, @@ -176,6 +178,19 @@ class TestKVCacheSpecRegistry: is spec_uniform_base_map[spec_cls] ) + @pytest.mark.parametrize("role", list(KVCacheGroupRole)) + def test_mla_manager_selection_by_role(self, role): + """Only the source role overrides ordinary MLA manager selection.""" + expected = ( + HiSparseSourceManager + if role == KVCacheGroupRole.HISPARSE_SOURCE + else FullAttentionManager + ) + assert ( + KVCacheSpecRegistry.get_manager_class(make_spec(MLAAttentionSpec), role) + is expected + ) + @pytest.mark.parametrize("spec_cls", list(spec_manager_map)) def test_custom_spec_register(self, spec_cls): """A decorated custom spec resolves to the declared manager.""" diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 906525494f..dc75ba4ea1 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -13,15 +13,18 @@ import pytest import torch from tests.v1.attention.utils import dense_kv_cache_views -from vllm.v1.attention.backend import AttentionCGSupport +from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy +from vllm.v1.hisparse.binding import allocate_hisparse_kv_caches from vllm.v1.kv_cache_interface import ( FullAttentionSpec, + HiSparseResidentSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheLayout, KVCacheTensor, MLAAttentionSpec, + SparseCacheRole, compute_layout_strides, ) from vllm.v1.worker.gpu import attn_utils @@ -36,6 +39,60 @@ from vllm.v1.worker.utils import ( ) +@pytest.mark.parametrize( + ("enabled", "block_size", "main_sizes", "indexer_sizes", "expected"), + [ + (True, 256, [64], [64], 64), + (True, 64, [32, 64], [16, 32], 32), + (True, 64, [MultipleOf(16)], [32], 32), + (True, 64, [64], [32], None), + (False, 256, [64], [64], 256), + ], +) +def test_get_kv_cache_spec_resolves_hisparse_block_size( + monkeypatch, enabled, block_size, main_sizes, indexer_sizes, expected +): + """Resolve shared MLA geometry before planning; leave other specs alone.""" + specs = { + "main": MLAAttentionSpec( + block_size=block_size, num_kv_heads=1, head_size=576, dtype=torch.bfloat16 + ), + "indexer": MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + cache_role=SparseCacheRole.INDEXER, + ), + "dense": FullAttentionSpec( + block_size=block_size, num_kv_heads=1, head_size=128, dtype=torch.bfloat16 + ), + } + layers = {} + for name, sizes in zip(specs, [main_sizes, indexer_sizes, [block_size]]): + backend = SimpleNamespace( + customize_spec=AttentionBackend.customize_spec, + get_supported_kernel_block_sizes=lambda sizes=sizes: sizes, + ) + layers[name] = SimpleNamespace( + get_kv_cache_spec=lambda _, spec=specs[name]: spec, + get_attn_backend=lambda backend=backend: backend, + ) + monkeypatch.setattr(attn_utils, "get_layers_from_vllm_config", lambda *_: layers) + config = SimpleNamespace( + attention_config=SimpleNamespace(hisparse_config=object() if enabled else None) + ) + if expected is None: + with pytest.raises(ValueError, match="supported by every sparse"): + attn_utils.get_kv_cache_spec(config) + return + + resolved = attn_utils.get_kv_cache_spec(config) + assert resolved["main"].block_size == resolved["indexer"].block_size == expected + assert resolved["dense"] is specs["dense"] + assert all(spec.block_size == block_size for spec in specs.values()) + + class _FakeMetadataBuilder: def __init__(self, support: AttentionCGSupport): self.support = support @@ -390,3 +447,71 @@ def test_copy_kv_cache_blocks_with_virtual_block_splitting( torch.testing.assert_close( cache[dst_start + physical_idx], expected[layer_idx][physical_idx] ) + + +def test_allocate_hisparse_kv_caches_host_pool_and_view_less_specs(): + """Host tensors get their own backing; view-less specs keep the raw one.""" + spec = FullAttentionSpec( + block_size=2, num_kv_heads=1, head_size=4, dtype=torch.float32 + ) + page = spec.page_size_bytes + resident_spec = HiSparseResidentSpec(block_size=2, page_size=page) + device_size = 4 * page + config = KVCacheConfig( + num_blocks=4, + hisparse_host_num_blocks=3, + kv_cache_tensors=[ + KVCacheTensor( + size=3 * page, + layers=["source"], + layer_stride=3 * page, + block_stride=page, + host_resident=True, + ), + KVCacheTensor( + size=device_size, + layers=["indexer"], + layer_stride=device_size, + block_stride=page, + ), + KVCacheTensor( + size=device_size, + layers=["resident"], + layer_stride=device_size, + block_stride=page, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["source"], spec, host_resident=True), + KVCacheGroupSpec(["indexer"], spec), + KVCacheGroupSpec(["resident"], resident_spec), + ], + ) + host_buffers: list[torch.Tensor] = [] + + def host_allocator(size: int) -> torch.Tensor: + host_buffers.append(torch.zeros(size, dtype=torch.int8)) + return host_buffers[-1] + + caches = allocate_hisparse_kv_caches( + config, + torch.device("cpu"), + KVCacheLayout.LBHNC, + [2, 2, 2], + SimpleNamespace(allocate=host_allocator), + ) + assert len(config.kv_cache_tensors) == 3 + + assert [buf.numel() for buf in host_buffers] == [3 * page] + assert caches["source"].shape[0] == 3 + assert ( + caches["source"].untyped_storage().data_ptr() + == host_buffers[0].untyped_storage().data_ptr() + ) + assert caches["indexer"].shape[0] == 4 + backing = caches["resident"] + assert backing.dtype == torch.int8 and backing.numel() >= device_size + assert ( + backing.untyped_storage().data_ptr() + == caches["indexer"].untyped_storage().data_ptr() + ) diff --git a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py index f84c91872f..4cc8f9ffe2 100644 --- a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py +++ b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py @@ -134,6 +134,7 @@ def test_overlaid_zeroer_dedups_segments_with_max_span(): vllm_config.cache_config = CacheConfig() vllm_config.cache_config.num_gpu_blocks_override = None vllm_config.cache_config.kv_cache_layout = "BLHNC" + vllm_config.attention_config.hisparse_config = None config = get_kv_cache_config_from_groups(vllm_config, groups, 8 * 1024 * 1024) views = allocate_kv_cache(config, torch.device("cpu"), KVCacheLayout.BLHNC, None) buf_ptr = views["g1.big"].data_ptr() diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index e9c7a3b290..920c9a25f0 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/tests/v1/worker/test_gpu_autoregressive_speculator.py @@ -351,6 +351,8 @@ def test_multi_step_decode_replays_captured_graph_as_expected( speculator = object.__new__(_TestSpeculator) speculator.num_speculative_steps = 4 speculator.current_draft_step = torch.tensor(0) + speculator.slot_mapping_observer = None + speculator.host_mirror_forward_observer = None speculator.input_buffers = SimpleNamespace( positions=torch.arange(2), query_start_loc=torch.arange(3), diff --git a/tests/v1/worker/test_gpu_kv_connector.py b/tests/v1/worker/test_gpu_kv_connector.py index 7fb40fff37..03c483dce1 100644 --- a/tests/v1/worker/test_gpu_kv_connector.py +++ b/tests/v1/worker/test_gpu_kv_connector.py @@ -5,9 +5,11 @@ from types import SimpleNamespace from unittest.mock import Mock import pytest +import torch import vllm.v1.worker.gpu.kv_connector as kv_connector_module from vllm.config import KVTransferConfig +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorTransferResults from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector @@ -18,9 +20,9 @@ def _make_connector( backend = Mock() backend.handle_preemptions.side_effect = lambda _: events.append("handle") backend.bind_connector_metadata.side_effect = lambda _: events.append("bind") - backend.start_load_kv.side_effect = lambda _: events.append("start") + backend.start_load_kv.side_effect = lambda *_args, **_kwargs: events.append("start") backend.wait_for_save.side_effect = lambda: events.append("wait") - backend.get_finished.side_effect = lambda _: (set(), set()) + backend.get_transfer_results.return_value = KVConnectorTransferResults() backend.get_block_ids_with_load_errors.return_value = set() backend.get_kv_connector_stats.return_value = None backend.get_kv_connector_kv_cache_events.return_value = None @@ -61,7 +63,15 @@ def test_load_start_phase( connector = _make_connector(monkeypatch, events) output = _scheduler_output(has_sync_kv_loads) - connector.pre_forward(output) # type: ignore[arg-type] + request_indices = torch.tensor([3, 1]) + request_ids = ["first", "second"] + attn_metadata = {"layer": object()} + connector.pre_forward( # type: ignore[arg-type] + output, + request_state_indices=request_indices, + request_ids=request_ids, + attn_metadata=attn_metadata, + ) assert events == ( ["handle", "bind", "start"] if has_sync_kv_loads else ["handle", "bind"] ) @@ -69,6 +79,16 @@ def test_load_start_phase( connector.post_forward(set()) assert events == ["handle", "bind", "start", "wait", "clear"] + kwargs = connector.kv_connector.start_load_kv.call_args.kwargs + assert kwargs["request_state_indices"] is request_indices + assert kwargs["request_ids"] is request_ids + assert kwargs["attn_metadata"] is attn_metadata + + # A subsequent step without a forward must not reuse the prior batch. + connector.no_forward(_scheduler_output(False)) # type: ignore[arg-type] + assert connector.kv_connector.start_load_kv.call_count == 2 + assert connector.kv_connector.start_load_kv.call_args.kwargs == {} + def test_no_forward_starts_deferred_load_once(monkeypatch: pytest.MonkeyPatch): events: list[str] = [] diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 99331e7b35..bb77f90ced 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -312,6 +312,10 @@ def test_select_common_block_size_uses_largest_shared_int(): assert selected_size == 64 +def test_select_common_block_size_without_active_backends_uses_manager_size(): + assert select_common_block_size(256, []) == 256 + + def test_select_common_block_size_accepts_rocm_sparse_block_size_16(monkeypatch): monkeypatch.setattr(current_platform, "is_rocm", lambda: True) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 1118ce589c..c9ce409f5c 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -237,6 +237,7 @@ def _make_capture_runner(captured: bool) -> GPUModelRunner: runner.attn_groups = None runner.kv_cache_config = None runner.use_aux_hidden_state_outputs = False + runner.kv_connector = model_runner_module.NO_OP_KV_CONNECTOR return runner diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py index fcace4ae4a..067f31e935 100644 --- a/tests/v1/worker/test_gpu_warmup_blocks.py +++ b/tests/v1/worker/test_gpu_warmup_blocks.py @@ -89,6 +89,7 @@ def _make_runner( return SimpleNamespace( num_speculative_steps=num_spec_steps, adaptive_verification=None, + rejection_sampler=None, decode_query_len=num_spec_steps + 1, is_pooling_model=False, is_encoder_decoder=False, diff --git a/tests/v1/worker/test_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py index a9579f1af6..bacd1cad84 100644 --- a/tests/v1/worker/test_kv_cache_allocation_scope.py +++ b/tests/v1/worker/test_kv_cache_allocation_scope.py @@ -5,8 +5,10 @@ from contextlib import AbstractContextManager from types import SimpleNamespace from typing import Any, cast +import pytest import torch +import vllm.v1.hisparse.binding as hisparse_binding import vllm.v1.worker.gpu.attn_utils as attn_utils import vllm.v1.worker.gpu_model_runner as gpu_model_runner from vllm.v1.worker.gpu_worker import Worker @@ -26,7 +28,8 @@ class _AllocationScope(AbstractContextManager): self.active = False -def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None: +@pytest.mark.parametrize("hisparse", [False, True]) +def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch, hisparse) -> None: scope = _AllocationScope() kv_caches = {"layer": torch.empty(0)} @@ -38,23 +41,48 @@ def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None: assert not scope.active monkeypatch.setattr(attn_utils, "allocate_kv_cache", allocate) - monkeypatch.setattr(attn_utils, "bind_kv_cache", bind) + monkeypatch.setattr(attn_utils, "bind_kv_cache_to_layers", bind) monkeypatch.setattr(attn_utils, "get_shared_kv_cache_layers", lambda config: {}) + hisparse_bindings = [] + if hisparse: + host_pool = SimpleNamespace() + + def bind_hisparse(**kwargs): + hisparse_bindings.append(kwargs) + assert scope.active + assert kwargs["kv_caches"] is kv_caches + assert kwargs["host_pool"] is host_pool + return [ + SimpleNamespace( + view=SimpleNamespace(cache=torch.empty(1, 1, 1), block_size=1), + runtime=SimpleNamespace(), + ) + ] + + monkeypatch.setattr(hisparse_binding, "HiSparseHostPool", lambda: host_pool) + monkeypatch.setattr(hisparse_binding, "allocate_hisparse_kv_caches", allocate) + monkeypatch.setattr(hisparse_binding, "bind_hisparse_kv_caches", bind_hisparse) + config = SimpleNamespace( + attention_config=SimpleNamespace( + hisparse_config=object() if hisparse else None + ), + scheduler_config=SimpleNamespace(max_num_seqs=1, max_num_batched_tokens=1), cache_config=SimpleNamespace(get_resolved_kv_cache_layout=lambda: None), model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type="test")), ) result = attn_utils.init_kv_cache( - [], {}, SimpleNamespace(kv_cache_groups=[]), torch.device("cpu"), [], config, kv_cache_allocation_context=scope, + block_tables=SimpleNamespace(), ) + assert len(hisparse_bindings) == int(hisparse) assert result is kv_caches assert not scope.active diff --git a/tests/v1/worker/test_mixed_warmup_gate.py b/tests/v1/worker/test_mixed_warmup_gate.py index df339135eb..04f456a7c9 100644 --- a/tests/v1/worker/test_mixed_warmup_gate.py +++ b/tests/v1/worker/test_mixed_warmup_gate.py @@ -35,10 +35,15 @@ def test_mixed_warmup_skipped_for_single_seq(max_num_reqs): def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_warmup): """Startup must warm fixed drafts before calibration and retain its manager.""" manager = SimpleNamespace(cost_tables=None) - runner = SimpleNamespace(adaptive_verification=manager) + rejection_sampler = SimpleNamespace(enable_adaptive_verification=True) + runner = SimpleNamespace( + adaptive_verification=manager, + rejection_sampler=rejection_sampler, + ) def run_steps(model_runner, execute, sample): assert model_runner.adaptive_verification is None + assert not model_runner.rejection_sampler.enable_adaptive_verification if fail_warmup: raise RuntimeError("warmup failed") @@ -50,3 +55,4 @@ def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_ warmup.warmup_kernels(runner, _fail, _fail) assert runner.adaptive_verification is manager assert manager.cost_tables is None + assert rejection_sampler.enable_adaptive_verification diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index a4474c8a58..86c8688906 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -1,13 +1,897 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections import deque +from contextlib import nullcontext from types import SimpleNamespace +from unittest.mock import MagicMock, call +import numpy as np +import pytest import torch +import vllm.v1.attention.backends.mla.index_group as index_group_module +import vllm.v1.hisparse.runtime as hisparse_runtime_module +from vllm.config import CUDAGraphMode from vllm.config.mamba import MambaBackendEnum, MambaConfig +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import ( + worker as hisparse_worker_module, +) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( + HiSparseConnectorWorker, + _flatten_row_mirrors, + _select_written_row_mirrors, + _SlotMappingStaging, +) from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 -from vllm.v1.worker.utils import bind_kv_cache +from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy +from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror +from vllm.v1.worker.utils import ( + bind_kv_cache, + bind_kv_cache_to_layers, + copy_kv_cache_blocks_inplace, +) + + +def _make_hisparse_worker() -> HiSparseConnectorWorker: + worker = object.__new__(HiSparseConnectorWorker) + worker.refines_row_mirrors = False + worker._slot_mapping_staging = None + worker._row_mirror_num_rows = 0 + worker._per_layer_mirrored = set() + worker._submitted_mirror_layers = set() + worker._pending_dma_descriptors = deque() + worker._dma_free_descriptors = [] + worker.host_write_events = (MagicMock(), MagicMock()) + worker.host_write_event = worker.host_write_events[1] + worker._next_host_write_event = 0 + worker.dma_stream = None + worker.shared_host_region = None + worker._completed_host_copy_dst_ids = [] + worker._slot_mapping_staging = None + return worker + + +def test_hisparse_row_mirrors_follow_runner_request_order(): + first = SparseKVRowMirror((1,), 10, 1) + second = SparseKVRowMirror((2,), 20, 1) + row_mirrors = {"first": (first,), "second": (second,)} + + assert _flatten_row_mirrors(row_mirrors, ["second", "first"]) == ( + second, + first, + ) + + +def test_hisparse_gpu_slots_select_written_rows_from_async_envelope(): + """Rejected speculative rows must never become durable host KV.""" + candidates = ( + SparseKVRowMirror((100, 500), 1000, 4), + SparseKVRowMirror((20, 300), 2000, 4), + ) + source_slots = np.array([102, 103, -1, 21, 22, 23], dtype=np.int64) + + assert _select_written_row_mirrors(candidates, source_slots, 0) == ( + SparseKVRowMirror((102, 502), 1002, 2), + SparseKVRowMirror((21, 301), 2001, 3), + ) + + +def test_hisparse_written_rows_preserve_scheduler_page_boundaries(): + candidates = ( + SparseKVRowMirror((0,), 100, 4), + SparseKVRowMirror((4,), 104, 4), + ) + + assert _select_written_row_mirrors( + candidates, np.array([2, 3, 4, 5], dtype=np.int64), 0 + ) == ( + SparseKVRowMirror((2,), 102, 2), + SparseKVRowMirror((4,), 104, 2), + ) + + +def test_hisparse_written_rows_without_host_destinations_are_skipped(): + assert ( + _select_written_row_mirrors((), np.array([100, 101], dtype=np.int64), 0) == () + ) + + +def test_hisparse_appends_reference_slots_within_a_mirror_phase(monkeypatch): + """Separate context and query writes must share one mirror phase.""" + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.refines_row_mirrors = True + state = _SlotMappingStaging( + MagicMock(), MagicMock(), torch.empty(4, dtype=torch.int64) + ) + worker._slot_mapping_staging = state + handle = SimpleNamespace(runtime=SimpleNamespace(resident_source_index=1)) + worker.cache_layer_names = ["layer"] + worker.cache_handles = [handle] + worker._layer_mirror_callbacks = (MagicMock(),) + worker._row_mirror_num_rows = 1 + worker._submitted_mirror_layers = set() + main_stream = MagicMock() + monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: main_stream) + monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext()) + + handle.slot_mapping = torch.tensor([7, 8], dtype=torch.int64) + worker._stage_row_mirror_mapping(2) + handle.slot_mapping = torch.tensor([9, 10], dtype=torch.int64) + worker._stage_row_mirror_mapping(2) + + torch.testing.assert_close(state.slots, torch.tensor([7, 8, 9, 10])) + assert state.stream.wait_stream.call_args_list == [call(main_stream)] * 2 + assert state.event.record.call_args_list == [call(state.stream)] * 2 + assert state.num_tokens == 4 + assert state.source_index == 1 + + +@pytest.mark.parametrize( + ("mode", "expected_calls"), + [(CUDAGraphMode.PIECEWISE, 1), (CUDAGraphMode.FULL, 0)], +) +def test_hisparse_submits_layer_mirror_at_replayed_attention_boundary( + monkeypatch, mode, expected_calls +): + """Piecewise replay must pipeline DMA; FULL replay uses the batch fallback.""" + handle = object.__new__(hisparse_runtime_module.HiSparseCacheHandle) + handle.dummy_batch = False + handle.decode_batch = False + handle.submit_layer_mirror = MagicMock() + monkeypatch.setattr( + hisparse_runtime_module, + "get_forward_context", + lambda: SimpleNamespace(cudagraph_runtime_mode=mode), + ) + + handle.finish_kv_update() + + assert handle.submit_layer_mirror.call_count == expected_calls + + +def test_copy_cpu_kv_cache_logical_blocks_ignores_storage_padding(): + backing = torch.full((10, 2, 3), -1, dtype=torch.float32) + cache = backing[1:9] + cache[2:4] = 7 + cache[6:8] = 11 + + copy_kv_cache_blocks_inplace( + [cache], + num_blocks=4, + kv_cache_block_copies=[ + KVCacheBlockCopy(1, 0), + KVCacheBlockCopy(3, 2), + ], + ) + + torch.testing.assert_close(cache[0:2], torch.full_like(cache[0:2], 7)) + torch.testing.assert_close(cache[4:6], torch.full_like(cache[4:6], 11)) + assert (backing[0] == -1).all() + assert (backing[9] == -1).all() + + +@pytest.mark.parametrize("shared,rank", [(False, 0), (True, 0), (True, 1)]) +@pytest.mark.parametrize("has_copies", [False, True]) +def test_hisparse_host_copy_waits_for_writes_on_writer_only( + monkeypatch, shared, rank, has_copies +): + """Host CoW must read completed writes; shared non-writers only barrier.""" + worker = _make_hisparse_worker() + cache = torch.zeros((2, 1, 1)) + worker.host_caches = [cache] + worker.host_num_blocks = 2 + worker.shared_host_region = object() if shared else None + event = MagicMock() + event.synchronize.side_effect = lambda: cache[0].fill_(7) + barrier = MagicMock() + monkeypatch.setattr( + hisparse_worker_module, "get_tensor_model_parallel_rank", lambda: rank + ) + monkeypatch.setattr( + hisparse_worker_module, "get_tp_group", lambda: SimpleNamespace(barrier=barrier) + ) + copies = [KVCacheBlockCopy(0, 1)] if has_copies else [] + + worker._copy_host_blocks(copies, event) + + writes = has_copies and (not shared or rank == 0) + assert cache[1].item() == (7 if writes else 0) + assert event.synchronize.call_count == int(writes) + assert barrier.call_count == int(shared and has_copies) + assert worker._completed_host_copy_dst_ids == ([1] if has_copies else []) + + +def test_hisparse_worker_updates_request_state_mapping_in_place(monkeypatch): + worker = _make_hisparse_worker() + worker.request_state_indices = torch.arange(4, dtype=torch.int32) + worker._pending_invalid_block_ids = [5] + invalidations = [] + worker.invalidate_blocks = lambda blocks, states: invalidations.append( + (blocks.copy(), states.clone()) + ) + original_ptr = worker.request_state_indices.data_ptr() + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + + worker.set_request_state_indices(torch.tensor([3, 1], dtype=torch.int32)) + + assert worker.request_state_indices.data_ptr() == original_ptr + assert worker.request_state_indices.tolist() == [3, 1, -1, -1] + assert len(invalidations) == 1 + assert invalidations[0][0] == [5] + torch.testing.assert_close( + invalidations[0][1], torch.tensor([3, 1], dtype=torch.int32) + ) + assert worker._pending_invalid_block_ids == [] + + +def test_hisparse_pre_forward_transfer_builds_page_descriptors(): + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.kernel_block_size = 2 + source = torch.empty((3, 2, 4), dtype=torch.uint8) + destination = torch.empty((6, 4), dtype=torch.uint8) + worker.resident_caches = (source,) + worker.host_caches = (destination,) + worker.cache_handles = [ + SimpleNamespace(runtime=SimpleNamespace(resident_source_index=0)) + ] + worker._dma_free_descriptors = [] + worker._submit_dma_descriptors = MagicMock() + + worker._enqueue_transfers([SparseKVPageTransfer(7, 2, (1,), after_forward=False)]) + + descriptors, count = worker._submit_dma_descriptors.call_args.args + assert count == 1 + assert descriptors.src[:count].tolist() == [ + source.data_ptr() + source.stride(0) * source.element_size() + ] + assert descriptors.dst[:count].tolist() == [destination.data_ptr() + 4 * 4] + assert descriptors.sizes[:count].tolist() == [2 * 4] + assert worker._submit_dma_descriptors.call_args.kwargs["transfer_ids"] == (7,) + + +def test_hisparse_eager_mirror_records_transfer_without_page_copy(): + worker = _make_hisparse_worker() + worker.cache_handles = [ + SimpleNamespace(runtime=SimpleNamespace(eager_host_mirror=True)) + ] + worker._record_transfer_completion = MagicMock() + worker._enqueue_transfers = MagicMock() + transfers = [SparseKVPageTransfer(7, 2, (1,), after_forward=False)] + + worker._submit_transfers(transfers) + + worker._record_transfer_completion.assert_called_once_with(transfers) + worker._enqueue_transfers.assert_not_called() + + +def test_hisparse_lazy_mirror_copies_transfer_pages(): + worker = _make_hisparse_worker() + worker.cache_handles = [ + SimpleNamespace(runtime=SimpleNamespace(eager_host_mirror=False)) + ] + worker._record_transfer_completion = MagicMock() + worker._enqueue_transfers = MagicMock() + transfers = [SparseKVPageTransfer(7, 2, (1,), after_forward=False)] + + worker._submit_transfers(transfers) + + worker._enqueue_transfers.assert_called_once_with(transfers) + worker._record_transfer_completion.assert_not_called() + + +def test_hisparse_dma_row_mirror_builds_descriptors(monkeypatch): + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.kernel_block_size = 2 + source = torch.empty((3, 2, 4), dtype=torch.uint8) + destination = torch.empty((6, 4), dtype=torch.uint8) + worker.resident_caches = (source,) + worker.host_caches = (destination,) + worker.cache_handles = [ + SimpleNamespace( + runtime=SimpleNamespace(resident_source_index=0), decode_batch=True + ) + ] + worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0")) + worker.dma_stream = MagicMock() + worker.host_write_event = MagicMock() + worker._dma_free_descriptors = [] + worker._pending_dma_descriptors = deque() + worker._pending_transfer_events = deque() + worker._enqueued_transfer_ids = [] + worker._dma_submitted = False + current_stream = MagicMock() + event = MagicMock() + swap_blocks_batch = MagicMock() + monkeypatch.setattr( + hisparse_worker_module, "current_stream", lambda: current_stream + ) + monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext()) + monkeypatch.setattr(torch, "Event", lambda: event) + monkeypatch.setattr( + hisparse_worker_module.ops, "swap_blocks_batch", swap_blocks_batch + ) + + worker._set_row_mirrors((SparseKVRowMirror((3,), 4, 1),)) + worker._enqueue_row_dma(range(1)) + + src_ptrs, dst_ptrs, sizes = swap_blocks_batch.call_args.args + assert src_ptrs.tolist() == [source.data_ptr() + 3 * 4] + assert dst_ptrs.tolist() == [destination.data_ptr() + 4 * 4] + assert sizes.tolist() == [4] + worker.dma_stream.wait_stream.assert_called_once_with(current_stream) + worker.host_write_event.record.assert_called_once_with(worker.dma_stream) + event.record.assert_called_once_with(worker.dma_stream) + assert worker._dma_submitted + + +def test_hisparse_row_dma_uses_resident_spans(): + worker = _make_hisparse_worker() + source = torch.empty((2, 2, 4), dtype=torch.uint8) + destination = torch.empty((12, 4), dtype=torch.uint8) + worker.is_host_writer = True + worker.kernel_block_size = 2 + worker.resident_caches = (source,) + worker.host_caches = (destination,) + worker.cache_handles = [ + SimpleNamespace( + runtime=SimpleNamespace(resident_source_index=0), decode_batch=True + ) + ] + worker._set_row_mirrors( + ( + SparseKVRowMirror((0,), 4, 2), + SparseKVRowMirror((2,), 9, 1), + ) + ) + worker._dma_free_descriptors = [] + worker._submit_dma_descriptors = MagicMock() + + worker._enqueue_row_dma(range(1)) + + descriptors, count = worker._submit_dma_descriptors.call_args.args + assert count == 2 + assert descriptors.src[:count].tolist() == [ + source.data_ptr(), + source.data_ptr() + 2 * 4, + ] + assert descriptors.dst[:count].tolist() == [ + destination.data_ptr() + 4 * 4, + destination.data_ptr() + 9 * 4, + ] + assert descriptors.sizes[:count].tolist() == [2 * 4, 4] + + +def test_hisparse_finish_forward_mirrors_all_layers_once(monkeypatch): + dst_slots = torch.tensor([7, 8, 9], dtype=torch.int64) + req_ids = torch.tensor([0, 1, 2], dtype=torch.int32) + leader = SimpleNamespace( + eager_host_mirror=True, + is_group_leader=True, + invalidate_written_slots=MagicMock(), + ) + follower = SimpleNamespace( + eager_host_mirror=True, + is_group_leader=False, + invalidate_written_slots=MagicMock(), + ) + handles = [ + SimpleNamespace( + runtime=runtime, + decode_batch=True, + host_mirror_required=True, + num_actual_tokens=3, + num_decode_tokens=2, + req_id_per_token=req_ids, + mirror_slot_mapping=dst_slots, + ) + for runtime in (leader, follower) + ] + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.cache_handles = handles + worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0")) + worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 3),)) + worker._dma_submitted = False + worker._per_layer_mirrored = set() + worker._submitted_mirror_layers = set() + worker._post_forward_transfers = [] + worker._enqueue_row_dma = MagicMock() + worker.host_write_event = MagicMock() + worker._forward_ready_event = MagicMock() + current_stream = MagicMock() + monkeypatch.setattr( + hisparse_worker_module, "current_stream", lambda: current_stream + ) + + worker.finish_forward() + + worker._forward_ready_event.record.assert_called_once_with() + worker._enqueue_row_dma.assert_called_once_with( + (0, 1), ready_event=worker._forward_ready_event + ) + leader.invalidate_written_slots.assert_called_once() + torch.testing.assert_close( + leader.invalidate_written_slots.call_args.args[0], dst_slots + ) + follower.invalidate_written_slots.assert_not_called() + worker.host_write_event.record.assert_called_once_with(current_stream) + + +def test_hisparse_finish_forward_does_not_repeat_per_layer_mirrors(): + slots = torch.tensor([7, 8], dtype=torch.int64) + runtime = SimpleNamespace( + eager_host_mirror=False, + is_group_leader=False, + invalidate_written_slots=MagicMock(), + ) + handles = [ + SimpleNamespace( + runtime=runtime, + decode_batch=False, + host_mirror_required=True, + num_actual_tokens=2, + num_decode_tokens=0, + req_id_per_token=torch.empty(0, dtype=torch.int32), + mirror_slot_mapping=slots, + ) + for _ in range(2) + ] + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.cache_handles = handles + worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 2),)) + worker._per_layer_mirrored = {0, 1} + worker._submitted_mirror_layers = {0, 1} + worker._enqueue_row_dma = MagicMock() + + worker._enqueue_host_mirror() + + worker._enqueue_row_dma.assert_not_called() + + +def test_hisparse_finish_forward_submits_lazy_post_forward_transfer(monkeypatch): + runtime = SimpleNamespace(eager_host_mirror=False) + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.cache_handles = [SimpleNamespace(runtime=runtime, num_actual_tokens=0)] + transfer = SparseKVPageTransfer(7, 2, (1,), after_forward=True) + worker._post_forward_transfers = [transfer] + worker._forward_ready_event = MagicMock() + worker._enqueue_host_mirror = MagicMock() + worker._submit_transfers = MagicMock() + worker._dma_submitted = False + worker._submitted_mirror_layers = set() + worker._finish_mirror_phase = MagicMock() + worker._release_completed_dma_descriptors = MagicMock() + stream = MagicMock() + monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream) + + worker.finish_forward() + + worker._finish_mirror_phase.assert_called_once_with(worker._forward_ready_event) + worker._submit_transfers.assert_called_once_with([transfer]) + worker.host_write_event.record.assert_called_once_with(stream) + + +def test_hisparse_prefill_mirrors_source_groups_and_flushes_partial_group(): + slots = torch.tensor([7, 8], dtype=torch.int64) + source_indices = [0, 0, 1, 1, 1, 1, 2] + handles = [ + SimpleNamespace( + runtime=SimpleNamespace( + eager_host_mirror=False, + is_group_leader=False, + resident_source_index=source_index, + ), + decode_batch=False, + host_mirror_required=layer_index < 5, + num_actual_tokens=2 if layer_index < 5 else 0, + num_decode_tokens=0, + req_id_per_token=torch.empty(0, dtype=torch.int32), + mirror_slot_mapping=slots, + ) + for layer_index, source_index in enumerate(source_indices) + ] + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.cache_handles = handles + worker._set_row_mirrors((SparseKVRowMirror((0, 0, 0), 7, 2),)) + worker._per_layer_mirrored = set() + worker._submitted_mirror_layers = set() + worker._layer_ready_events = tuple(MagicMock() for _ in handles) + worker._enqueue_row_dma = MagicMock() + + for layer_index in range(5): + worker._enqueue_layer_mirror(layer_index) + + worker._enqueue_row_dma.assert_called_once_with( + (0, 1), ready_event=worker._layer_ready_events[1] + ) + worker._enqueue_host_mirror(ready_event=worker._layer_ready_events[-1]) + assert worker._enqueue_row_dma.call_args_list[1].args == ((2, 3, 4),) + assert worker._enqueue_row_dma.call_args_list[1].kwargs == { + "ready_event": worker._layer_ready_events[-1] + } + + +def test_hisparse_finish_forward_rejects_partial_per_layer_mirror(): + slots = torch.tensor([7, 8], dtype=torch.int64) + runtime = SimpleNamespace( + eager_host_mirror=False, + is_group_leader=False, + invalidate_written_slots=MagicMock(), + ) + handles = [ + SimpleNamespace( + runtime=runtime, + decode_batch=False, + host_mirror_required=True, + num_actual_tokens=2, + num_decode_tokens=0, + req_id_per_token=torch.empty(0, dtype=torch.int32), + mirror_slot_mapping=slots, + ) + for _ in range(2) + ] + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.cache_handles = handles + worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 2),)) + worker._per_layer_mirrored = {0} + worker._submitted_mirror_layers = set() + worker._enqueue_row_dma = MagicMock() + + with pytest.raises(RuntimeError, match="did not mirror every active layer"): + worker._enqueue_host_mirror() + + +def test_hisparse_finish_forward_orders_next_forward_after_dma(monkeypatch): + current_stream = MagicMock() + worker = _make_hisparse_worker() + worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0")) + worker.is_host_writer = True + worker._dma_submitted = True + worker.host_write_event = MagicMock() + worker.cache_handles = [ + SimpleNamespace( + runtime=SimpleNamespace(eager_host_mirror=True), num_actual_tokens=0 + ) + ] + worker._post_forward_transfers = [] + worker._pending_dma_descriptors = deque() + worker._dma_free_descriptors = [] + worker._forward_ready_event = MagicMock() + worker._finish_mirror_phase = MagicMock() + monkeypatch.setattr( + hisparse_worker_module, "current_stream", lambda: current_stream + ) + + worker.finish_forward() + + current_stream.wait_event.assert_called_once_with(worker.host_write_event) + assert not worker._dma_submitted + + +def test_hisparse_finish_forward_mirrors_standalone_mtp_cache(monkeypatch): + dst_slots = torch.tensor([7, 8], dtype=torch.int64) + runtime = SimpleNamespace( + eager_host_mirror=True, + is_group_leader=False, + invalidate_written_slots=MagicMock(), + ) + inactive = SimpleNamespace( + runtime=runtime, + decode_batch=False, + host_mirror_required=False, + num_actual_tokens=0, + num_decode_tokens=0, + req_id_per_token=None, + mirror_slot_mapping=dst_slots, + ) + mtp = SimpleNamespace( + runtime=runtime, + decode_batch=True, + host_mirror_required=True, + num_actual_tokens=2, + num_decode_tokens=2, + req_id_per_token=torch.tensor([0, 1], dtype=torch.int32), + mirror_slot_mapping=dst_slots, + ) + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.cache_handles = [inactive, mtp] + worker._set_row_mirrors((SparseKVRowMirror((0,), 7, 2),)) + worker._per_layer_mirrored = set() + worker._submitted_mirror_layers = set() + worker._enqueue_row_dma = MagicMock() + + worker._enqueue_host_mirror() + + worker._enqueue_row_dma.assert_called_once_with((1,), ready_event=None) + + +def test_hisparse_shared_host_reader_skips_mirror(monkeypatch): + """A non-writer TP rank must not mirror rows into the shared host pool.""" + dst_slots = torch.tensor([7, 8], dtype=torch.int64) + leader = SimpleNamespace( + eager_host_mirror=True, + is_group_leader=True, + invalidate_written_slots=MagicMock(), + ) + handle = SimpleNamespace( + runtime=leader, + decode_batch=True, + host_mirror_required=True, + num_actual_tokens=2, + num_decode_tokens=1, + req_id_per_token=torch.tensor([0], dtype=torch.int32), + mirror_slot_mapping=dst_slots, + ) + worker = _make_hisparse_worker() + worker.is_host_writer = False + worker.cache_handles = [handle] + worker._set_row_mirrors((SparseKVRowMirror((0,), 7, 2),)) + worker._per_layer_mirrored = set() + worker._enqueue_row_dma = MagicMock() + + worker._enqueue_host_mirror() + + worker._enqueue_row_dma.assert_not_called() + leader.invalidate_written_slots.assert_called_once() + + +def test_hisparse_shared_host_reader_skips_transfer_completion(): + """A non-writer TP rank must not acknowledge host writes.""" + worker = _make_hisparse_worker() + worker.is_host_writer = False + worker.kernel_block_size = 1 + worker._enqueued_transfer_ids = [] + worker._pending_transfer_events = [] + + worker._record_transfer_completion([SparseKVPageTransfer(1, 2, (3,), True)]) + + assert worker._enqueued_transfer_ids == [] + assert worker._pending_transfer_events == [] + + +def test_hisparse_writer_records_transfer_completion_after_dma(monkeypatch): + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.dma_stream = MagicMock() + worker._enqueued_transfer_ids = [] + worker._pending_transfer_events = deque() + event = MagicMock() + monkeypatch.setattr(torch, "Event", lambda: event) + + worker._record_transfer_completion( + [ + SparseKVPageTransfer(3, 2, (1,), False), + SparseKVPageTransfer(7, 4, (5,), True), + ] + ) + + event.record.assert_called_once_with(worker.dma_stream) + assert worker._enqueued_transfer_ids == [3, 7] + assert worker._pending_transfer_events == deque([(event, (3, 7))]) + + +@pytest.mark.parametrize("is_host_writer", [False, True]) +def test_hisparse_step_waits_for_previous_host_write(monkeypatch, is_host_writer): + worker = _make_hisparse_worker() + worker.is_host_writer = is_host_writer + worker.hot_backing = SimpleNamespace(device=torch.device("cuda:1")) + host_write_events = (MagicMock(), MagicMock()) + worker.host_write_events = host_write_events + worker.host_write_event = host_write_events[1] + worker._next_host_write_event = 0 + worker.host_caches = () + worker.host_num_blocks = 1 + worker._post_forward_transfers = [] + worker._pending_invalid_block_ids = [] + worker.cache_handles = [ + SimpleNamespace( + runtime=SimpleNamespace(eager_host_mirror=True), + decode_batch=False, + host_mirror_required=False, + num_actual_tokens=0, + num_decode_tokens=0, + req_id_per_token=None, + ) + ] + worker._per_layer_mirrored = set() + worker._submitted_mirror_layers = set() + worker._layer_mirror_callbacks = () + stream = MagicMock() + monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream) + + worker.start_step( + SimpleNamespace( + host_block_copies=[], + command=None, + source_block_ids=[], + row_mirrors={}, + all_context_pages_resident=True, + row_mirrors_from_resident=False, + ), + None, + ) + worker.start_step( + SimpleNamespace( + host_block_copies=[], + command=None, + source_block_ids=[], + row_mirrors={}, + all_context_pages_resident=True, + row_mirrors_from_resident=False, + ), + None, + ) + + assert stream.wait_event.call_args_list == [ + call(host_write_events[1]), + call(host_write_events[0]), + ] + assert worker.host_write_event is host_write_events[1] + + +@pytest.mark.parametrize("tp_rank, expected_copies", [(0, 1), (1, 0)]) +def test_hisparse_shared_host_block_copy_has_one_writer( + monkeypatch, tp_rank, expected_copies +): + worker = _make_hisparse_worker() + worker.shared_host_region = object() + worker.host_caches = (torch.empty(4, 1),) + worker.host_num_blocks = 4 + previous_event = MagicMock() + copies = (KVCacheBlockCopy(0, 1),) + copy_blocks = MagicMock() + tp_group = MagicMock() + monkeypatch.setattr( + hisparse_worker_module, "get_tensor_model_parallel_rank", lambda: tp_rank + ) + monkeypatch.setattr(hisparse_worker_module, "get_tp_group", lambda: tp_group) + monkeypatch.setattr( + hisparse_worker_module, "copy_kv_cache_blocks_inplace", copy_blocks + ) + + worker._copy_host_blocks(copies, previous_event) + + assert copy_blocks.call_count == expected_copies + tp_group.barrier.assert_called_once_with() + assert worker.take_completed_host_copies() == [1] + assert worker.take_completed_host_copies() == [] + + +def test_hisparse_empty_step_does_not_replay_stale_host_mirror(monkeypatch): + handle = SimpleNamespace( + runtime=SimpleNamespace(eager_host_mirror=True), + decode_batch=True, + num_actual_tokens=2, + num_decode_tokens=2, + req_id_per_token=torch.tensor([0, 1]), + mirror_slot_mapping=torch.tensor([4, 5]), + ) + worker = _make_hisparse_worker() + worker.is_host_writer = True + worker.hot_backing = SimpleNamespace(device=torch.device("cpu")) + worker.host_write_event = MagicMock() + worker._forward_ready_event = MagicMock() + worker.host_caches = () + worker.host_num_blocks = 1 + worker.cache_handles = [handle] + worker._per_layer_mirrored = set() + worker._submitted_mirror_layers = set() + worker._layer_mirror_callbacks = (MagicMock(),) + worker._post_forward_transfers = [] + worker._pending_invalid_block_ids = [] + worker._enqueue_host_mirror = MagicMock(wraps=worker._enqueue_host_mirror) + stream = MagicMock() + monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream) + + worker.start_step( + SimpleNamespace( + host_block_copies=[], + command=None, + source_block_ids=[], + row_mirrors={}, + all_context_pages_resident=True, + row_mirrors_from_resident=False, + ), + None, + ) + worker.finish_forward() + + worker._enqueue_host_mirror.assert_called_once_with(worker._forward_ready_event) + assert handle.num_actual_tokens == 0 + torch.testing.assert_close(handle.mirror_slot_mapping, torch.tensor([4, 5])) + + +def test_hisparse_cache_handles_join_index_groups_during_construction(monkeypatch): + """Followers must not allocate duplicate runtime state before profiling.""" + config = SimpleNamespace( + scheduler_config=SimpleNamespace( + max_num_seqs=2, + max_num_batched_tokens=2, + async_scheduling=False, + ), + speculative_config=None, + kv_transfer_config=None, + ) + resolved = hisparse_runtime_module.ResolvedHiSparseConfig( + top_k=4, + device_buffer_size=8, + ) + monkeypatch.setattr(hisparse_runtime_module, "_has_hisparse_ops", lambda: True) + monkeypatch.setattr( + hisparse_runtime_module.ResolvedHiSparseConfig, + "from_vllm_config", + classmethod(lambda cls, vllm_config, model_top_k: resolved), + ) + shared_states: list[object] = [] + streams: list[object] = [] + + def create_shared_state(_device, _max_rows, _top_k): + shared_states.append(object()) + return shared_states[-1] + + def create_stream(_device): + streams.append(object()) + return streams[-1] + + monkeypatch.setattr( + hisparse_runtime_module, "_create_shared_topk_state", create_shared_state + ) + monkeypatch.setattr(hisparse_runtime_module, "_create_copy_stream", create_stream) + monkeypatch.setattr(index_group_module, "_create_side_stream", lambda _: object()) + monkeypatch.setattr(index_group_module, "_create_event", lambda: object()) + index_group_builder = index_group_module.SparseMLAIndexGroupBuilder( + torch.empty((2, 4), dtype=torch.int32) + ) + + def make_cache_handle(is_leader: bool): + index_group, _ = index_group_builder.register_layer(is_leader) + cache_handle = hisparse_runtime_module.create_hisparse_cache_handle( + config, + model_top_k=4, + is_index_group_leader=is_leader, + row_width=8, + kv_dtype=torch.float32, + index_group=index_group, + device="cpu", + ) + assert cache_handle is not None + return cache_handle + + first_leader = make_cache_handle(True) + first_index_group = index_group_builder.current_group + first_follower = make_cache_handle(False) + second_leader = make_cache_handle(True) + second_index_group = index_group_builder.current_group + second_follower = make_cache_handle(False) + + assert first_index_group is not None + assert second_index_group is not None + assert first_follower.runtime.index_group is first_leader.runtime.index_group + assert second_follower.runtime.index_group is second_leader.runtime.index_group + assert first_leader.runtime.index_group is not second_leader.runtime.index_group + assert first_leader.runtime.index_group.copy_stream is first_index_group.side_stream + assert ( + first_leader.runtime.index_group.logical_topk_ready + is first_index_group.logical_topk_ready + ) + assert ( + second_leader.runtime.index_group.copy_stream is second_index_group.side_stream + ) + assert ( + second_leader.runtime.index_group.logical_topk_ready + is second_index_group.logical_topk_ready + ) + assert len(shared_states) == 2 + assert streams == [] class _TestReplaySSMMixer(MambaMixer2): @@ -30,7 +914,8 @@ def _packed_replayssm_cache(num_blocks: int) -> torch.Tensor: return torch.full((num_blocks, 1, 1, 80), 0, dtype=torch.int8) -def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(): +@pytest.mark.parametrize("layers_only", [False, True]) +def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(layers_only): mixers = [_TestReplaySSMMixer() for _ in range(3)] layer_names = [f"layers.{i}.mixer" for i in range(3)] ctx = dict(zip(layer_names, mixers)) @@ -45,7 +930,10 @@ def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(): SimpleNamespace(layer_names=[layer_names[1]]), ] - bind_kv_cache(kv_cache, ctx, [], kv_cache_groups=kv_cache_groups) + if layers_only: + bind_kv_cache_to_layers(kv_cache, ctx, kv_cache_groups=kv_cache_groups) + else: + bind_kv_cache(kv_cache, ctx, [], kv_cache_groups=kv_cache_groups) assert ( mixers[0]._replayssm_ring_start.data_ptr() diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index c5fbb76d98..3311215016 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3008,6 +3008,9 @@ def cp_gather_and_upconvert_fp8_kv_cache( workspace_starts: torch.Tensor, batch_size: int, seq_starts: torch.Tensor | None = None, + host_cache: torch.Tensor | None = None, + host_row_ids: torch.Tensor | None = None, + device_row_ids: torch.Tensor | None = None, ) -> None: """Gather and upconvert FP8 KV cache to BF16 workspace. @@ -3018,9 +3021,20 @@ def cp_gather_and_upconvert_fp8_kv_cache( workspace_starts: Workspace start offsets [num_reqs] batch_size: Number of requests seq_starts: Optional source sequence offsets [num_reqs] + host_cache: Optional pinned host rows used for non-resident entries + host_row_ids: Host source row for each remapped cache row + device_row_ids: Device source row, or -1, for each remapped cache row """ torch.ops._C_cache_ops.cp_gather_and_upconvert_fp8_kv_cache( - src_cache, dst, block_table, workspace_starts, batch_size, seq_starts + src_cache, + dst, + block_table, + workspace_starts, + batch_size, + seq_starts, + host_cache, + host_row_ids, + device_row_ids, ) diff --git a/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py b/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py index 5998a03138..e89864fe43 100644 --- a/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py +++ b/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py @@ -35,7 +35,10 @@ def fused_rope_unified_mla_kv_cache_update_impl( layer_name: LayerNameType, ) -> torch.Tensor: layer_name = _resolve_layer_name(layer_name) - attn_metadata, _, kv_cache, layer_slot_mapping = get_attention_context(layer_name) + attn_metadata, attn_layer, kv_cache, layer_slot_mapping = get_attention_context( + layer_name + ) + attn_layer.prepare_kv_cache_update(attn_metadata) if layer_slot_mapping is not None: ops.concat_and_cache_mla_rope_fused( positions, diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index 1cfa372d60..ff2ba0493d 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm.config.attention import AttentionConfig +from vllm.config.attention import AttentionConfig, HiSparseConfig from vllm.config.cache import CacheConfig from vllm.config.compilation import ( CompilationConfig, @@ -68,6 +68,7 @@ from vllm.config.weight_transfer import WeightTransferConfig __all__ = [ # From vllm.config.attention "AttentionConfig", + "HiSparseConfig", # From vllm.config.cache "CacheConfig", # From vllm.config.compilation diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 799614a699..05da048bdb 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -4,7 +4,7 @@ from dataclasses import field from typing import Any, Literal -from pydantic import field_validator +from pydantic import Field, field_validator from vllm.config.utils import config from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum @@ -14,6 +14,25 @@ IndexerKVDType = Literal["auto", "bf16", "fp8", "mxfp4", "nvfp4"] MiniMaxM3MSADecodeBackend = Literal["triton", "cutlass"] +@config +class HiSparseConfig: + """Configuration for HiSparse sparse-MLA KV offloading.""" + + device_buffer_size: int | None = Field(default=None, gt=0) + """Total per-request GPU hot-buffer rows, including the newest-token slot. + + Defaults to one top-k per decode query plus one top-k of LRU slack. The + physical allocation is rounded up to the GPU cache block size selected + from the active backends. + """ + + eager_host_mirror: bool = True + """Mirror decode-written KV rows to the host pool during the forward so + page spills complete without moving data. When disabled, decode rows stay + resident-only and evicted pages are copied to host at spill time. Prefill + rows are always mirrored during the forward.""" + + @config class AttentionConfig: """Configuration for attention mechanisms in vLLM.""" @@ -67,6 +86,10 @@ class AttentionConfig: indexer). Quantized formats (fp8, mxfp4, nvfp4) require indexer kernel support in the backend.""" + hisparse_config: HiSparseConfig | None = None + """HiSparse host-resident KV configuration. Setting this enables experimental + Model Runner V2-only HiSparse sparse-MLA decode hot-buffering.""" + use_non_causal: bool = False """Whether to use non-causal (bidirectional) attention.""" diff --git a/vllm/config/kv_transfer.py b/vllm/config/kv_transfer.py index d1344a5211..b299aae6f0 100644 --- a/vllm/config/kv_transfer.py +++ b/vllm/config/kv_transfer.py @@ -13,6 +13,40 @@ KVConsumer = Literal["kv_consumer", "kv_both"] KVRole = Literal[KVProducer, KVConsumer] +def hisparse_host_pool_gib( + kv_transfer_config: "KVTransferConfig | None", +) -> float | None: + if kv_transfer_config is None: + return None + if kv_transfer_config.kv_connector == "MultiConnector": + connectors = kv_transfer_config.kv_connector_extra_config.get("connectors", []) + else: + connectors = [ + { + "kv_connector": kv_transfer_config.kv_connector, + "kv_connector_extra_config": ( + kv_transfer_config.kv_connector_extra_config + ), + } + ] + entries = [ + connector.get("kv_connector_extra_config", {}) + for connector in connectors + if connector.get("kv_connector") == "HiSparseConnector" + ] + if len(entries) > 1: + raise ValueError("Only one HiSparseConnector may be configured") + if not entries: + return None + host_pool_gib = entries[0].get("host_pool_gib") + if host_pool_gib is None: + raise ValueError("HiSparseConnector requires host_pool_gib") + host_pool_gib = float(host_pool_gib) + if host_pool_gib <= 0: + raise ValueError("HiSparseConnector host_pool_gib must be positive") + return host_pool_gib + + def kv_buffer_device_default_factory() -> str: from vllm.platforms import current_platform diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index ae56e58803..cfecff72b6 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1919,6 +1919,13 @@ class SpeculativeConfig: def uses_draft_model(self) -> bool: return self.method == "draft_model" + def uses_draft_kv_cache(self) -> bool: + return ( + self.use_eagle() + or self.uses_draft_model() + or self.uses_extract_hidden_states() + ) + def uses_extract_hidden_states(self) -> bool: return self.method == "extract_hidden_states" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 05377a2735..3c197df0da 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -666,6 +666,14 @@ class VllmConfig: @property def use_v2_model_runner(self) -> bool: + if self.attention_config.hisparse_config is not None: + if envs.VLLM_USE_V2_MODEL_RUNNER is False: + raise ValueError( + "HiSparse requires Model Runner V2; remove " + "VLLM_USE_V2_MODEL_RUNNER=0." + ) + return True + if getattr(self, "watermark_config", None) is not None: if envs.VLLM_USE_V2_MODEL_RUNNER is False: logger.info_once( @@ -1559,6 +1567,39 @@ class VllmConfig: self._maybe_disable_dynamic_sd_for_data_parallel() self._maybe_override_dynamic_sd_cudagraph_mode() + if self.attention_config.hisparse_config is not None: + if not current_platform.is_cuda(): + raise ValueError("HiSparse currently requires NVIDIA CUDA.") + if self.parallel_config.pipeline_parallel_size > 1: + raise ValueError("HiSparse does not support pipeline parallelism.") + if self.parallel_config.decode_context_parallel_size > 1: + raise ValueError( + "HiSparse does not support decode context parallelism." + ) + if self.model_config is not None and not hasattr( + self.model_config.hf_config, "index_topk" + ): + raise ValueError( + "HiSparse is only supported for DSA models with index_topk." + ) + if self.kv_transfer_config is not None and ( + self.kv_transfer_config.kv_connector + not in ( + None, + "NixlConnector", + "MooncakeStoreConnector", + "MultiConnector", + ) + ): + logger.warning( + "HiSparse host-resident KV is configured with connector " + "%s. NixlConnector (GPU-staged host imports) and " + "MooncakeStoreConnector (shared-store offload) are the " + "validated paths; other connectors are treated as " + "debug/fallback paths.", + self.kv_transfer_config.kv_connector, + ) + self._normalize_piecewise_cudagraph_mode( breakable_cudagraph_enabled=breakable_cudagraph_enabled ) @@ -1937,6 +1978,16 @@ class VllmConfig: # Default to enable HMA if not explicitly disabled by user or logic above. self.scheduler_config.disable_hybrid_kv_cache_manager = False + if ( + self.attention_config.hisparse_config is not None + and self.scheduler_config.disable_hybrid_kv_cache_manager + ): + raise ValueError( + "HiSparse requires the hybrid KV cache manager; remove " + "--disable-hybrid-kv-cache-manager or use connectors that " + "support HMA." + ) + if self.compilation_config.debug_dump_path: self.compilation_config.debug_dump_path = ( self.compilation_config.debug_dump_path.absolute().expanduser() diff --git a/vllm/distributed/kv_transfer/kv_connector/factory.py b/vllm/distributed/kv_transfer/kv_connector/factory.py index aad7999d08..c08064bd84 100644 --- a/vllm/distributed/kv_transfer/kv_connector/factory.py +++ b/vllm/distributed/kv_transfer/kv_connector/factory.py @@ -197,6 +197,12 @@ KVConnectorFactory.register_connector( "MultiConnector", ) +KVConnectorFactory.register_connector( + "HiSparseConnector", + "vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector", + "HiSparseConnector", +) + KVConnectorFactory.register_connector( "MoRIIOConnector", "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector", diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index 39f8d8a1b9..93552d51b5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -60,6 +60,7 @@ class KVOutputAggregator: # [req_id -> n_remaining_workers] self._recv_remaining_count = dict[str, int]() self._send_remaining_count = dict[str, int]() + self._failed_recving_pending = set[str]() self._expected_finished_count = expected_finished_count @classmethod @@ -156,6 +157,10 @@ class KVOutputAggregator: combined_kv_cache_events.increment_workers(1) invalid_block_ids |= kv_output.invalid_block_ids + self._failed_recving_pending |= kv_output.failed_recving + + failed_recving = self._failed_recving_pending & finished_recving + self._failed_recving_pending -= failed_recving # select output of the worker specified by output_rank output = outputs[output_rank] @@ -168,6 +173,7 @@ class KVOutputAggregator: kv_cache_events=combined_kv_cache_events or None, kv_connector_worker_meta=aggregated_kv_connector_worker_meta or None, invalid_block_ids=invalid_block_ids, + failed_recving=failed_recving, expected_finished_count=self._expected_finished_count, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index fa764571ac..c97374832f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -34,8 +34,8 @@ The class provides the following primitives: save_kv_layer() - starts saving KV for layer i (maybe async) wait_for_save() - blocks until all saves are done - get_finished() - called with ids of finished requests, returns - ids of requests that have completed async sending/recving. + get_transfer_results() - returns async send/receive completions and + receive failures in one snapshot. build_connector_worker_meta() - builds metadata to be sent back to the scheduler-side connector """ @@ -43,6 +43,7 @@ The class provides the following primitives: import enum from abc import ABC, abstractmethod from collections.abc import Callable, Iterable +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal import torch @@ -63,7 +64,7 @@ if TYPE_CHECKING: ) from vllm.forward_context import ForwardContext from vllm.v1.core.block_pool import BlockPool - from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request @@ -82,6 +83,19 @@ CopyBlocksOp = Callable[ logger = init_logger(__name__) +@dataclass +class KVConnectorTransferResults: + """Asynchronous transfer completions from one worker snapshot. + + Failed receives also appear in ``finished_recving`` so the scheduler can + release the request from its transfer wait state. + """ + + finished_sending: set[str] = field(default_factory=set) + finished_recving: set[str] = field(default_factory=set) + failed_recving: set[str] = field(default_factory=set) + + class SupportsHMA(ABC): """ The class that indicates the corresponding connector supports hybrid memory @@ -211,6 +225,7 @@ class KVConnectorBase_V1(ABC): else: raise ValueError("kv_transfer_config must be set for KVConnectorBase_V1") self._kv_cache_config = kv_cache_config + self._kv_cache_manager: KVCacheManager | None = None self._role = role @property @@ -278,6 +293,15 @@ class KVConnectorBase_V1(ABC): """ return + # TODO(NickLucche): group model-runner lifecycle hooks in the interface. + def finish_forward(self) -> None: + """Notify the connector that the model no longer reads this step's KV.""" + return + + def reset_capture_state(self) -> None: + """Reset worker state mutated while capturing CUDA graphs.""" + return + def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata): """ Handle preempted requests or evicted blocks BEFORE they are overwritten. @@ -368,6 +392,16 @@ class KVConnectorBase_V1(ABC): """ return None, None + def get_transfer_results( + self, finished_req_ids: set[str] + ) -> KVConnectorTransferResults: + """Return completed sends, receives, and receive failures together.""" + finished_sending, finished_recving = self.get_finished(finished_req_ids) + return KVConnectorTransferResults( + finished_sending=set(finished_sending or ()), + finished_recving=set(finished_recving or ()), + ) + def get_block_ids_with_load_errors(self) -> set[int]: """ Get the set of block IDs that failed to load. @@ -380,9 +414,9 @@ class KVConnectorBase_V1(ABC): - Applies to both sync- and async-loading requests. - Async loading: failed blocks may be reported in any forward pass up to and including the pass where the request ID is returned by - `get_finished()`. Even if failures occur, the request must still - be reported via `get_finished()`, and the failed block IDs must - appear here no later than that same pass. + `get_transfer_results()`. Even if failures occur, the request + must still be reported as finished receiving, and the failed + block IDs must appear here no later than that same pass. - Sync loading: failed blocks should be reported in the forward pass in which they are detected. """ @@ -436,6 +470,11 @@ class KVConnectorBase_V1(ABC): # Scheduler-side methods # ============================== + def bind_kv_cache_manager(self, kv_cache_manager: "KVCacheManager") -> None: + """Bind the scheduler's cache manager after it has been constructed.""" + self._kv_cache_manager = kv_cache_manager + self.bind_gpu_block_pool(kv_cache_manager.block_pool) + def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None: """ Bind the GPU block pool to the connector for per-GPU block status tracking. @@ -583,6 +622,10 @@ class KVConnectorBase_V1(ABC): """ return () + def has_pending_block_frees(self) -> bool: + """Whether pending transfers can release blocks instead of preemption.""" + return False + def has_pending_push_work(self) -> bool: """Return True if the connector has push-mode work that requires the engine main loop to keep stepping (e.g. a P-side request whose diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py new file mode 100644 index 0000000000..338fd5cc6e --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py @@ -0,0 +1,329 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Connector boundary for HiSparse scheduler and worker state.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorBase_V1, + KVConnectorMetadata, + KVConnectorRole, + KVConnectorWorkerMetadata, + SupportsHMA, +) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( + HiSparseConnectorWorker, +) +from vllm.v1.attention.backend import AttentionMetadata +from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.hisparse.coordinator import get_hisparse_coordinator +from vllm.v1.hisparse.types import SparseKVOffloadCommand, SparseKVRowMirror +from vllm.v1.outputs import KVConnectorOutput + +if TYPE_CHECKING: + from vllm.forward_context import ForwardContext + from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager + from vllm.v1.hisparse.coordinator import HiSparseCoordinator + from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.request import Request + + +@dataclass +class HiSparseConnectorMetadata(KVConnectorMetadata): + command: SparseKVOffloadCommand | None + host_block_copies: tuple[KVCacheBlockCopy, ...] + source_block_ids: tuple[int, ...] + row_mirrors: dict[str, tuple[SparseKVRowMirror, ...]] + all_context_pages_resident: bool + + +@dataclass +class HiSparseConnectorWorkerMetadata(KVConnectorWorkerMetadata): + enqueued_transfer_counts: dict[int, int] + completed_transfer_counts: dict[int, int] + completed_host_copy_dst_ids: tuple[int, ...] = () + + def aggregate(self, other: KVConnectorWorkerMetadata) -> KVConnectorWorkerMetadata: + assert isinstance(other, HiSparseConnectorWorkerMetadata) + + def add_counts(first: dict[int, int], second: dict[int, int]) -> dict[int, int]: + combined = first.copy() + for transfer_id, count in second.items(): + combined[transfer_id] = combined.get(transfer_id, 0) + count + return combined + + return HiSparseConnectorWorkerMetadata( + enqueued_transfer_counts=add_counts( + self.enqueued_transfer_counts, other.enqueued_transfer_counts + ), + completed_transfer_counts=add_counts( + self.completed_transfer_counts, other.completed_transfer_counts + ), + # Every worker runs the same copies, so a union is the ack. + completed_host_copy_dst_ids=tuple( + {*self.completed_host_copy_dst_ids, *other.completed_host_copy_dst_ids} + ), + ) + + +class HiSparseConnectorScheduler: + def __init__( + self, + *, + async_speculative: bool, + draft_kv_lookahead: int = 0, + ) -> None: + self.coordinator: HiSparseCoordinator | None = None + self.async_speculative = async_speculative + self.draft_kv_lookahead = draft_kv_lookahead + self.requests: dict[str, Request] = {} + + def bind_coordinator(self, coordinator: HiSparseCoordinator) -> None: + assert self.coordinator is None + self.coordinator = coordinator + + def build_connector_meta( + self, scheduler_output: SchedulerOutput + ) -> HiSparseConnectorMetadata: + assert self.coordinator is not None + # HiSparse rebinds worker state every step, so the load must start + # before the forward rather than being deferred to post-forward. + scheduler_output.has_sync_kv_loads = True + scheduler_output.block_table_updates = ( + self.coordinator.take_block_table_updates() or None + ) + command = self.coordinator.build_offload_command() + host_block_copies = self.coordinator.take_host_block_copies() + source_group_id = self.coordinator.host_group_id + assert source_group_id is not None + source_block_ids = [ + block_id + for request in scheduler_output.scheduled_new_reqs + for block_id in request.block_ids[source_group_id] + ] + for new_block_ids in scheduler_output.scheduled_cached_reqs.new_block_ids: + if new_block_ids is not None: + source_block_ids.extend(new_block_ids[source_group_id]) + num_computed_tokens = { + request.req_id: request.num_computed_tokens + for request in scheduler_output.scheduled_new_reqs + } + num_computed_tokens.update( + zip( + scheduler_output.scheduled_cached_reqs.req_ids, + scheduler_output.scheduled_cached_reqs.num_computed_tokens, + ) + ) + scheduled_requests = tuple( + ( + request_id, + num_computed_tokens[request_id], + scheduled_count, + ) + for request_id, scheduled_count in ( + scheduler_output.num_scheduled_tokens.items() + ) + ) + row_mirrors = {} + for request_id, scheduled_start, scheduled_count in scheduled_requests: + mirror_start = scheduled_start + if self.async_speculative: + mirror_start = max( + 0, + scheduled_start - self.requests[request_id].num_output_placeholders, + ) + row_mirrors[request_id] = self.coordinator.build_row_mirrors( + ( + ( + request_id, + mirror_start, + scheduled_count + + scheduled_start + - mirror_start + + self.draft_kv_lookahead, + ), + ), + ) + return HiSparseConnectorMetadata( + command, + host_block_copies, + tuple(source_block_ids), + row_mirrors, + self.coordinator.all_context_pages_resident(scheduled_requests), + ) + + def update_connector_output(self, connector_output: KVConnectorOutput) -> None: + assert self.coordinator is not None + for request_id in connector_output.finished_recving or (): + self.coordinator.finish_host_import( + request_id, failed=request_id in connector_output.failed_recving + ) + metadata = connector_output.kv_connector_worker_meta + if metadata is None: + return + assert isinstance(metadata, HiSparseConnectorWorkerMetadata) + self.coordinator.release_completed_host_copies( + metadata.completed_host_copy_dst_ids + ) + self.coordinator.update_spills( + metadata.enqueued_transfer_counts, + metadata.completed_transfer_counts, + ) + + +class HiSparseConnector(KVConnectorBase_V1, SupportsHMA): + """Join the scheduler coordinator to the worker's transfer engine.""" + + @classmethod + def get_required_kvcache_layout(cls, vllm_config: VllmConfig) -> str: + return "BLHNC" + + def __init__( + self, + vllm_config: VllmConfig, + role: KVConnectorRole, + kv_cache_config: KVCacheConfig, + ) -> None: + super().__init__(vllm_config, role, kv_cache_config) + if kv_cache_config.hisparse_host_num_blocks is None: + raise ValueError("HiSparseConnector requires a HiSparse host pool") + self.connector_scheduler: HiSparseConnectorScheduler | None = None + self.connector_worker: HiSparseConnectorWorker | None = None + if role == KVConnectorRole.SCHEDULER: + speculative_config = vllm_config.speculative_config + self.connector_scheduler = HiSparseConnectorScheduler( + async_speculative=bool( + vllm_config.scheduler_config.async_scheduling + and speculative_config is not None + ), + draft_kv_lookahead=vllm_config.num_lookahead_tokens, + ) + elif role == KVConnectorRole.WORKER: + self.connector_worker = HiSparseConnectorWorker( + vllm_config, kv_cache_config + ) + else: + raise ValueError(f"Unsupported KV connector role: {role}") + + def bind_kv_cache_manager(self, kv_cache_manager: KVCacheManager) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.bind_coordinator( + get_hisparse_coordinator(kv_cache_manager) + ) + + @property + def requires_kv_delivery(self) -> bool: + return False + + def has_pending_push_work(self) -> bool: + assert self.connector_scheduler is not None + assert self.connector_scheduler.coordinator is not None + return self.connector_scheduler.coordinator.has_pending_work() + + def has_pending_block_frees(self) -> bool: + assert self.connector_scheduler is not None + assert self.connector_scheduler.coordinator is not None + return self.connector_scheduler.coordinator.has_pending_reclamation() + + def finish_forward(self) -> None: + assert self.connector_worker is not None + self.connector_worker.finish_forward() + + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: + assert self.connector_worker is not None + self.connector_worker.register_kv_caches(kv_caches) + + def reset_capture_state(self) -> None: + assert self.connector_worker is not None + self.connector_worker.reset_hot_state() + + def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None: + assert self.connector_worker is not None + metadata = self._get_connector_metadata() + assert isinstance(metadata, HiSparseConnectorMetadata) + request_state_indices = kwargs.get("request_state_indices") + assert request_state_indices is None or isinstance( + request_state_indices, torch.Tensor + ) + request_ids = kwargs.get("request_ids") + assert request_ids is None or isinstance(request_ids, list) + attn_metadata = kwargs.get("attn_metadata") + assert attn_metadata is None or isinstance(attn_metadata, dict) + self.connector_worker.start_step( + metadata, + request_state_indices, + request_ids, + num_tokens=int(kwargs.get("num_tokens") or 0), + ) + self.connector_worker.prepare_forward(attn_metadata) + + def wait_for_layer_load(self, layer_name: str) -> None: + return + + def save_kv_layer( + self, + layer_name: str, + kv_layer: torch.Tensor, + attn_metadata: AttentionMetadata, + **kwargs: Any, + ) -> None: + return + + def wait_for_save(self) -> None: + return + + def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None: + assert self.connector_worker is not None + enqueued, completed = self.connector_worker.take_transfer_updates() + host_copies = self.connector_worker.take_completed_host_copies() + if not enqueued and not completed and not host_copies: + return None + return HiSparseConnectorWorkerMetadata( + enqueued_transfer_counts={transfer_id: 1 for transfer_id in enqueued}, + completed_transfer_counts={transfer_id: 1 for transfer_id in completed}, + completed_host_copy_dst_ids=tuple(host_copies), + ) + + def shutdown(self) -> None: + if self.connector_worker is not None: + self.connector_worker.shutdown() + + def get_num_new_matched_tokens( + self, request: Request, num_computed_tokens: int + ) -> tuple[int | None, bool]: + return 0, False + + def update_state_after_alloc( + self, + request: Request, + blocks: KVCacheBlocks, + num_external_tokens: int, + ) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.requests[request.request_id] = request + + def build_connector_meta( + self, scheduler_output: SchedulerOutput + ) -> KVConnectorMetadata: + assert self.connector_scheduler is not None + return self.connector_scheduler.build_connector_meta(scheduler_output) + + def update_connector_output(self, connector_output: KVConnectorOutput) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.update_connector_output(connector_output) + + def request_finished_all_groups( + self, + request: Request, + block_ids: tuple[list[int], ...], + ) -> tuple[bool, dict[str, Any] | None]: + assert self.connector_scheduler is not None + self.connector_scheduler.requests.pop(request.request_id, None) + return False, None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py new file mode 100644 index 0000000000..1f723e5344 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py @@ -0,0 +1,832 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HiSparse worker-side host/hot state and data movement.""" + +from __future__ import annotations + +from collections import deque +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import partial +from typing import TYPE_CHECKING, Any, NamedTuple + +import numpy as np +import torch + +from vllm import _custom_ops as ops +from vllm.config import VllmConfig +from vllm.distributed.parallel_state import ( + get_tensor_model_parallel_rank, + get_tp_group, +) +from vllm.utils.torch_utils import current_stream +from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy +from vllm.v1.hisparse.layout import HISPARSE_HOT_SUFFIX +from vllm.v1.hisparse.runtime import HiSparseCacheHandle, release_pinned_state +from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror +from vllm.v1.kv_cache_interface import ( + HiSparseHotSpec, + KVCacheConfig, +) +from vllm.v1.worker.utils import copy_kv_cache_blocks_inplace + +if TYPE_CHECKING: + from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( + HiSparseConnectorMetadata, + ) + from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion + + +class _DMADescriptors(NamedTuple): + src: torch.Tensor + dst: torch.Tensor + sizes: torch.Tensor + src_np: np.ndarray + dst_np: np.ndarray + sizes_np: np.ndarray + + +@dataclass +class _SlotMappingStaging: + stream: torch.Stream + event: torch.Event + slots: torch.Tensor + candidates: tuple[SparseKVRowMirror, ...] = () + num_tokens: int = 0 + source_index: int = 0 + + +def _allocate_dma_descriptors(size: int) -> _DMADescriptors: + src, dst, sizes = (torch.empty(size, dtype=torch.int64) for _ in range(3)) + return _DMADescriptors( + src, + dst, + sizes, + src.numpy(), + dst.numpy(), + sizes.numpy(), + ) + + +def _get_hisparse_cache( + forward_context: dict[str, Any], layer_name: str +) -> HiSparseCacheHandle: + attention_layer = forward_context[layer_name] + hisparse_cache = attention_layer.hisparse_cache + assert hisparse_cache is not None + return hisparse_cache + + +def _flatten_row_mirrors( + row_mirrors: Mapping[str, tuple[SparseKVRowMirror, ...]], + request_ids: Sequence[str] | None, +) -> tuple[SparseKVRowMirror, ...]: + ordered_ids = row_mirrors if request_ids is None else request_ids + return tuple( + mirror + for request_id in ordered_ids + for mirror in row_mirrors.get(request_id, ()) + ) + + +def _select_written_row_mirrors( + candidates: tuple[SparseKVRowMirror, ...], + source_slots: np.ndarray, + source_index: int, +) -> tuple[SparseKVRowMirror, ...]: + """Select and coalesce GPU-written rows from a scheduler-owned envelope.""" + source_slots = source_slots[source_slots >= 0] + if source_slots.size == 0 or not candidates: + return () + starts = np.asarray([mirror.source_starts for mirror in candidates], dtype=np.int64) + counts = np.fromiter((mirror.num_rows for mirror in candidates), dtype=np.int64) + destinations = np.fromiter( + [mirror.destination_start for mirror in candidates], dtype=np.int64 + ) + if source_index >= starts.shape[1]: + raise RuntimeError("HiSparse row-mirror source index is out of range.") + source_starts = starts[:, source_index] + order = np.argsort(source_starts) + matches = np.searchsorted(source_starts[order], source_slots, side="right") - 1 + candidate_ids = order[np.maximum(matches, 0)] + offsets = source_slots - source_starts[candidate_ids] + valid = (matches >= 0) & (offsets < counts[candidate_ids]) + if not np.all(valid): + # A page the scheduler could not map has no host destination, so its + # rows are simply not mirrorable this step. + source_slots = source_slots[valid] + candidate_ids = candidate_ids[valid] + offsets = offsets[valid] + if source_slots.size == 0: + return () + sources = starts[candidate_ids] + offsets[:, None] + destinations = destinations[candidate_ids] + offsets + boundaries = ( + np.flatnonzero( + (candidate_ids[1:] != candidate_ids[:-1]) + | (np.diff(destinations) != 1) + | np.any(np.diff(sources, axis=0) != 1, axis=1) + ) + + 1 + ) + run_starts = np.concatenate(([0], boundaries)) + run_ends = np.concatenate((boundaries, [destinations.size])) + return tuple( + SparseKVRowMirror( + source_starts=tuple(int(value) for value in sources[start]), + destination_start=int(destinations[start]), + num_rows=int(end - start), + ) + for start, end in zip(run_starts, run_ends, strict=True) + ) + + +def _create_hisparse_host_events( + shared_host_region: SharedOffloadRegion | None, + is_host_writer: bool, + device: torch.device, +) -> tuple[torch.Event, torch.Event]: + if shared_host_region is None: + return torch.Event(), torch.Event() + + events: tuple[torch.Event, torch.Event] | None = None + ipc_handles = None + if is_host_writer: + events = ( + torch.cuda.Event(interprocess=True), + torch.cuda.Event(interprocess=True), + ) + stream = current_stream() + for event in events: + event.record(stream) + ipc_handles = tuple(event.ipc_handle() for event in events) + ipc_handles = get_tp_group().broadcast_object(ipc_handles, src=0) + if events is None: + events = ( + torch.cuda.Event.from_ipc_handle(device, ipc_handles[0]), + torch.cuda.Event.from_ipc_handle(device, ipc_handles[1]), + ) + return events + + +class HiSparseConnectorWorker: + """Own HiSparse host/hot state and execute its worker-side transfers.""" + + def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> None: + self.vllm_config = vllm_config + self.kv_cache_config = kv_cache_config + self._initialized = False + + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: + forward_context = self.vllm_config.compilation_config.static_forward_context + cache_handles: list[HiSparseCacheHandle] = [] + cache_layer_names: list[str] = [] + for group in self.kv_cache_config.kv_cache_groups: + if not isinstance(group.kv_cache_spec, HiSparseHotSpec): + continue + for cache_name in group.layer_names: + assert cache_name.endswith(HISPARSE_HOT_SUFFIX) + layer_name = cache_name[: -len(HISPARSE_HOT_SUFFIX)] + cache_handles.append(_get_hisparse_cache(forward_context, layer_name)) + cache_layer_names.append(layer_name) + + if not cache_handles: + raise RuntimeError("HiSparse connector found no hot-cache handles.") + hot_backings: dict[int, torch.Tensor] = {} + registered_host_pools: dict[int, torch.Tensor] = {} + for cache in cache_handles: + hot_backing = cache.runtime.hot_backing + hot_backings[hot_backing.untyped_storage().data_ptr()] = hot_backing + registered_pool = cache.runtime.registered_host_pool + registered_host_pools[registered_pool.data_ptr()] = registered_pool + if len(hot_backings) != 1: + raise RuntimeError("HiSparse hot tensors must share one GPU backing.") + hot_backing = next(iter(hot_backings.values())) + pinned_host_pools = list(registered_host_pools.values()) + + resident = cache_handles[0].view + assert resident is not None + host_num_blocks = self.kv_cache_config.hisparse_host_num_blocks + assert host_num_blocks is not None + try: + self.initialize( + cache_handles, + cache_layer_names, + hot_backing, + self.vllm_config.scheduler_config.max_num_seqs, + host_num_blocks, + hot_backing.device, + pinned_host_pools, + ) + except Exception: + release_pinned_state( + [cache.runtime for cache in cache_handles], + pinned_host_pools, + ) + raise + + def initialize( + self, + cache_handles: list[HiSparseCacheHandle], + cache_layer_names: list[str], + hot_backing: torch.Tensor, + max_num_reqs: int, + host_num_blocks: int, + device: torch.device, + pinned_host_pools: list[torch.Tensor], + *, + shared_host_region: SharedOffloadRegion | None = None, + is_host_writer: bool = True, + ) -> None: + if self._initialized: + raise RuntimeError("HiSparse connector worker is already initialized.") + resident = cache_handles[0].view + assert resident is not None + self.kernel_block_size = resident.block_size + self.host_num_blocks = host_num_blocks + self.pinned_host_pools = pinned_host_pools + self.shared_host_region = shared_host_region + self.is_host_writer = is_host_writer + self.dma_stream = ( + torch.cuda.Stream(device=device) if self.is_host_writer else None + ) + self._slot_mapping_staging = None + if self.is_host_writer: + max_mirror_rows = ( + self.vllm_config.scheduler_config.max_num_batched_tokens + + max_num_reqs * (self.vllm_config.num_lookahead_tokens + 1) + ) + self._slot_mapping_staging = _SlotMappingStaging( + stream=torch.Stream(device=device), + event=torch.Event(), + slots=torch.empty( + max_mirror_rows, + dtype=torch.int64, + pin_memory=True, + ), + ) + self._dma_free_descriptors: list[_DMADescriptors] = [] + self._pending_dma_descriptors: deque[tuple[torch.Event, _DMADescriptors]] = ( + deque() + ) + self._dma_submitted = False + self._per_layer_mirrored: set[int] = set() + self._submitted_mirror_layers: set[int] = set() + self._layer_ready_events = tuple(torch.Event() for _ in cache_handles) + self._forward_ready_event = torch.Event() + self._set_row_mirrors(()) + self.cache_layer_names = cache_layer_names + self._group_leaders = tuple( + (layer_name, cache) + for layer_name, cache in zip(cache_layer_names, cache_handles, strict=True) + if cache.runtime.is_group_leader + ) + self.host_write_events = _create_hisparse_host_events( + shared_host_region, is_host_writer, device + ) + self.host_write_event = self.host_write_events[1] + self._next_host_write_event = 0 + self.cache_handles = cache_handles + self.leader_runtimes = [cache.runtime for _, cache in self._group_leaders] + request_state_indices = { + indices.data_ptr(): indices + for cache in cache_handles + if (indices := cache.runtime.request_state_indices) is not None + } + if len(request_state_indices) != 1: + raise RuntimeError( + "HiSparse runtimes must share one request-state mapping." + ) + self.request_state_indices = next(iter(request_state_indices.values())) + if self.request_state_indices.numel() != max_num_reqs: + raise RuntimeError( + "HiSparse request-state mapping does not match max_num_seqs." + ) + self.hot_backing = hot_backing + self._pending_invalid_block_ids: list[int] = [] + # Destination block ids of host copies this worker has run. + self._completed_host_copy_dst_ids: list[int] = [] + self._post_forward_transfers: list[SparseKVPageTransfer] = [] + self._enqueued_transfer_ids: list[int] = [] + self._pending_transfer_events: deque[tuple[torch.Event, tuple[int, ...]]] = ( + deque() + ) + self._init_dma() + if self.is_host_writer: + for layer_index, handle in enumerate(cache_handles): + handle.submit_layer_mirror = partial( + self._enqueue_layer_mirror, layer_index + ) + self._initialized = True + + def set_request_state_indices(self, indices: torch.Tensor) -> None: + if indices.numel() > self.request_state_indices.numel(): + raise ValueError( + "HiSparse request-state mapping exceeds max_num_seqs: " + f"{indices.numel()} > {self.request_state_indices.numel()}." + ) + if torch.cuda.is_current_stream_capturing(): + return + # Attention indexes persistent request state by input-batch row. Refresh + # that indirection after every batch compaction or reorder. + self.request_state_indices.fill_(-1) + self.request_state_indices[: indices.numel()].copy_(indices) + if self._pending_invalid_block_ids: + self.invalidate_blocks(self._pending_invalid_block_ids, indices) + self._pending_invalid_block_ids.clear() + + def _init_dma(self) -> None: + host_caches = tuple(cache.runtime.host_cache for cache in self.cache_handles) + resident_caches = [] + for cache in self.cache_handles: + assert cache.view is not None and cache.slot_mapping is not None + resident_caches.append(cache.view.cache) + self.host_caches = host_caches + self.resident_caches = tuple(resident_caches) + + for resident_cache, host_cache in zip(resident_caches, host_caches): + row_bytes = resident_cache.shape[-1] * resident_cache.element_size() + if ( + resident_cache.ndim != 3 + or resident_cache.shape[1] != self.kernel_block_size + or resident_cache.stride(1) * resident_cache.element_size() != row_bytes + ): + raise RuntimeError("HiSparse DMA requires contiguous resident rows.") + if ( + host_cache.ndim != 2 + or not host_cache.is_contiguous() + or host_cache.shape[1] * host_cache.element_size() != row_bytes + ): + raise RuntimeError("HiSparse DMA requires contiguous host rows.") + + def start_step( + self, + metadata: HiSparseConnectorMetadata, + request_state_indices: torch.Tensor | None, + request_ids: list[str] | None = None, + num_tokens: int = 0, + ) -> None: + self._stage_row_mirror_mapping(num_tokens) + previous_host_write_event = self.host_write_event + self.host_write_event = self.host_write_events[self._next_host_write_event] + self._next_host_write_event ^= 1 + current_stream().wait_event(previous_host_write_event) + self._release_completed_dma_descriptors() + mirrors = _flatten_row_mirrors(metadata.row_mirrors, request_ids) + if self._slot_mapping_staging is not None: + self._slot_mapping_staging.candidates = mirrors + self._set_row_mirrors(mirrors) + self._dma_submitted = False + self._clear_forward_mirror_state() + for handle in self.cache_handles: + handle.all_context_pages_resident = metadata.all_context_pages_resident + handle.mirror_from_resident = True + self._copy_host_blocks(metadata.host_block_copies, previous_host_write_event) + transfers = ( + metadata.command.page_transfers if metadata.command is not None else [] + ) + self._post_forward_transfers = [ + transfer for transfer in transfers if transfer.after_forward + ] + self._submit_transfers( + [transfer for transfer in transfers if not transfer.after_forward] + ) + self._pending_invalid_block_ids.extend(metadata.source_block_ids) + if request_state_indices is not None: + self.set_request_state_indices(request_state_indices) + + def _clear_forward_mirror_state(self) -> None: + self._per_layer_mirrored.clear() + self._submitted_mirror_layers.clear() + for handle in self.cache_handles: + handle.decode_batch = False + handle.host_mirror_required = False + handle.num_actual_tokens = 0 + handle.num_decode_tokens = 0 + handle.req_id_per_token = None + + def prepare_forward(self, attn_metadata: Mapping[str, Any] | None) -> None: + if attn_metadata is None: + return + for layer_name, handle in self._group_leaders: + metadata = attn_metadata.get(layer_name) + if metadata is not None: + handle.prepare_group_for_batch(metadata) + + def _stage_row_mirror_mapping(self, num_tokens: int) -> None: + """Snapshot the rows this forward will write, off the compute stream. + + The resident slot mapping is a persistent view bound at registration + time, so the rows are already staged by the time the forward launches. + """ + state = self._slot_mapping_staging + if state is None or not num_tokens: + return + handle = self.cache_handles[0] + slots = handle.slot_mapping + assert slots is not None + source_index = handle.runtime.resident_source_index + start = state.num_tokens + end = start + num_tokens + if end > state.slots.shape[0]: + raise ValueError( + "HiSparse row mapping exceeds staging capacity: " + f"{end} > {state.slots.shape[0]}." + ) + main_stream = current_stream() + state.stream.wait_stream(main_stream) + with torch.cuda.stream(state.stream): + state.slots[start:end].copy_( + slots[:num_tokens], + non_blocking=True, + ) + state.event.record(state.stream) + state.num_tokens = end + state.source_index = source_index + + def _resolve_row_mirrors(self, state: _SlotMappingStaging) -> None: + """Compute the final row mirrors from the staged GPU slot mapping.""" + self._set_row_mirrors( + _select_written_row_mirrors( + state.candidates, + state.slots[: state.num_tokens].numpy(), + state.source_index, + ) + ) + state.num_tokens = 0 + + def _copy_host_blocks( + self, + host_block_copies: Sequence[KVCacheBlockCopy], + previous_host_write_event: torch.Event, + ) -> None: + if not host_block_copies: + return + self._completed_host_copy_dst_ids.extend( + copy.dst_block_id for copy in host_block_copies + ) + if self.shared_host_region is None or get_tensor_model_parallel_rank() == 0: + if self.host_caches: + previous_host_write_event.synchronize() + copy_kv_cache_blocks_inplace( + self.host_caches, + self.host_num_blocks, + host_block_copies, + ) + if self.shared_host_region is not None: + get_tp_group().barrier() + + def invalidate_blocks( + self, block_ids: list[int], request_state_indices: torch.Tensor + ) -> None: + """Invalidate recycled host slots in this worker's leader runtimes.""" + if not block_ids: + return + device = self.cache_handles[0].runtime.device + staging = torch.tensor(block_ids, dtype=torch.int32, pin_memory=True) + blocks = staging.to(device, dtype=torch.int32, non_blocking=True) + offsets = torch.arange(self.kernel_block_size, dtype=torch.int32, device=device) + slots = (blocks[:, None] * self.kernel_block_size + offsets[None, :]).flatten() + sorted_slots = torch.sort(slots).values + state_indices = request_state_indices.to(device=device, dtype=torch.long) + for runtime in self.leader_runtimes: + runtime.invalidate_sorted_slots(sorted_slots, state_indices) + + def reset_hot_state(self) -> None: + for runtime in self.leader_runtimes: + runtime.reset_hot_state() + + def _release_completed_dma_descriptors(self) -> None: + pending = self._pending_dma_descriptors + while pending and pending[0][0].query(): + _, descriptors = pending.popleft() + self._dma_free_descriptors.append(descriptors) + + def _acquire_dma_descriptors(self, size: int) -> _DMADescriptors: + for index, descriptors in enumerate(self._dma_free_descriptors): + if descriptors.src.numel() >= size: + return self._dma_free_descriptors.pop(index) + return _allocate_dma_descriptors(size) + + def _submit_dma_descriptors( + self, + descriptors: _DMADescriptors, + descriptor_count: int, + transfer_ids: tuple[int, ...] = (), + ready_event: torch.Event | None = None, + ) -> None: + stream = self.dma_stream + assert stream is not None + if ready_event is None: + stream.wait_stream(current_stream()) + else: + stream.wait_event(ready_event) + completion_event = torch.Event() + with torch.cuda.stream(stream): + ops.swap_blocks_batch( + descriptors.src[:descriptor_count], + descriptors.dst[:descriptor_count], + descriptors.sizes[:descriptor_count], + ) + self.host_write_event.record(stream) + completion_event.record(stream) + self._pending_dma_descriptors.append((completion_event, descriptors)) + self._dma_submitted = True + if transfer_ids: + self._pending_transfer_events.append((completion_event, transfer_ids)) + self._enqueued_transfer_ids.extend(transfer_ids) + + def _set_row_mirrors(self, mirrors: tuple[SparseKVRowMirror, ...]) -> None: + self._row_mirrors = mirrors + self._row_mirror_destination_starts = np.fromiter( + (mirror.destination_start for mirror in mirrors), + dtype=np.int64, + count=len(mirrors), + ) + self._row_mirror_counts = np.fromiter( + (mirror.num_rows for mirror in mirrors), + dtype=np.int64, + count=len(mirrors), + ) + if mirrors: + self._row_mirror_source_starts = np.asarray( + [mirror.source_starts for mirror in mirrors], dtype=np.int64 + ) + if self._row_mirror_source_starts.ndim != 2: + raise RuntimeError("HiSparse DMA source mappings must be rectangular.") + else: + self._row_mirror_source_starts = np.empty((0, 0), dtype=np.int64) + self._row_mirror_num_rows = int(self._row_mirror_counts.sum()) + + def _enqueue_row_dma( + self, layer_indices: Sequence[int], ready_event: torch.Event | None = None + ) -> None: + if ( + not layer_indices + or not self._row_mirror_num_rows + or not self.is_host_writer + ): + return + mirrors = self._row_mirrors + num_layers = len(layer_indices) + descriptor_count = len(mirrors) * num_layers + descriptors = self._acquire_dma_descriptors(descriptor_count) + destination_starts = self._row_mirror_destination_starts + row_counts = self._row_mirror_counts + for descriptor_offset, layer_index in enumerate(layer_indices): + cache = self.cache_handles[layer_index] + source_index = cache.runtime.resident_source_index + if source_index >= self._row_mirror_source_starts.shape[1]: + raise RuntimeError("HiSparse row DMA source index is out of range.") + source_rows = self._row_mirror_source_starts[:, source_index] + source = self.resident_caches[layer_index] + destination = self.host_caches[layer_index] + row_bytes = source.shape[-1] * source.element_size() + if ( + source.stride(1) * source.element_size() != row_bytes + or destination.shape[1] * destination.element_size() != row_bytes + ): + raise RuntimeError("HiSparse row DMA requires contiguous rows.") + source_blocks, source_row_offsets = np.divmod( + source_rows, self.kernel_block_size + ) + source_out_of_range = ( + np.any(source_blocks < 0) + or np.any(source_blocks >= source.shape[0]) + or np.any(source_row_offsets + row_counts > self.kernel_block_size) + ) + if source_out_of_range: + raise RuntimeError("HiSparse row DMA source is out of range.") + if np.any(destination_starts < 0) or np.any( + destination_starts + row_counts > destination.shape[0] + ): + raise RuntimeError("HiSparse row DMA index is out of range.") + descriptor_slice = slice(descriptor_offset, descriptor_count, num_layers) + descriptors.src_np[descriptor_slice] = ( + source.data_ptr() + + source_blocks * source.stride(0) * source.element_size() + + source_row_offsets * row_bytes + ) + descriptors.dst_np[descriptor_slice] = ( + destination.data_ptr() + destination_starts * row_bytes + ) + descriptors.sizes_np[descriptor_slice] = row_counts * row_bytes + + self._submit_dma_descriptors( + descriptors, descriptor_count, ready_event=ready_event + ) + + def _enqueue_layer_mirror(self, layer_index: int) -> None: + handle = self.cache_handles[layer_index] + if not handle.host_mirror_required: + return + if layer_index in self._per_layer_mirrored: + raise RuntimeError(f"HiSparse layer {layer_index} mirrored twice.") + state = self._slot_mapping_staging + if state is not None and state.num_tokens: + # The staging stream waited on the compute stream before the + # forward launched, so this only drains pre-forward work that is + # already queued ahead of the layer kernels: no GPU bubble. + state.event.synchronize() + self._resolve_row_mirrors(state) + self._per_layer_mirrored.add(layer_index) + next_layer = layer_index + 1 + if ( + next_layer < len(self.cache_handles) + and self.cache_handles[next_layer].runtime.resident_source_index + == handle.runtime.resident_source_index + ): + return + ready_event = self._layer_ready_events[layer_index] + ready_event.record() + pending_layers = tuple( + sorted(self._per_layer_mirrored - self._submitted_mirror_layers) + ) + self._enqueue_row_dma( + pending_layers, + ready_event=ready_event, + ) + self._submitted_mirror_layers.update(pending_layers) + + def _record_transfer_completion( + self, transfers: list[SparseKVPageTransfer] + ) -> None: + if not transfers or not self.is_host_writer: + return + stream = self.dma_stream + assert stream is not None + completion_event = torch.Event() + completion_event.record(stream) + transfer_ids = tuple(transfer.transfer_id for transfer in transfers) + self._pending_transfer_events.append((completion_event, transfer_ids)) + self._enqueued_transfer_ids.extend(transfer_ids) + + def _submit_transfers(self, transfers: list[SparseKVPageTransfer]) -> None: + if self.cache_handles[0].runtime.eager_host_mirror: + self._record_transfer_completion(transfers) + else: + self._enqueue_transfers(transfers) + + def _enqueue_transfers(self, transfers: list[SparseKVPageTransfer]) -> None: + if not transfers or not self.is_host_writer: + return + num_layers = len(self.cache_handles) + descriptor_count = len(transfers) * num_layers + descriptors = self._acquire_dma_descriptors(descriptor_count) + destination_rows = np.fromiter( + ( + transfer.destination_block_id * self.kernel_block_size + for transfer in transfers + ), + dtype=np.int64, + count=len(transfers), + ) + source_blocks_by_transfer = np.asarray( + [transfer.source_block_ids for transfer in transfers], dtype=np.int64 + ) + if source_blocks_by_transfer.ndim != 2: + raise RuntimeError( + "HiSparse spill DMA source mappings must be rectangular." + ) + for layer_index, cache in enumerate(self.cache_handles): + source_index = cache.runtime.resident_source_index + if source_index >= source_blocks_by_transfer.shape[1]: + raise RuntimeError("HiSparse spill DMA source index is out of range.") + source_blocks = source_blocks_by_transfer[:, source_index] + source = self.resident_caches[layer_index] + destination = self.host_caches[layer_index] + if np.any(source_blocks < 0) or np.any(source_blocks >= source.shape[0]): + raise RuntimeError("HiSparse spill DMA source is out of range.") + if np.any(destination_rows < 0) or np.any( + destination_rows + self.kernel_block_size > destination.shape[0] + ): + raise RuntimeError("HiSparse spill DMA destination is out of range.") + row_bytes = source.shape[-1] * source.element_size() + descriptor_slice = slice(layer_index, descriptor_count, num_layers) + descriptors.src_np[descriptor_slice] = ( + source.data_ptr() + + source_blocks * source.stride(0) * source.element_size() + ) + descriptors.dst_np[descriptor_slice] = ( + destination.data_ptr() + destination_rows * row_bytes + ) + descriptors.sizes_np[descriptor_slice] = self.kernel_block_size * row_bytes + self._submit_dma_descriptors( + descriptors, + descriptor_count, + transfer_ids=tuple(transfer.transfer_id for transfer in transfers), + ) + + def _enqueue_host_mirror( + self, + ready_event: torch.Event | None = None, + ) -> None: + active = [ + (index, handle) + for index, handle in enumerate(self.cache_handles) + if handle.num_actual_tokens != 0 + ] + if not active: + return + + def mirror_key(handle: HiSparseCacheHandle) -> tuple: + slots = handle.mirror_slot_mapping + return ( + handle.num_actual_tokens, + handle.num_decode_tokens, + handle.decode_batch, + handle.host_mirror_required, + handle.runtime.eager_host_mirror, + None if slots is None else slots.data_ptr(), + ) + + keys = {mirror_key(handle) for _, handle in active} + if len(keys) > 1: + raise RuntimeError( + f"HiSparse cache layers disagree on mirror metadata: {list(keys)}." + ) + cache = active[0][1] + if not cache.host_mirror_required: + return + dst_slots = cache.mirror_slot_mapping + if dst_slots is None: + raise RuntimeError("HiSparse host mirror has no source slot mapping.") + num_rows = min(cache.num_actual_tokens, dst_slots.numel()) + if num_rows == 0: + return + if self.is_host_writer: + expected_layers = {index for index, _ in active} + if self._per_layer_mirrored and self._per_layer_mirrored != expected_layers: + raise RuntimeError( + "HiSparse per-layer DMA did not mirror every active layer: " + f"expected {sorted(expected_layers)}, got " + f"{sorted(self._per_layer_mirrored)}." + ) + pending_layers = tuple( + sorted(expected_layers - self._submitted_mirror_layers) + ) + if pending_layers: + self._enqueue_row_dma(pending_layers, ready_event=ready_event) + self._submitted_mirror_layers.update(pending_layers) + assert cache.req_id_per_token is not None + for _, handle in active: + if handle.runtime.is_group_leader: + handle.runtime.invalidate_written_slots( + dst_slots[:num_rows], + cache.req_id_per_token[:num_rows], + ) + + def _finish_mirror_phase(self, ready_event: torch.Event | None = None) -> None: + state = self._slot_mapping_staging + if state is not None and state.num_tokens: + # Measured free: blocking here costs no throughput (26.7 vs 27.6 + # gen tok/s) and lets every step mirror exactly the written rows. + state.event.synchronize() + self._resolve_row_mirrors(state) + self._enqueue_host_mirror(ready_event) + self._clear_forward_mirror_state() + + def finish_forward(self) -> None: + compute_stream = current_stream() + self._forward_ready_event.record() + self._finish_mirror_phase(self._forward_ready_event) + transfers = self._post_forward_transfers + self._post_forward_transfers = [] + self._submit_transfers(transfers) + if self.is_host_writer: + if self._dma_submitted: + compute_stream.wait_event(self.host_write_event) + self._dma_submitted = False + else: + self.host_write_event.record(compute_stream) + self._release_completed_dma_descriptors() + + def take_completed_host_copies(self) -> list[int]: + """Drain host copies this worker has enqueued for this step.""" + completed = self._completed_host_copy_dst_ids + self._completed_host_copy_dst_ids = [] + return completed + + def take_transfer_updates(self) -> tuple[list[int], list[int]]: + enqueued = self._enqueued_transfer_ids + self._enqueued_transfer_ids = [] + completed: list[int] = [] + while self._pending_transfer_events: + event, transfer_ids = self._pending_transfer_events[0] + if not event.query(): + break + self._pending_transfer_events.popleft() + completed.extend(transfer_ids) + return enqueued, completed + + def shutdown(self) -> None: + if not self._initialized: + return + if self._slot_mapping_staging is not None: + self._slot_mapping_staging.stream.synchronize() + if self.dma_stream is not None: + self.dma_stream.synchronize() + release_pinned_state( + [cache.runtime for cache in self.cache_handles], self.pinned_host_pools + ) + self._initialized = False diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py index f465fcad92..e28076c89b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py @@ -702,9 +702,9 @@ class LMCacheMPConnectorUpstream(KVConnectorBase_V1): - Applies to both sync- and async-loading requests. - Async loading: failed blocks may be reported in any forward pass up to and including the pass where the request ID is returned by - `get_finished()`. Even if failures occur, the request must still - be reported via `get_finished()`, and the failed block IDs must - appear here no later than that same pass. + `get_transfer_results()`. Even if failures occur, the request + must still be reported as finished receiving, and the failed + block IDs must appear here no later than that same pass. - Sync loading: failed blocks should be reported in the forward pass in which they are detected. """ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 8fa9a8836a..4009af7009 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -635,6 +635,7 @@ class MooncakeConnectorScheduler: ): self.vllm_config = vllm_config self.block_size = vllm_config.cache_config.block_size + self.kv_cache_config = kv_cache_config assert vllm_config.kv_transfer_config self.is_kv_producer: bool = ( @@ -649,7 +650,7 @@ class MooncakeConnectorScheduler: not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager and any( not isinstance(g.kv_cache_spec, FullAttentionSpec) - for g in kv_cache_config.kv_cache_groups + for g in kv_cache_config.transfer_groups ) ) # GDN is represented as a MambaSpec in vLLM. This Mooncake MambaSpec @@ -670,7 +671,7 @@ class MooncakeConnectorScheduler: (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size) if isinstance(g.kv_cache_spec, SlidingWindowSpec) else (0, self.block_size) - for g in kv_cache_config.kv_cache_groups + for g in kv_cache_config.transfer_groups ] # cdiv(n_tokens, block_size) gives blocks/window; add 1 to # conservatively account for boundary overlap. @@ -684,11 +685,17 @@ class MooncakeConnectorScheduler: block_ids: tuple[list[int], ...] | list[list[int]], ) -> list[list[int]]: """Clip per-group block IDs to sliding window size.""" - if len(block_ids) == 0 or not self._is_hma_required: - return list(block_ids) + if len(block_ids) == 0: + return [] + if len(block_ids) == len(self.kv_cache_config.transfer_group_ids): + selected = list(block_ids) + else: + selected = list(self.kv_cache_config.select_transfer_block_ids(block_ids)) + if len(selected) == 0 or not self._is_hma_required: + return selected return [ blocks[-self.blocks_per_sw[i] :] if self.blocks_per_sw[i] > 0 else blocks - for i, blocks in enumerate(block_ids) + for i, blocks in enumerate(selected) ] def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int: @@ -1046,18 +1053,13 @@ class MooncakeConnectorWorker: self._tp_size: dict[EngineId, int] = {self.engine_id: self.tp_size} self._layer_specs: dict[str, KVCacheSpec] = {} - for group in kv_cache_config.kv_cache_groups: + for group in kv_cache_config.transfer_groups: group_spec = group.kv_cache_spec specs_by_layer = getattr(group_spec, "kv_cache_specs", {}) for layer_name in group.layer_names: self._layer_specs[layer_name] = specs_by_layer.get( layer_name, group_spec ) - self._layer_group_indices: dict[str, int] = { - layer: group_index - for group_index, group in enumerate(kv_cache_config.kv_cache_groups) - for layer in group.layer_names - } self.transfer_topo = TransferTopology( tp_rank=self.tp_rank, tp_size=self.tp_size, @@ -1441,7 +1443,7 @@ class MooncakeConnectorWorker: block_arange = np.arange(self._physical_blocks_per_logical_kv_block).reshape( 1, -1 ) - group_specs = self.kv_cache_config.kv_cache_groups + group_specs = self.kv_cache_config.transfer_groups return [ BlockTable.map_to_kernel_blocks( np.array(group), @@ -1494,7 +1496,7 @@ class MooncakeConnectorWorker: local_block_ids_by_group: list[list[int]] = [] remote_block_ids_by_group: list[list[int]] = [] has_block_error = False - group_specs = self.kv_cache_config.kv_cache_groups + group_specs = self.kv_cache_config.transfer_groups for group_index, (local_group, remote_group) in enumerate( zip(send_meta.local_block_ids, remote_block_ids_per_group) ): @@ -1750,7 +1752,7 @@ class MooncakeConnectorWorker: self.registered_layer_names.append(layer_name) self.registered_layer_indices.append(layer_index) self.registered_group_indices.append( - self._layer_group_indices[layer_name] + self.kv_cache_config.transfer_group_index_by_layer[layer_name] ) storage = cache.untyped_storage() storage_addr = storage.data_ptr() @@ -2121,7 +2123,7 @@ class MooncakeConnectorWorker: ) -> list[TransferRegion]: if not group_indices: group_indices = [ - self._layer_group_indices.get(layer_name, 0) + self.kv_cache_config.transfer_group_index_by_layer.get(layer_name, 0) for layer_name in layer_names ] return _expand_transfer_regions( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index 67a935d7ee..375730807d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -3,18 +3,20 @@ """External-store cache-hit coordinator for MooncakeStoreConnector.""" from collections.abc import Sequence -from typing import cast +from typing import NamedTuple, cast from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( chunk_hashes_for_block_size, ) from vllm.utils.math_utils import cdiv from vllm.v1.core.block_pool import BlockPool -from vllm.v1.core.kv_cache_coordinator import SpecGroup from vllm.v1.core.kv_cache_utils import ( BlockHash, KVCacheBlock, ) +from vllm.v1.core.single_type_kv_cache_manager import ( + SingleTypeKVCacheManager, +) from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, @@ -25,6 +27,13 @@ from vllm.v1.kv_cache_interface import ( from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry +class StoreSpecGroup(NamedTuple): + spec: KVCacheSpec + group_ids: list[int] + manager_cls: type[SingleTypeKVCacheManager] + use_eagle: bool + + class ExternalCachedBlockPool: """Duck-typed BlockPool backed by a ``(group_id, hash)`` exists set.""" @@ -121,7 +130,7 @@ class MooncakeStoreCoordinator: """Mirrors KVCacheCoordinator.verify_and_split_kv_cache_groups but dispatches via spec_manager_map (we don't allocate managers). """ - attention_groups: list[SpecGroup] = [] + attention_groups: list[StoreSpecGroup] = [] for i, g in enumerate(self.kv_cache_groups): # Skip groups that opt out of prefix caching (e.g. GLM-5.3-Flash # kpool tail): per-request scratch, never shareable, so they must @@ -143,7 +152,7 @@ class MooncakeStoreCoordinator: break else: attention_groups.append( - SpecGroup(spec, [i], manager_cls, g.is_eagle_group) + StoreSpecGroup(spec, [i], manager_cls, g.is_eagle_group) ) # Full attention first (matches upstream convergence ordering). attention_groups.sort(key=lambda g: not isinstance(g.spec, FullAttentionSpec)) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 870500e840..ebcf5708c8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -67,6 +67,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import get_ip, make_zmq_socket +from vllm.utils.torch_utils import is_non_overlapping_and_dense from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.kv_cache_utils import ( BlockHash, @@ -1909,21 +1910,31 @@ class MooncakeStoreWorker: logger.warning("No KV caches to offload.") return + # Resolve each entry to a representative tensor for storage + # deduplication. For attention layers the value is already a tensor; + # for Mamba layers it is a list of tensors that all share the same + # underlying raw storage, so we take the first one. + def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: + assert isinstance(v, torch.Tensor | list) + return v if isinstance(v, torch.Tensor) else v[0] + assert self.cache_config.num_gpu_blocks is not None self.num_blocks = self.cache_config.num_gpu_blocks + use_group_regions = self._kv_cache_config.hisparse_host_num_blocks is not None - seen_storage_ptrs: set[int] = set() - cache_tensors: list[torch.Tensor] = [] - - for cache in kv_caches.values(): - cache = group_kernel_blocks(cache, self.num_blocks) - cache_tensors.append(cache) - cache_storage = cache.untyped_storage() - base_addr = cache_storage.data_ptr() - region_len = cache_storage.nbytes() - - if base_addr not in seen_storage_ptrs: + if not use_group_regions: + seen_storage_ptrs: set[int] = set() + cache_tensors = [ + group_kernel_blocks(_repr_tensor(cache), self.num_blocks) + for cache in kv_caches.values() + ] + for cache in cache_tensors: + cache_storage = cache.untyped_storage() + base_addr = cache_storage.data_ptr() + if base_addr in seen_storage_ptrs: + continue seen_storage_ptrs.add(base_addr) + region_len = cache_storage.nbytes() ret = self.store.register_buffer(base_addr, region_len) if ret != 0: logger.error( @@ -1932,17 +1943,110 @@ class MooncakeStoreWorker: region_len, ret, ) + for db in self.token_dbs: + db.store_layout.register_kv_caches(cache_tensors, self.num_blocks) + + registered_buffers: dict[int, int] = {} + num_host_segments = 0 + num_device_segments = 0 + + for group, db in zip(self._kv_cache_groups, self.token_dbs, strict=True): + if not use_group_regions: + continue + missing_layers = [ + layer_name + for layer_name in group.layer_names + if layer_name not in kv_caches + ] + if missing_layers: + raise ValueError( + f"Missing transferable KV caches for layers: {missing_layers}" + ) + group_caches = [ + (layer_name, kv_caches[layer_name]) for layer_name in group.layer_names + ] + + seen_group_ptrs: set[int] = set() + addrs: list[int] = [] + block_lens: list[int] = [] + for layer_name, value in group_caches: + cache = _repr_tensor(value) + is_host_resident = group.host_resident + if is_host_resident: + num_blocks = self._kv_cache_config.hisparse_host_num_blocks + assert num_blocks is not None + else: + num_blocks = self._kv_cache_config.num_blocks + cache = group_kernel_blocks(cache, num_blocks) + cache_storage = cache.untyped_storage() + if is_host_resident: + base_addr = cache.data_ptr() + region_len = cache.nbytes + else: + base_addr = cache_storage.data_ptr() + region_len = cache_storage.nbytes() + if region_len % num_blocks != 0: + raise ValueError( + f"KV cache region for {layer_name} has size {region_len}, " + f"which is not divisible by its {num_blocks} blocks." + ) + + registered_len = registered_buffers.get(base_addr) + if registered_len is None: + ret = self.store.register_buffer(base_addr, region_len) + if ret != 0: + mem_kind = "host" if is_host_resident else "device" + raise RuntimeError( + "Mooncake register_buffer failed for addr " + f"{base_addr:#x} len {region_len} ({mem_kind}): {ret}" + ) + registered_buffers[base_addr] = region_len + elif registered_len != region_len: + raise ValueError( + f"KV cache views at {base_addr:#x} expose inconsistent " + f"region sizes: {registered_len} and {region_len}." + ) + + regions: list[tuple[int, int]] = [] + if not is_non_overlapping_and_dense(cache[0]): + for head_idx in range(cache.shape[1]): + head_cache = cache[:, head_idx] + assert is_non_overlapping_and_dense(head_cache[0]) + regions.append( + ( + head_cache.data_ptr(), + head_cache.stride(0) * head_cache.element_size(), + ) + ) + elif cache.stride(0) * cache.element_size() * num_blocks == region_len: + regions.append((base_addr, region_len // num_blocks)) + else: + regions.append( + (cache.data_ptr(), cache.stride(0) * cache.element_size()) + ) + + for region_addr, block_len in regions: + if region_addr in seen_group_ptrs: + continue + seen_group_ptrs.add(region_addr) + addrs.append(region_addr) + block_lens.append(block_len) + if is_host_resident: + num_host_segments += 1 + else: + num_device_segments += 1 + + db.set_kv_caches_base_addr(addrs) + db.set_block_len(block_lens) logger.info( - "Registered KV caches: num_groups=%d, num_tensors=%d, num_blocks=%d", + "Registered KV caches: num_groups=%d, num_segments=%d (host=%d, device=%d)", len(self.token_dbs), - len(cache_tensors), - self.num_blocks, + num_host_segments + num_device_segments, + num_host_segments, + num_device_segments, ) - for db in self.token_dbs: - db.store_layout.register_kv_caches(cache_tensors, self.num_blocks) - # Start transfer threads if self.can_put: ready_event_sending = threading.Event() @@ -2316,6 +2420,19 @@ class MooncakeStoreWorker: if store is None: return self.store = None + deadline = time.monotonic() + 5.0 + for thread in (self.kv_send_thread, *self.kv_recv_threads): + if thread is None: + continue + while thread.request_queue.unfinished_tasks and time.monotonic() < deadline: + time.sleep(0.05) + if thread.request_queue.unfinished_tasks: + logger.warning( + "Mooncake store %s still has %d in-flight requests at " + "close; closing anyway.", + thread.name, + thread.request_queue.unfinished_tasks, + ) try: store.close() except Exception as e: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index 591ded1c67..d291d38528 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -17,6 +17,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorHandshakeMetadata, KVConnectorMetadata, KVConnectorRole, + KVConnectorTransferResults, KVConnectorWorkerMetadata, SupportsHMA, ) @@ -35,7 +36,7 @@ if TYPE_CHECKING: from vllm.distributed.kv_events import KVCacheEvent from vllm.forward_context import ForwardContext from vllm.v1.core.block_pool import BlockPool - from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request @@ -208,6 +209,10 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA): # Propagated from scheduler to worker side via the connector metadata. self._extra_async_saves: dict[str, int] = {} + @property + def sub_connectors(self) -> list[KVConnectorBase_V1]: + return list(self._connectors) + @property def supports_divergent_local_hybrid_hits(self) -> bool: return bool(self._connectors) and all( @@ -248,6 +253,10 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA): for c in self._connectors: c.register_kv_caches(kv_caches) + def bind_kv_cache_manager(self, kv_cache_manager: "KVCacheManager") -> None: + for connector in self._connectors: + connector.bind_kv_cache_manager(kv_cache_manager) + def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None: for c in self._connectors: c.bind_gpu_block_pool(gpu_block_pool) @@ -292,6 +301,14 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA): for c in self._connectors: c.start_load_kv(forward_context, **kwargs) + def finish_forward(self) -> None: + for c in self._connectors: + c.finish_forward() + + def reset_capture_state(self) -> None: + for c in self._connectors: + c.reset_capture_state() + def wait_for_layer_load(self, layer_name: str) -> None: for c in self._connectors: c.wait_for_layer_load(layer_name) @@ -313,29 +330,31 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA): def get_finished( self, finished_req_ids: set[str] ) -> tuple[set[str] | None, set[str] | None]: - finished_sending: set[str] = set() - finished_recving: set[str] = set() - for c in self._connectors: - sending, recving = c.get_finished(finished_req_ids) - if not recving and not sending: - continue - # Aggregate finished recving request ids. - finished_recving.update(recving or ()) - # Aggregate finished sending request ids - only include - # once we've drained the "extra" count (for cases where - # more than one connector is async-saving the same request). - for req_id in sending or (): + results = self.get_transfer_results(finished_req_ids) + return ( + results.finished_sending or None, + results.finished_recving or None, + ) + + def get_transfer_results( + self, finished_req_ids: set[str] + ) -> KVConnectorTransferResults: + results = KVConnectorTransferResults() + for connector in self._connectors: + child_results = connector.get_transfer_results(finished_req_ids) + results.finished_recving.update(child_results.finished_recving) + results.failed_recving.update(child_results.failed_recving) + for req_id in child_results.finished_sending: extra_pending = self._extra_async_saves.get(req_id) if extra_pending is None: - finished_sending.add(req_id) - continue - assert extra_pending > 0 - if extra_pending == 1: - del self._extra_async_saves[req_id] + results.finished_sending.add(req_id) else: - self._extra_async_saves[req_id] = extra_pending - 1 - - return finished_sending or None, finished_recving or None + assert extra_pending > 0 + if extra_pending == 1: + del self._extra_async_saves[req_id] + else: + self._extra_async_saves[req_id] = extra_pending - 1 + return results def get_block_ids_with_load_errors(self) -> set[int]: agg_block_ids: set[int] = set() @@ -563,6 +582,9 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA): for c in self._connectors: yield from c.take_events() + def has_pending_block_frees(self) -> bool: + return any(c.has_pending_block_frees() for c in self._connectors) + def has_pending_push_work(self) -> bool: return any(c.has_pending_push_work() for c in self._connectors) @@ -630,7 +652,7 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA): return MultiKVConnectorStats(data=reconstructed_data) - def get_kv_connector_stats(self) -> MultiKVConnectorStats | None: + def get_kv_connector_stats(self) -> KVConnectorStats | None: # Group connector stats by connector type. stats_by_connector: MultiKVConnectorStats | None = None for c in self._connectors: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index ddfcf35e03..00579c40f0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -93,10 +93,7 @@ class NixlBaseConnectorScheduler: for g in kv_cache_config.transfer_groups ) ) - self._has_mamba = any( - isinstance(g.kv_cache_spec, MambaSpec) - for g in kv_cache_config.transfer_groups - ) + self._has_mamba = kv_cache_config.has_mamba_layers logger.info("Initializing NIXL Scheduler %s", engine_id) if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index e4b46cb64a..17660da451 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -2,8 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Base worker-side logic for the NIXL connector.""" +import contextlib import itertools import logging +import math import os import queue import threading @@ -31,7 +33,10 @@ from vllm.distributed.kv_transfer.kv_connector.utils import ( kv_postprocess_blksize_on_receive, kv_postprocess_layout_on_receive, ) -from vllm.distributed.kv_transfer.kv_connector.v1.base import CopyBlocksOp +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + CopyBlocksOp, + KVConnectorTransferResults, +) from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( GET_META_MSG, @@ -186,23 +191,28 @@ class NixlBaseConnectorWorker: num_blocks = dst_num_blocks if block_size_ratio is not None: num_blocks = int(num_blocks * block_size_ratio) + num_regions = self.num_regions or len(self.block_len_per_layer) if region_num_blocks is None: - region_num_blocks = [num_blocks] * self.num_regions + region_num_blocks = [num_blocks] * num_regions elif block_size_ratio is not None: region_num_blocks = [ int(count * block_size_ratio) for count in region_num_blocks ] - assert len(region_num_blocks) == self.num_regions + if region_num_blocks: + num_regions = len(region_num_blocks) + assert len(region_num_blocks) == num_regions region_offsets = np.cumsum([0, *region_num_blocks[:-1]]) num_fa_descs = sum(region_num_blocks) - # All-attention fast path: single vectorized broadcast. + # All-attention fast path. if num_ssm_regions == 0: if region_group_ids is None: region_group_ids = self.region_group_ids if uses_region_group_mapping is None: uses_region_group_mapping = self._uses_region_group_mapping - assert len(region_group_ids) == self.num_regions + if not region_group_ids and num_regions == 1: + region_group_ids = [0] + assert len(region_group_ids) == num_regions if uses_region_group_mapping is None: uses_region_group_mapping = len(set(region_group_ids)) > 1 region_group_ids_array = np.asarray(region_group_ids, dtype=np.int32) @@ -397,8 +407,6 @@ class NixlBaseConnectorWorker: # Config. self.vllm_config = vllm_config - # mypy will complain on re-assignment otherwise. - self.block_size: int = cast(int, vllm_config.cache_config.block_size) if vllm_config.kv_transfer_config is None: raise ValueError("kv_transfer_config must be set for NixlConnector") @@ -420,6 +428,17 @@ class NixlBaseConnectorWorker: ) self.kv_cache_config = kv_cache_config + transfer_block_sizes = [ + group.kv_cache_spec.block_size + for group in kv_cache_config.transfer_groups + if get_representative_spec_type(group.kv_cache_spec) + not in (MambaSpec, CircularBufferSpec) + ] + self.block_size = ( + math.lcm(*transfer_block_sizes) + if transfer_block_sizes + else cast(int, vllm_config.cache_config.block_size) + ) # Per-layer specs, unwrapping UniformTypeKVCacheSpecs group wrappers. self._layer_specs: dict[str, KVCacheSpec] = {} for group in kv_cache_config.transfer_groups: @@ -599,12 +618,15 @@ class NixlBaseConnectorWorker: # Number of NIXL regions. Currently one region per cache # (so 1 per layer for MLA, otherwise 2 per layer) self.num_regions = 0 + self.region_mem_types: list[str] = [] self.region_group_ids: list[int] = [] self._uses_region_group_mapping = False self.region_names: list[str] = [] self.region_num_blocks: list[int] = [] - self.region_mem_types: list[str] = [] self._mixed_mem_types = False + self._desc_is_dram_by_block_size: dict[int, np.ndarray] = {} + self._desc_pos_by_block_size: dict[int, np.ndarray] = {} + self._dram_src_handles_by_block_size: dict[int, int] = {} # PP>1 (push mode): this worker holds a contiguous layer slice and # transfers into the matching sub-range of a PP=1 remote's regions. @@ -644,9 +666,6 @@ class NixlBaseConnectorWorker: self.dst_region_group_ids: dict[EngineId, list[int]] = {} self.dst_uses_region_group_mapping: dict[EngineId, bool] = {} self.dst_region_mem_types: dict[EngineId, list[str]] = {} - self._desc_is_dram_by_block_size: dict[int, np.ndarray] = {} - self._desc_pos_by_block_size: dict[int, np.ndarray] = {} - self._dram_src_handles_by_block_size: dict[int, int] = {} self._registered_descs: list[Any] = [] # In progress transfers. @@ -667,6 +686,7 @@ class NixlBaseConnectorWorker: # Uses Queue for thread-safe cross-thread coordination with the # background handshake thread, matching the _ready_requests pattern. self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue() + self._failed_inflight_recvs: set[ReqId] = set() self._pending_recv_notifs: dict[ReqId, list[tuple[str, bytes]]] = {} # Handshake metadata of this worker for NIXL transfers. @@ -690,7 +710,6 @@ class NixlBaseConnectorWorker: "engine_ttl", 3600.0 ) - self.block_size = vllm_config.cache_config.block_size self.model_config = vllm_config.model_config self.use_mla = self.model_config.use_mla @@ -1240,7 +1259,7 @@ class NixlBaseConnectorWorker: self.use_host_buffer, ) - registration_ranges: dict[tuple[int, str], tuple[int, int, int, str]] = {} + registration_ranges: dict[tuple[int, str], tuple[int, int, int]] = {} region_mem_types: list[str] = [] seen_base_addresses: list[int] = [] self._ssm_region_indices = [] @@ -1280,10 +1299,16 @@ class NixlBaseConnectorWorker: # K and V are packed into the content dim, so each attention layer is a # single NIXL region whose block transfers as one unit. Mamba layers instead # register separate conv/ssm sub-regions (see `_build_mamba_local`). + # P and D may allocate equivalent transferable layers in different + # cache-group orders. Keep their region lists aligned without putting + # layers.10 before layers.2, which would break PP region slicing. layer_names = ( xfer_buffers if self._is_csa_linear - else sorted(xfer_buffers, key=_region_sort_key) + else sorted( + xfer_buffers, + key=_region_sort_key, + ) ) for layer_name in layer_names: cache = xfer_buffers[layer_name] @@ -1299,6 +1324,7 @@ class NixlBaseConnectorWorker: layer_name, ) continue + group_index = self.kv_cache_config.transfer_group_index_by_layer[layer_name] # `layer_spec.page_size_bytes` only accounts for logical page_size, that is # the page_size assuming constant `self._logical_num_blocks`. physical_page_size = ( @@ -1307,17 +1333,48 @@ class NixlBaseConnectorWorker: else layer_spec.page_size_bytes // self._physical_blocks_per_logical_kv_block ) - group_id = self.kv_cache_config.transfer_group_index_by_layer[layer_name] + group = self.kv_cache_config.transfer_groups[group_index] + if group.host_resident: + logical_num_blocks = self.kv_cache_config.hisparse_host_num_blocks + assert logical_num_blocks is not None + else: + logical_num_blocks = self.kv_cache_config.num_blocks + group_id = group_index num_blocks = ( - self._logical_num_blocks + logical_num_blocks if isinstance(layer_spec, MambaSpec) - else self.num_blocks + else logical_num_blocks * self._physical_blocks_per_logical_kv_block + ) + base_addr = cache.data_ptr() + is_mla_region = isinstance( + layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec) ) logger.debug( "Registering layer %s with cache shape: %s", layer_name, cache.shape ) storage = cache.untyped_storage() storage_addr = storage.data_ptr() + is_host_resident = group.host_resident + if cache.device.type == "cpu": + mem_type = "DRAM" + region_device_id = 0 + else: + mem_type = self.nixl_memory_type + region_device_id = max(cache.get_device(), 0) + self.device_id = region_device_id + if is_host_resident: + registration_base = cache.data_ptr() + registration_len = cache.nbytes + else: + registration_base = storage_addr + registration_len = storage.nbytes() + storage_key = (registration_base, mem_type) + if storage_key not in registration_ranges: + registration_ranges[storage_key] = ( + registration_base, + registration_base + registration_len, + region_device_id, + ) if isinstance(layer_spec, KpoolTailSpec): compressed_owner = compressed_region_owners.get(cache.data_ptr()) @@ -1345,22 +1402,6 @@ class NixlBaseConnectorWorker: ) continue - if cache.device.type == "cpu": - mem_type = "DRAM" - region_device_id = 0 - else: - mem_type = self.nixl_memory_type - region_device_id = max(cache.get_device(), 0) - self.device_id = region_device_id - registration_ranges.setdefault( - (storage_addr, mem_type), - (storage_addr, storage.nbytes(), region_device_id, ""), - ) - - is_mla_region = isinstance( - layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec) - ) - if isinstance(layer_spec, MambaSpec): physical_ratio = self._physical_blocks_per_logical_kv_block block_len = physical_page_size // physical_ratio @@ -1381,14 +1422,13 @@ class NixlBaseConnectorWorker: block_stride = physical_page_size else: block_stride = cache.stride(0) * cache.element_size() - storage_is_block_major = num_blocks * block_stride == storage.nbytes() - # A layer whose [H, N, C] interior is dense addresses its own page - # as one chunk. Otherwise the block's whole row is the only - # contiguous transfer unit. - hnc_contiguous = ( - cache.ndim == 4 - and cache.stride(2) == cache.shape[3] - and cache.stride(1) == cache.shape[2] * cache.shape[3] + storage_is_block_major = num_blocks * block_stride == registration_len + # A view with a contiguous page can own a strided region even + # when sibling layers share the same backing allocation. + page_contiguous = ( + cache.ndim > 1 + and cache[0].is_contiguous() + and cache[0].nbytes == physical_page_size ) virtual_transfer_pages = _uses_dense_virtual_transfer_pages( layer_spec, cache, physical_page_size, num_blocks @@ -1399,13 +1439,14 @@ class NixlBaseConnectorWorker: (cache.data_ptr(), physical_page_size, physical_page_size) ] elif storage_is_block_major and ( - (packed_storage and is_mla_region) - or (not hnc_contiguous and not self._is_csa_linear) + not page_contiguous + and ((packed_storage and is_mla_region) or not self._is_csa_linear) ): - # Packed MLA layouts transfer the complete storage row. - storage_block_len = storage.nbytes() // num_blocks + # TODO(Lucas): handle TP slicing for packed_storage; for now + # restrict to MLA (DSv4) where kv is replicated. + storage_block_len = registration_len // num_blocks region_specs = [ - (storage_addr, storage_block_len, storage_block_len) + (registration_base, storage_block_len, storage_block_len) ] elif storage_is_block_major: region_specs = [ @@ -1440,6 +1481,10 @@ class NixlBaseConnectorWorker: region_index = seen_base_addresses.index(base_addr) assert region_mem_types[region_index] == mem_type self._region_is_mla[region_index] |= is_mla_region + if is_mla_region: + self.block_len_per_layer[region_index] = block_len + self.block_stride_per_layer[region_index] = block_stride + self.region_num_blocks[region_index] = num_blocks if self.region_group_ids[region_index] != group_id: self.region_group_ids[region_index] = _SHARED_REGION_GROUP_ID else: @@ -1450,8 +1495,8 @@ class NixlBaseConnectorWorker: self.region_group_ids.append(group_id) self.region_names.append(layer_name) self.region_num_blocks.append(num_blocks) - region_mem_types.append(mem_type) self._region_is_mla.append(is_mla_region) + region_mem_types.append(mem_type) if isinstance(layer_spec, MambaSpec): if layer_spec.tp_replicated: @@ -1496,7 +1541,6 @@ class NixlBaseConnectorWorker: == len(self.region_group_ids) == len(self.region_names) == len(self.region_num_blocks) - == len(region_mem_types) ) # Descriptor ids must be region-ordered, matching the remote side. self._scratch_region_indices.sort() @@ -1521,16 +1565,23 @@ class NixlBaseConnectorWorker: # Total local FA descriptors (boundary between FA and mamba descs). self.num_descs = sum(self.region_num_blocks) + self._mixed_mem_types = len(set(region_mem_types)) > 1 + if self._mixed_mem_types: + assert self.use_mla and not self._has_mamba, ( + "Mixed-device KV registration is only supported for MLA " + "models without Mamba layers." + ) for mem_type in sorted(set(region_mem_types)): - ranges = [ - registration - for ( - _, - registration_mem_type, - ), registration in registration_ranges.items() - if registration_mem_type == mem_type + ranges_for_mem_type = [ + (start, end - start, device_id, "") + for (_, cache_mem_type), ( + start, + end, + device_id, + ) in registration_ranges.items() + if cache_mem_type == mem_type ] - descs = self.nixl_wrapper.get_reg_descs(ranges, mem_type) + descs = self.nixl_wrapper.get_reg_descs(ranges_for_mem_type, mem_type) self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends) self._registered_descs.append(descs) @@ -1854,7 +1905,11 @@ class NixlBaseConnectorWorker: if self._mixed_mem_types: desc_is_dram = np.concatenate( [ - np.full(count * block_size_ratio, mem_type == "DRAM", dtype=bool) + np.full( + count * block_size_ratio, + mem_type == "DRAM", + dtype=bool, + ) for mem_type, count in zip( self.region_mem_types, self.region_num_blocks, @@ -1864,24 +1919,29 @@ class NixlBaseConnectorWorker: ) assert len(desc_is_dram) == len(blocks_data) desc_pos = np.empty(len(desc_is_dram), dtype=np.int64) - dram_indices = np.flatnonzero(desc_is_dram) - device_indices = np.flatnonzero(~desc_is_dram) - desc_pos[dram_indices] = np.arange(len(dram_indices), dtype=np.int64) - desc_pos[device_indices] = np.arange(len(device_indices), dtype=np.int64) + dram_idx = np.where(desc_is_dram)[0] + vram_idx = np.where(~desc_is_dram)[0] + desc_pos[dram_idx] = np.arange(len(dram_idx), dtype=np.int64) + desc_pos[vram_idx] = np.arange(len(vram_idx), dtype=np.int64) self._desc_is_dram_by_block_size[block_size] = desc_is_dram self._desc_pos_by_block_size[block_size] = desc_pos - blocks_data[desc_is_dram, 2] = 0 - dram_descs = self.nixl_wrapper.get_xfer_descs( - blocks_data[desc_is_dram], "DRAM" - ) + + # DRAM descriptors are registered under CPU device 0. + blocks_data = [ + (addr, length, 0) if is_dram else (addr, length, dev) + for (addr, length, dev), is_dram in zip( + blocks_data, desc_is_dram, strict=True + ) + ] + dram_blocks = [blocks_data[i] for i in dram_idx] + vram_blocks = [blocks_data[i] for i in vram_idx] + dram_descs = self.nixl_wrapper.get_xfer_descs(dram_blocks, "DRAM") self._dram_src_handles_by_block_size[block_size] = ( self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", dram_descs) ) - device_descs = self.nixl_wrapper.get_xfer_descs( - blocks_data[~desc_is_dram], self.nixl_memory_type - ) + descs = self.nixl_wrapper.get_xfer_descs(vram_blocks, self.nixl_memory_type) return ( - self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", device_descs), + self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs), blocks_data, ) @@ -1968,10 +2028,6 @@ class NixlBaseConnectorWorker: ] nixl_agent_meta.block_lens = nixl_agent_meta.block_lens[start:end] nixl_agent_meta.block_strides = nixl_agent_meta.block_strides[start:end] - if nixl_agent_meta.region_num_blocks is not None: - nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[ - start:end - ] if nixl_agent_meta.region_group_ids is not None: nixl_agent_meta.region_group_ids = nixl_agent_meta.region_group_ids[ start:end @@ -1982,6 +2038,10 @@ class NixlBaseConnectorWorker: nixl_agent_meta.region_mem_types = nixl_agent_meta.region_mem_types[ start:end ] + if nixl_agent_meta.region_num_blocks is not None: + nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[ + start:end + ] ### Register remote engine in TransferTopology (idempotent). assert self.transfer_topo is not None @@ -2026,11 +2086,12 @@ class NixlBaseConnectorWorker: nixl_agent_meta.region_num_blocks or [nixl_agent_meta.num_blocks] * num_remote_regions ) - self.dst_region_group_ids[engine_id] = nixl_agent_meta.region_group_ids or ( + remote_region_group_ids = nixl_agent_meta.region_group_ids or ( self.region_group_ids if len(self.region_group_ids) == num_remote_regions else [0] * num_remote_regions ) + self.dst_region_group_ids[engine_id] = remote_region_group_ids self.dst_uses_region_group_mapping[engine_id] = ( len(set(self.dst_region_group_ids[engine_id])) > 1 ) @@ -2138,8 +2199,6 @@ class NixlBaseConnectorWorker: mamba = self._build_mamba_remote(nixl_agent_meta, tp_ratio, transfer_info) blocks_data = np.concatenate([blocks_data, mamba]) - # A mixed local destination is split into multiple READs below. The - # producer remains a single descriptor list in this prerequisite. remote_mem_types = set(self.dst_region_mem_types[engine_id]) if len(remote_mem_types) != 1: raise NotImplementedError( @@ -2540,9 +2599,10 @@ class NixlBaseConnectorWorker: indices=indices, ) - def get_finished(self) -> tuple[set[str], set[str]]: + def get_transfer_results(self) -> KVConnectorTransferResults: """ - Get requests that are done sending or recving on this specific worker. + Get transfers that completed on this specific worker. + The scheduler process (via the MultiprocExecutor) will use this output to track which workers are done. """ @@ -2564,7 +2624,6 @@ class NixlBaseConnectorWorker: # Add failed requests to done_recving for scheduler tracking # (blocks are already marked invalid, scheduler will handle recompute) done_recving.update(failed_recv_reqs) - if len(done_sending) > 0 or len(done_recving) > 0: logger.debug( "Rank %s, get_finished: %s requests done sending " @@ -2577,6 +2636,7 @@ class NixlBaseConnectorWorker: block_ids_for_blocksize_post_process = defaultdict(list) block_ids_for_heterogeneous_attn_post_process = list[list[int]]() + direct_device_recving = set[str]() for req_id in done_recving: # clean up metadata for completed requests meta = self._recving_metadata.pop(req_id, None) @@ -2594,6 +2654,8 @@ class NixlBaseConnectorWorker: if self.use_host_buffer: self.sync_recved_kv_to_device(req_id, meta) + direct_device_recving.add(req_id) + # Post processing for heteroblocksize/layout, and for blocks the # transfer clipped. The latter happens either at remote-block # granularity (block_size_ratio > 1) or at kernel-block @@ -2642,7 +2704,7 @@ class NixlBaseConnectorWorker: for block_ids in block_ids_for_heterogeneous_attn_post_process: self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids) - self._sync_device_after_mamba_recv(done_recving, failed_recv_reqs) + self._sync_device_after_direct_recv(direct_device_recving) # Handle timeout to avoid stranding blocks on remote. now = time.perf_counter() @@ -2664,20 +2726,21 @@ class NixlBaseConnectorWorker: del self._reqs_to_send[req_id] done_sending.add(req_id) - return done_sending, done_recving + return KVConnectorTransferResults( + finished_sending=done_sending, + finished_recving=done_recving, + failed_recving=failed_recv_reqs, + ) - def _sync_device_after_mamba_recv( - self, - done_recving: set[str], - failed_recv_reqs: set[str], - ) -> None: - """Synchronize ROCm direct-GPU Mamba receives before model execution.""" - if ( - not current_platform.is_rocm() - or not self._has_mamba - or self.use_host_buffer - or not (done_recving - failed_recv_reqs) - ): + def get_finished(self) -> tuple[set[str], set[str]]: + """Compatibility wrapper for the legacy completion API.""" + results = self.get_transfer_results() + return results.finished_sending, results.finished_recving + + def _sync_device_after_direct_recv(self, done_recving: set[str]) -> None: + """Make direct NIXL writes visible before model execution.""" + requires_sync = current_platform.is_rocm() and self._has_mamba + if self.use_host_buffer or not done_recving or not requires_sync: return torch.accelerator.synchronize() @@ -2731,11 +2794,11 @@ class NixlBaseConnectorWorker: self.nixl_wrapper.release_xfer_handle(handle) elif xfer_state == "PROC": in_progress.append(handle) - continue else: self._log_failure( failure_type="transfer_failed", - msg="Marking blocks as invalid", + msg="Deferring request completion until its last " + "xfer is terminal", req_id=req_id, xfer_state=xfer_state, ) @@ -2743,21 +2806,20 @@ class NixlBaseConnectorWorker: except Exception as e: self._log_failure( failure_type="transfer_exception", - msg="Marking blocks as invalid", + msg="Handle is unpollable; treating it as terminal", req_id=req_id, error=e, ) self._handle_failed_transfer(req_id, handle) if not in_progress: - # Only report request as completed when all transfers are done. - # A request failed in an earlier poll was already reported via - # _failed_recv_reqs and its metadata popped by get_finished(); - # don't report it again, just drop the remaining handles. - if req_id in self._recving_metadata: - done_req_ids.add(req_id) - self._send_pending_recv_notifs(req_id) del transfers[req_id] + done_req_ids.add(req_id) + if req_id in self._failed_inflight_recvs: + self._failed_inflight_recvs.remove(req_id) + self._report_failed_recv(req_id) + else: + self._send_pending_recv_notifs(req_id) else: transfers[req_id] = in_progress return done_req_ids @@ -2771,22 +2833,25 @@ class NixlBaseConnectorWorker: req_id: The request ID. handle: The transfer handle. """ - # (multi-read) One handle is created per remote rank, and they do not - # all fail in the same _pop_done_transfers poll. The request is - # reported failed on the first one, which pops its metadata in - # get_finished(); on the later failures only the handle cleanup is left. - # TODO (NickLucche) handle failed transfer for HMA. - # A split READ's notification is sent only after every handle succeeds. + # A sibling READ may still be writing these blocks. Retain metadata + # and defer invalidation until every handle is terminal. self._pending_recv_notifs.pop(req_id, None) - if (meta := self._recving_metadata.get(req_id)) is not None: - if not self._is_hma_required: - self._invalid_block_ids.put(set(meta.local_block_ids[0])) - self._failed_recv_reqs.put(req_id) + if req_id in self._recving_transfers: + self._failed_inflight_recvs.add(req_id) + else: + self._report_failed_recv(req_id) if handle is not None: self.nixl_wrapper.release_xfer_handle(handle) self.xfer_stats.record_failed_transfer() + def _report_failed_recv(self, req_id: str) -> None: + if (meta := self._recving_metadata.get(req_id)) is not None: + if not self._is_hma_required: + self._invalid_block_ids.put(set(meta.local_block_ids[0])) + self._failed_recv_reqs.put(req_id) + def _send_pending_recv_notifs(self, req_id: str) -> None: + """Send notifications deferred by split DRAM/VRAM reads.""" for agent_name, notif_id in self._pending_recv_notifs.pop(req_id, []): try: self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id) @@ -2909,10 +2974,10 @@ class NixlBaseConnectorWorker: return block_ids block_arange = np.arange(0, ratio).reshape(1, -1) # Mamba blocks have no logical<>physical discrepancy (block-size=1) - group_specs = self.kv_cache_config.transfer_groups physical_block_ids = [] for i, group in enumerate(block_ids): - if _is_ssm_spec(get_representative_spec_type(group_specs[i].kv_cache_spec)): + spec = self.kv_cache_config.transfer_groups[i].kv_cache_spec + if _is_ssm_spec(get_representative_spec_type(spec)): physical_block_ids.append(group) else: physical_block_ids.append( @@ -3215,9 +3280,44 @@ class NixlBaseConnectorWorker: ) def __del__(self): - self.shutdown() + with contextlib.suppress(Exception): + self.shutdown() - def shutdown(self): + def _finish_shutdown(self) -> None: + self._recving_transfers.clear() + try: + for handle in self.src_xfer_handles_by_block_size.values(): + self.nixl_wrapper.release_dlist_handle(handle) + for handles in self.src_xfer_handles_by_tp_ratio.values(): + for handle in handles: + self.nixl_wrapper.release_dlist_handle(handle) + for handles in self._dram_src_handles_by_tp_ratio.values(): + for handle in handles: + self.nixl_wrapper.release_dlist_handle(handle) + for handle in self._dram_src_handles_by_block_size.values(): + self.nixl_wrapper.release_dlist_handle(handle) + except Exception: + logger.exception("NIXL dlist-handle release failed at shutdown.") + self.src_xfer_handles_by_block_size.clear() + self.src_xfer_handles_by_tp_ratio.clear() + self._dram_src_handles_by_tp_ratio.clear() + self._dram_src_handles_by_block_size.clear() + try: + for engine_id in list(self._remote_agents): + self._cleanup_remote_engine(engine_id, log_eviction=False) + except Exception: + logger.exception("NIXL remote-engine cleanup failed at shutdown.") + try: + for desc in self._registered_descs: + self.nixl_wrapper.deregister_memory(desc) + finally: + self._registered_descs.clear() + # Drop cache references before their owners release registered + # host memory; handshake futures may outlive model-runner shutdown. + self.device_kv_caches = {} + self.host_xfer_buffers = {} + + def shutdown(self) -> None: """Shutdown the connector worker.""" if not hasattr(self, "_handshake_initiation_executor"): # error happens during init, no need to shutdown @@ -3227,22 +3327,4 @@ class NixlBaseConnectorWorker: for handle in handles: self.nixl_wrapper.release_xfer_handle(handle) self._recving_transfers.clear() - for handle in self.src_xfer_handles_by_block_size.values(): - self.nixl_wrapper.release_dlist_handle(handle) - self.src_xfer_handles_by_block_size.clear() - for handles in self.src_xfer_handles_by_tp_ratio.values(): - for handle in handles: - self.nixl_wrapper.release_dlist_handle(handle) - self.src_xfer_handles_by_tp_ratio.clear() - for handles in self._dram_src_handles_by_tp_ratio.values(): - for handle in handles: - self.nixl_wrapper.release_dlist_handle(handle) - self._dram_src_handles_by_tp_ratio.clear() - for handle in self._dram_src_handles_by_block_size.values(): - self.nixl_wrapper.release_dlist_handle(handle) - self._dram_src_handles_by_block_size.clear() - for engine_id in list(self._remote_agents): - self._cleanup_remote_engine(engine_id, log_eviction=False) - for desc in self._registered_descs: - self.nixl_wrapper.deregister_memory(desc) - self._registered_descs.clear() + self._finish_shutdown() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py index 7914b7ab3a..105617bb8d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py @@ -26,6 +26,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorHandshakeMetadata, KVConnectorMetadata, KVConnectorRole, + KVConnectorTransferResults, SupportsHMA, ) from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( @@ -234,6 +235,18 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA): assert self.connector_worker is not None return self.connector_worker.get_finished() + def get_transfer_results( + self, finished_req_ids: set[str] + ) -> KVConnectorTransferResults: + assert self.connector_worker is not None + results = self.connector_worker.get_transfer_results() + if ( + self.kv_transfer_config.kv_role == "kv_producer" + and self.connector_worker.pcp_rank > 0 + ): + results.finished_sending.clear() + return results + def get_block_ids_with_load_errors(self) -> set[int]: """Get block IDs that failed to load via NIXL.""" assert self.connector_worker is not None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index 526d9c6881..f61fdfa835 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -305,7 +305,9 @@ class NixlConnectorMetadata(KVConnectorMetadata): local_num_computed_blocks: tuple[int, ...] = (), ): req = self._add_new_req( - local_block_ids, kv_transfer_params, local_num_computed_blocks + local_block_ids, + kv_transfer_params, + local_num_computed_blocks, ) req.remote = RemoteMeta( block_ids=kv_transfer_params["remote_block_ids"], diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index d39a9fe287..ceaa4fce57 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -166,11 +166,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): dcp_active = self.dcp_size > 1 or remote_info.remote_dcp_size > 1 local_block_ids = meta.local_physical_block_ids - local_region_groups = getattr(self, "region_group_ids", []) - remote_region_groups = getattr(self, "dst_region_group_ids", {}).get( - engine_id, local_region_groups - ) - if local_region_groups != remote_region_groups: + remote_region_groups = self.dst_region_group_ids[engine_id] + local_region_groups = self.region_group_ids or remote_region_groups + groups_differ = local_region_groups != remote_region_groups + if groups_differ: if not self.use_mla or self._has_mamba: raise NotImplementedError( "Different NIXL cache-group layouts are only supported for " @@ -292,6 +291,9 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): spec.remote_rank ] + # Once a read routes the request to failure reporting, the + # scheduler may free and reuse its blocks, so no sibling READs + # may be posted (and P must not be notified). if not self._read_blocks( read_spec=spec, request_id=req_id, @@ -329,6 +331,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): """ Post a READ point-to-point xfer request from a single local worker to a single remote worker. + + Returns True when the read was posted (or was unnecessary), False + when the request was routed to failure reporting — the caller must + not post further transfers for it. """ assert self.transfer_topo is not None remote_rank = read_spec.remote_rank @@ -416,11 +422,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): dst_num_blocks=self.dst_num_blocks[dst_engine_id], block_size_ratio=None, physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical, - region_num_blocks=self.dst_region_num_blocks[dst_engine_id], + region_num_blocks=(self.dst_region_num_blocks.get(dst_engine_id) or None), region_group_ids=( list(range(self.num_regions)) if read_spec.block_ids_by_region - else self.dst_region_group_ids[dst_engine_id] + else (self.dst_region_group_ids.get(dst_engine_id) or None) ), uses_region_group_mapping=( self.num_regions > 1 @@ -433,11 +439,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): dst_num_blocks=self.dst_num_blocks[self.engine_id], block_size_ratio=block_size_ratio, physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block, - region_num_blocks=self.dst_region_num_blocks[self.engine_id], + region_num_blocks=(self.dst_region_num_blocks.get(self.engine_id) or None), region_group_ids=( list(range(self.num_regions)) if read_spec.block_ids_by_region - else self.region_group_ids + else (self.region_group_ids or None) ), uses_region_group_mapping=( self.num_regions > 1 @@ -537,11 +543,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): self._pending_recv_notifs.setdefault(request_id, []).append( (notif_agent, notif_id) ) - for index, handle in enumerate(handles): + for i, handle in enumerate(handles): try: self.nixl_wrapper.transfer(handle) except Exception: - for unstarted in handles[index:]: + for unstarted in handles[i:]: self.nixl_wrapper.release_xfer_handle(unstarted) raise self._recving_transfers[request_id].append(handle) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 21bd956c67..c3bcd2c8f9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -41,6 +41,9 @@ from typing import TYPE_CHECKING, Any import msgspec from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorTransferResults, +) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import ( NixlBaseConnectorWorker, ) @@ -145,7 +148,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): self._push_writer_thread.start() logger.info("nixl-push-writer thread started (rank=%d)", self.tp_rank) - def shutdown(self): + def shutdown(self) -> None: self._push_writer_stop.set() # Unblock the writer if it's waiting in the no-active-state branch. self._push_writer_wake.set() @@ -522,7 +525,14 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): ) remote_block_ids = meta.remote.block_ids local_block_ids = meta.local_physical_block_ids - num_groups = len(local_block_ids) + local_region_groups = self.region_group_ids + remote_region_groups = self.dst_region_group_ids[engine_id] + groups_differ = local_region_groups != remote_region_groups + if groups_differ: + raise NotImplementedError( + "NixlPushConnector does not support different producer and " + "consumer cache-group layouts" + ) # MLA latent is replicated across D's TP ranks: the tp-mapping # collapses it to one rank (fine for reads), but push must WRITE every @@ -537,6 +547,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): else: write_ranks = list(plan.all_source_ranks) + num_groups = len(local_block_ids) + def group_ids(block_ids: BlockIds, rank: int) -> BlockIds: return [ list(block_ids[g]) @@ -777,13 +789,14 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): self._reqs_to_send.pop(req_id, None) return notified_req_ids - def get_finished(self) -> tuple[set[str], set[str]]: + def get_transfer_results(self) -> KVConnectorTransferResults: # Engine main thread asking for completions: also wake the writer # so it gets a chance to drain NIXL notifs (heartbeats, completion # notifs, late PUSH_REGs) even if it had been parked. self._push_writer_wake.set() - done_sending, done_recving = super().get_finished() + results = super().get_transfer_results() + done_sending = results.finished_sending # ``_pop_done_transfers`` mutates ``_sending_transfers``; the # writer thread also appends to it, so guard the pop. @@ -803,4 +816,4 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): if done_sending: self._push_writer_wake.set() - return done_sending, done_recving + return results diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index bc78b72988..42b9dd6aa3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -12,6 +12,7 @@ from vllm.v1.core.kv_cache_utils import ( from vllm.v1.kv_cache_interface import ( AttentionSpec, FullAttentionSpec, + KVCacheGroupRole, KVCacheSpec, MLAAttentionSpec, SlidingWindowMLASpec, @@ -31,6 +32,16 @@ if TYPE_CHECKING: from vllm.v1.kv_cache_interface import KVCacheConfig +def get_offloading_group_ids(kv_cache_config: "KVCacheConfig") -> tuple[int, ...]: + if kv_cache_config.hisparse_host_num_blocks is None: + return tuple(range(len(kv_cache_config.kv_cache_groups))) + return tuple( + group_id + for group_id, group in enumerate(kv_cache_config.kv_cache_groups) + if group.role is KVCacheGroupRole.HISPARSE_INDEXER + ) + + def build_offloading_config( vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig", @@ -43,15 +54,22 @@ def build_offloading_config( engine_id = kv_transfer_config.engine_id parallel_config = vllm_config.parallel_config + selected_groups = tuple( + (group_id, kv_cache_config.kv_cache_groups[group_id]) + for group_id in get_offloading_group_ids(kv_cache_config) + ) + if not selected_groups: + raise ValueError("KV offloading found no eligible cache groups.") groups = tuple( OffloadingGroupConfig( + group_id=group_id, tokens_per_block=resolve_dcp_kv_block_size( group.kv_cache_spec, parallel_config.decode_context_parallel_size, ), layer_names=tuple(group.layer_names), ) - for group in kv_cache_config.kv_cache_groups + for group_id, group in selected_groups ) _, tokens_per_hash = resolve_kv_cache_block_sizes(kv_cache_config, vllm_config) @@ -103,11 +121,20 @@ def build_offloading_config( ) worker_kv_bytes_per_block = 0 - if kv_cache_config.num_blocks > 0 and kv_cache_config.kv_cache_tensors: + all_groups_selected = len(selected_groups) == len(kv_cache_config.kv_cache_groups) + if ( + all_groups_selected + and kv_cache_config.num_blocks > 0 + and kv_cache_config.kv_cache_tensors + ): # Every KVCacheTensor describes placement within the same backing allocation, # so its size is the total, not a per-tensor share. total_gpu_kv_bytes = kv_cache_config.kv_cache_tensors[0].size worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks + elif kv_cache_config.num_blocks > 0: + worker_kv_bytes_per_block = sum( + group.kv_cache_spec.page_size_bytes for _, group in selected_groups + ) single_group_spec = ( kv_cache_config.kv_cache_groups[0].kv_cache_spec diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index a3edb78750..86e3510939 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -152,8 +152,9 @@ def resolve_mamba_align_size( such groups agree on the same value. """ mamba_align_size: int | None = None - for idx, tokens_per_block in enumerate(spec.tokens_per_block): - kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec + for group in spec.config.groups: + tokens_per_block = group.tokens_per_block + kv_spec = kv_cache_config.kv_cache_groups[group.group_id].kv_cache_spec if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode in ( "align", "all", @@ -188,8 +189,9 @@ class SchedulerOffloadConfig(NamedTuple): # each segment can never serve a load hit. Relevant for hybrid # architectures like DeepSeek V4 (MLA + SWA groups). full_attn_tokens_per_chunk: set[int] = set() - for idx, tokens_per_block in enumerate(spec.tokens_per_block): - kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec + for group in spec.config.groups: + tokens_per_block = group.tokens_per_block + kv_spec = kv_cache_config.kv_cache_groups[group.group_id].kv_cache_spec sw = get_sliding_window_size_in_chunks( kv_spec, tokens_per_block * spec.blocks_per_chunk ) @@ -214,9 +216,9 @@ class SchedulerOffloadConfig(NamedTuple): retention_interval = vllm_config.cache_config.prefix_cache_retention_interval eagle_groups = { - idx - for idx, g in enumerate(kv_cache_config.kv_cache_groups) - if g.is_eagle_group + group.group_id + for group in spec.config.groups + if kv_cache_config.kv_cache_groups[group.group_id].is_eagle_group } use_eagle_block_drop = ( @@ -249,8 +251,10 @@ class SchedulerOffloadConfig(NamedTuple): ) kv_group_configs_list: list[GroupOffloadConfig] = [] - for idx, tokens_per_block in enumerate(spec.tokens_per_block): - kv_cache_group = kv_cache_config.kv_cache_groups[idx] + for group in spec.config.groups: + group_id = group.group_id + tokens_per_block = group.tokens_per_block + kv_cache_group = kv_cache_config.kv_cache_groups[group_id] kv_spec = kv_cache_group.kv_cache_spec sw = get_sliding_window_size_in_chunks( kv_spec, tokens_per_block * spec.blocks_per_chunk @@ -261,7 +265,7 @@ class SchedulerOffloadConfig(NamedTuple): ) kv_group_configs_list.append( GroupOffloadConfig( - group_idx=idx, + group_idx=group_id, tokens_per_block=tokens_per_block, tokens_per_chunk=tokens_per_block * spec.blocks_per_chunk, hashes_per_chunk=( @@ -272,7 +276,7 @@ class SchedulerOffloadConfig(NamedTuple): kv_cache_spec=kv_spec, manager_cls=manager_cls, kv_event_group_spec=get_offloading_event_group_spec(kv_cache_group), - is_eagle_group=idx in eagle_groups, + is_eagle_group=group_id in eagle_groups, requires_cow_source=( isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode == "align" @@ -408,9 +412,10 @@ class RequestOffloadState: if new_block_id_groups is None: return - assert len(new_block_id_groups) == len(self.group_states) - for group_state, new_blocks in zip(self.group_states, new_block_id_groups): - group_state.block_ids.extend(new_blocks) + for group_config, group_state in zip( + self.config.kv_group_configs, self.group_states + ): + group_state.block_ids.extend(new_block_id_groups[group_config.group_idx]) def storable_chunks( self, @@ -541,11 +546,11 @@ class OffloadingConnectorScheduler: full_attention_groups: list[int] = [] sliding_window_groups: list[int] = [] - for group_config in self.config.kv_group_configs: + for config_idx, group_config in enumerate(self.config.kv_group_configs): if group_config.sliding_window_size_in_chunks is None: - full_attention_groups.append(group_config.group_idx) + full_attention_groups.append(config_idx) else: - sliding_window_groups.append(group_config.group_idx) + sliding_window_groups.append(config_idx) # sort sliding window groups by window size in decreasing order def _sliding_window_sort_key(i: int) -> int: @@ -737,7 +742,11 @@ class OffloadingConnectorScheduler: req_status.req_context, ) - def _lookup_complete_chunks(self, req_status: RequestOffloadState) -> int | None: + def _lookup_complete_chunks( + self, + req_status: RequestOffloadState, + max_num_new_tokens: int | None = None, + ) -> int | None: """ Find how many tokens beyond num_locally_computed_tokens can be loaded. @@ -748,6 +757,10 @@ class OffloadingConnectorScheduler: """ num_computed_tokens = req_status.num_locally_computed_tokens max_hit_size_tokens: int = req_status.req.num_tokens + if max_num_new_tokens is not None: + max_hit_size_tokens = min( + max_hit_size_tokens, num_computed_tokens + max_num_new_tokens + ) if self._sliding_window_groups: # the last prompt token has to be recomputed to get the logprobs # for sliding window attention, we must reduce by 1 to make sure @@ -929,8 +942,12 @@ class OffloadingConnectorScheduler: hash_idx = boundary_tokens // self.config.tokens_per_hash - 1 return make_offload_key(request.block_hashes[hash_idx], group_idx) - def _lookup(self, req_status: RequestOffloadState) -> int | None: - complete_hit = self._lookup_complete_chunks(req_status) + def _lookup( + self, + req_status: RequestOffloadState, + max_num_new_tokens: int | None = None, + ) -> int | None: + complete_hit = self._lookup_complete_chunks(req_status, max_num_new_tokens) req_status.partial_tail_boundary = None if complete_hit is None or not self.config.supports_partial_tail: return complete_hit @@ -939,9 +956,10 @@ class OffloadingConnectorScheduler: complete_boundary = local_tokens + complete_hit tokens_per_hash = self.config.tokens_per_hash block_end = complete_boundary + self._partial_tail_block_size - max_boundary = round_down( - min(req_status.req.num_prompt_tokens - 1, block_end - 1), tokens_per_hash - ) + max_boundary = min(req_status.req.num_prompt_tokens - 1, block_end - 1) + if max_num_new_tokens is not None: + max_boundary = min(max_boundary, local_tokens + max_num_new_tokens) + max_boundary = round_down(max_boundary, tokens_per_hash) if max_boundary <= complete_boundary: return complete_hit @@ -990,7 +1008,10 @@ class OffloadingConnectorScheduler: self._req_status[request.request_id] = req_status def get_num_new_matched_tokens( - self, request: Request, num_computed_tokens: int + self, + request: Request, + num_computed_tokens: int, + max_num_new_tokens: int | None = None, ) -> tuple[int | None, bool]: """ Get number of new tokens that can be loaded beyond the @@ -1030,7 +1051,7 @@ class OffloadingConnectorScheduler: num_hit_tokens = 0 else: lookup_start = time.monotonic() - num_hit_tokens = self._lookup(req_status) + num_hit_tokens = self._lookup(req_status, max_num_new_tokens) self._connector_stats.observe_histogram( _ConnectorMetricName.LOOKUP_SYNC_DELAY, time.monotonic() - lookup_start, @@ -1065,11 +1086,11 @@ class OffloadingConnectorScheduler: # per group group_sizes: list[int] = [] block_indices: list[int] = [] - for group_config, group_state, group_blocks in zip( + for group_config, group_state in zip( self.config.kv_group_configs, req_status.group_states, - blocks.blocks, ): + group_blocks = blocks.blocks[group_config.group_idx] self._current_batch_allocated_block_ids.update( block.block_id for block in group_blocks if block.block_id != 0 ) @@ -1184,7 +1205,8 @@ class OffloadingConnectorScheduler: for grp_idx in self._sliding_window_groups ) req_status.update_block_id_groups(new_block_id_groups) - for new_blocks in new_block_id_groups: + for group_config in self.config.kv_group_configs: + new_blocks = new_block_id_groups[group_config.group_idx] for bid in new_blocks: if bid != 0: self._current_batch_allocated_block_ids.add(bid) @@ -1218,6 +1240,10 @@ class OffloadingConnectorScheduler: ) -> dict[int, TransferJob]: store_jobs: dict[int, TransferJob] = {} num_groups = len(self.config.kv_group_configs) + config_idx_by_group = { + config.group_idx: idx + for idx, config in enumerate(self.config.kv_group_configs) + } for req_id, entries in handoffs.items(): req_status = self._req_status.get(req_id) if req_status is None: @@ -1225,7 +1251,10 @@ class OffloadingConnectorScheduler: req = req_status.req max_boundary = self._calc_num_offloadable_tokens(req_status, req.num_tokens) for group_idx, block_id, boundary in entries: - group_config = self.config.kv_group_configs[group_idx] + config_idx = config_idx_by_group.get(group_idx) + if config_idx is None: + continue + group_config = self.config.kv_group_configs[config_idx] if ( block_id == 0 or boundary > max_boundary @@ -1254,9 +1283,11 @@ class OffloadingConnectorScheduler: fenced_block_ids=[block_id], ) group_sizes = [0] * num_groups - group_sizes[group_idx] = 1 + group_sizes[config_idx] = 1 block_indices = [0] * num_groups - block_indices[group_idx] = boundary // group_config.tokens_per_block - 1 + block_indices[config_idx] = ( + boundary // group_config.tokens_per_block - 1 + ) store_jobs[job_id] = TransferJob( req_id=req_id, src_spec=GPULoadStoreSpec( @@ -1300,6 +1331,12 @@ class OffloadingConnectorScheduler: assert len(boundaries) == 1 boundary = boundaries.pop() req = req_status.req + group_states = { + group.group_idx: state + for group, state in zip( + self.config.kv_group_configs, req_status.group_states + ) + } max_boundary = min( req.num_prompt_tokens, req_status.max_offload_tokens or req.num_prompt_tokens, @@ -1315,7 +1352,7 @@ class OffloadingConnectorScheduler: block_idx = boundary // self._partial_tail_block_size if any( group.group_idx not in self._cow_source_groups - and block_idx >= len(req_status.group_states[group.group_idx].block_ids) + and block_idx >= len(group_states[group.group_idx].block_ids) for group in self.config.kv_group_configs ): continue @@ -1326,7 +1363,7 @@ class OffloadingConnectorScheduler: block_ids = [ cow_blocks[group.group_idx] if group.group_idx in self._cow_source_groups - else req_status.group_states[group.group_idx].block_ids[block_idx] + else group_states[group.group_idx].block_ids[block_idx] for group in self.config.kv_group_configs ] assert all(block_id != 0 for block_id in block_ids) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index ecdae9c4f2..14e826e552 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -14,6 +14,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import ( OffloadingWorkerMetadata, ReqId, ) +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( + get_offloading_group_ids, +) from vllm.logger import init_logger from vllm.v1.kv_cache_interface import ( AttentionSpec, @@ -65,7 +68,10 @@ class OffloadingConnectorWorker: def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): kv_cache_config = self.kv_cache_config - num_blocks = kv_cache_config.num_blocks + selected_group_ids = get_offloading_group_ids(kv_cache_config) + selected_groups = tuple( + kv_cache_config.kv_cache_groups[group_id] for group_id in selected_group_ids + ) mappings = derive_canonical_mappings( self.vllm_config, kv_cache_config, kv_caches ) @@ -76,7 +82,9 @@ class OffloadingConnectorWorker: unpadded_page_size_bytes: dict[str, int] = {} # layer_name -> size of page in bytes page_size_bytes: dict[str, int] = {} - for kv_cache_group in kv_cache_config.kv_cache_groups: + for kv_cache_group in selected_groups: + assert not kv_cache_group.host_resident + num_blocks = kv_cache_config.num_blocks group_layer_names = kv_cache_group.layer_names group_kv_cache_spec = kv_cache_group.kv_cache_spec if isinstance(group_kv_cache_spec, UniformTypeKVCacheSpecs): @@ -131,8 +139,11 @@ class OffloadingConnectorWorker: ), None, ) - if packed_layer_name is not None: + if packed_layer_name is not None and len(selected_groups) == len( + kv_cache_config.kv_cache_groups + ): (tensor,) = tensors_per_block[packed_layer_name] + num_blocks = tensor.shape[0] block_stride = tensor.stride(0) packed_tensor = tensor.as_strided( (num_blocks, block_stride), @@ -142,10 +153,7 @@ class OffloadingConnectorWorker: self._init_worker( CanonicalKVCaches( [CanonicalKVCacheTensor(packed_tensor, block_stride)], - [ - [CanonicalKVCacheRef(0, block_stride)] - for _ in kv_cache_config.kv_cache_groups - ], + [[CanonicalKVCacheRef(0, block_stride)] for _ in selected_groups], ) ) return @@ -192,7 +200,7 @@ class OffloadingConnectorWorker: ) group_data_refs: list[list[CanonicalKVCacheRef]] = [] - for kv_cache_group in kv_cache_config.kv_cache_groups: + for kv_cache_group in selected_groups: group_refs: list[CanonicalKVCacheRef] = [] for layer_name in kv_cache_group.layer_names: group_refs += block_data_refs[layer_name] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index a64f97805e..7357088175 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable +from functools import cached_property from typing import Any import torch @@ -25,6 +26,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import ( ) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( build_offloading_config, + get_offloading_group_ids, ) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, @@ -38,6 +40,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import ( ) from vllm.forward_context import ForwardContext from vllm.v1.attention.backend import AttentionMetadata +from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig @@ -47,6 +50,25 @@ from vllm.v1.request import Request class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): + @cached_property + def _bounding_group_ids(self) -> tuple[int, ...]: + """Prefix-cacheable groups this connector does not offload. + + Never offer tokens for a range some group cannot cover: past such a + group's own cached prefix its KV would be left unwritten. + """ + offloaded = set(get_offloading_group_ids(self._kv_cache_config)) + return tuple( + group_id + for group_id, group in enumerate(self._kv_cache_config.kv_cache_groups) + if group.kv_cache_spec.prefix_cacheable and group_id not in offloaded + ) + + @property + def scheduler(self) -> OffloadingConnectorScheduler: + assert self.connector_scheduler is not None + return self.connector_scheduler + @property def requires_kv_delivery(self) -> bool: # Runs as kv_both, but is a best-effort cache: a dropped save is just a @@ -138,9 +160,30 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): ) -> tuple[int | None, bool]: assert self.connector_scheduler is not None return self.connector_scheduler.get_num_new_matched_tokens( - request, num_computed_tokens + request, + num_computed_tokens, + max_num_new_tokens=self._max_loadable_tokens(request, num_computed_tokens), ) + def _max_loadable_tokens( + self, request: "Request", num_computed_tokens: int + ) -> int | None: + """How far past ``num_computed_tokens`` a load may reach, if bounded. + + Bounded by the deepest prefix the groups this connector does not + offload already hold, since nothing refills them beyond it. + """ + if not self._bounding_group_ids: + return None + assert self._kv_cache_manager is not None + coordinator = self._kv_cache_manager.coordinator + assert isinstance(coordinator, HybridKVCacheCoordinator) + _, per_group_hits = coordinator.find_longest_cache_hit_per_group( + request.block_hashes, request.num_tokens - 1 + ) + bound = min(per_group_hits[group_id] for group_id in self._bounding_group_ids) + return max(0, bound - num_computed_tokens) + def update_state_after_alloc( self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int ): @@ -185,6 +228,8 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): @classmethod def get_required_kvcache_layout(cls, vllm_config: VllmConfig) -> str | None: + if vllm_config.attention_config.hisparse_config is not None: + return "BLHNC" return "LBHNC" def reset_cache(self) -> bool | None: diff --git a/vllm/distributed/kv_transfer/kv_transfer_state.py b/vllm/distributed/kv_transfer/kv_transfer_state.py index 211d33fff2..690eb9fa23 100644 --- a/vllm/distributed/kv_transfer/kv_transfer_state.py +++ b/vllm/distributed/kv_transfer/kv_transfer_state.py @@ -78,13 +78,11 @@ def ensure_kv_transfer_initialized( global _KV_CONNECTOR_AGENT - if vllm_config.kv_transfer_config is None: + if _KV_CONNECTOR_AGENT is not None: return - if ( - vllm_config.kv_transfer_config.is_kv_transfer_instance - and _KV_CONNECTOR_AGENT is None - ): + kv_transfer_config = vllm_config.kv_transfer_config + if kv_transfer_config is not None and kv_transfer_config.is_kv_transfer_instance: # NIXL P/D requires an interleave_size equal to block_size. vllm_config.adjust_dcp_kv_cache_interleave_size(kv_cache_config) _sync_engine_id_across_tp(vllm_config) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index f9b86bd810..21596eb043 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -207,7 +207,7 @@ from collections.abc import Callable from dataclasses import dataclass, replace from enum import Enum from math import lcm -from typing import ClassVar, Generic, TypeVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast import numpy as np import torch @@ -284,6 +284,7 @@ from vllm.v1.attention.backend import ( CommonAttentionMetadata, MLAAttentionImpl, ) +from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.prefill import ( MLAPrefillBackend, get_mla_prefill_backend, @@ -309,6 +310,11 @@ from vllm.v1.kv_cache_interface import ( get_kv_quant_mode, ) +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroupBuilder, + ) + logger = init_logger(__name__) _FP8_DTYPE = current_platform.fp8_dtype() @@ -433,6 +439,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): use_sparse: bool = False, indexer: object | None = None, topk_indices_buffer: torch.Tensor | None = None, + index_group_builder: "SparseMLAIndexGroupBuilder | None" = None, non_causal_multi_token_decode: bool = False, sliding_window: int | None = None, prefill_backend_cls: type[MLAPrefillBackend] | None = None, @@ -538,9 +545,11 @@ class MLAAttention(nn.Module, AttentionLayerBase): # explicitly so backbone "skip" layers (indexer=None) still find it. if use_sparse: extra_impl_args["topk_indices_buffer"] = topk_indices_buffer + if index_group_builder is not None: + extra_impl_args["index_group_builder"] = index_group_builder impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls()) - self.impl = impl_cls( # type: ignore[assignment] # impl_cls always returns an MLAAttentionImpl subclass + impl = impl_cls( num_heads=self.num_heads, head_size=self.head_size, scale=self.scale, @@ -562,6 +571,13 @@ class MLAAttention(nn.Module, AttentionLayerBase): indexer=indexer, **extra_impl_args, ) + self.impl = impl # type: ignore[assignment] + index_group = getattr(impl, "index_group", None) + self.hisparse_cache = ( + index_group.cache(cast(Any, impl).index_group_index) + if isinstance(index_group, HiSparseMLAIndexGroup) + else None + ) self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None) self.is_amx_bmm_enabled = getattr(self.impl, "uses_amx_bmm", False) # MLA reads this weight directly to build W_UK/W_UV, so a backend must @@ -723,6 +739,61 @@ class MLAAttention(nn.Module, AttentionLayerBase): ) return self._chunked_prefill_workspace_size + def update_kv_cache( + self, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor | None, + attn_metadata: "MLACommonMetadata | None", + kv_cache_dtype: str, + k_scale: torch.Tensor, + ) -> None: + cache = self.hisparse_cache + if slot_mapping is None or (cache is not None and cache.dummy_batch): + return + kv_c_normed, k_pe, slot_mapping = maybe_gather_mla_latent_cache_inputs( + kv_c_normed, + k_pe, + slot_mapping, + attn_metadata.num_decode_tokens if attn_metadata is not None else None, + self.use_pcp, + ) + assert slot_mapping is not None + if cache is not None: + kv_cache, slot_mapping, num_rows = cache.write_target( + kv_c_normed.shape[0], slot_mapping.numel() + ) + kv_c_normed = kv_c_normed[:num_rows] + k_pe = k_pe[:num_rows] + self.impl.do_kv_cache_update( # type: ignore[attr-defined] + kv_c_normed, + k_pe, + kv_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + ) + if cache is not None: + mirror_target = cache.mirror_write_target(kv_c_normed.shape[0]) + if mirror_target is not None: + mirror_cache, mirror_slots = mirror_target + self.impl.do_kv_cache_update( # type: ignore[attr-defined] + kv_c_normed, + k_pe, + mirror_cache, + mirror_slots, + kv_cache_dtype, + k_scale, + ) + + def prepare_kv_cache_update( + self, attn_metadata: "MLACommonMetadata | None" + ) -> None: + cache = self.hisparse_cache + if cache is not None and cache.runtime.is_group_leader: + cache.prepare_group_for_batch(attn_metadata) + def forward( self, q: torch.Tensor, @@ -750,25 +821,18 @@ class MLAAttention(nn.Module, AttentionLayerBase): f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. " ) layer_slot_mapping = slot_mapping.get(self.layer_name) - kv_for_cache, kpe_for_cache, layer_slot_mapping = ( - maybe_gather_mla_latent_cache_inputs( - kv_c_normed, - k_pe, - layer_slot_mapping, - attn_metadata.num_decode_tokens - if attn_metadata is not None - else None, - self.use_pcp, - ) - ) - self.impl.do_kv_cache_update( # type: ignore[attr-defined] - kv_for_cache, - kpe_for_cache, + self.prepare_kv_cache_update(attn_metadata) + self.update_kv_cache( + kv_c_normed, + k_pe, self_kv_cache, layer_slot_mapping, + attn_metadata, self.kv_cache_dtype, self._k_scale, ) + if self.hisparse_cache is not None: + self.hisparse_cache.finish_kv_update() output = torch.empty(output_shape, dtype=q.dtype, device=q.device) self.forward_impl( q, @@ -1113,6 +1177,8 @@ class MLAAttention(nn.Module, AttentionLayerBase): return output_padded def _use_sparse_mha(self, attn_metadata: "MLACommonMetadata") -> bool: + if self.hisparse_cache is not None: + return False prefill = attn_metadata.prefill if prefill is None: return False @@ -1296,6 +1362,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): ) return MLAAttentionSpec( **common_kwargs, + is_index_group_leader=self.indexer is not None, non_causal_multi_token_decode=self.non_causal_multi_token_decode, ) @@ -1349,19 +1416,14 @@ def unified_mla_kv_cache_update( attn_metadata, attn_layer, kv_cache, layer_slot_mapping = get_attention_context( layer_name ) + attn_layer.prepare_kv_cache_update(attn_metadata) if layer_slot_mapping is not None: - kv_c_normed, k_pe, layer_slot_mapping = maybe_gather_mla_latent_cache_inputs( - kv_c_normed, - k_pe, - layer_slot_mapping, - attn_metadata.num_decode_tokens if attn_metadata is not None else None, - attn_layer.use_pcp, - ) - attn_layer.impl.do_kv_cache_update( # type: ignore[attr-defined] + attn_layer.update_kv_cache( kv_c_normed, k_pe, kv_cache, layer_slot_mapping, + attn_metadata, kv_cache_dtype, k_scale, ) @@ -1409,6 +1471,8 @@ def unified_mla_attention_with_output( del kv_cache_dummy_dep layer_name = _resolve_layer_name(layer_name) attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name) + if layer.hisparse_cache is not None: + layer.hisparse_cache.finish_kv_update() layer.forward_impl( q, kv_c_normed, @@ -1501,6 +1565,7 @@ class MLACommonBackend(AttentionBackend): def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": """Per-token-head modes pack an inline fp32 scale pair after the latent data (single-sided: ``head_size_v == 0`` for MLA).""" + spec = super().customize_spec(spec) mode = spec.kv_quant_mode if spec.state_content_bytes is not None or not mode.is_per_token_head: return spec @@ -2090,7 +2155,11 @@ def build_mla_chunked_context_metadata( local_token_offset += num_local_tokens seq_lens_cpu = _flat_int32(seq_lens_flat) - starts = _flat_int32(starts_flat).to(device, non_blocking=True) + starts_and_context_lens = _flat_int32(starts_flat + context_lens).to( + device, non_blocking=True + ) + starts = starts_and_context_lens[: len(starts_flat)] + context_lens_gpu = starts_and_context_lens[len(starts_flat) :] cu_seq_lens = _flat_int32(cu_seq_lens_flat).to(device, non_blocking=True) cu_seqlens_q = _flat_int32(cu_seqlens_q_flat).to(device, non_blocking=True) token_to_seq = _flat_int32(np.concatenate(token_to_seq_parts)).to( @@ -2145,7 +2214,7 @@ def build_mla_chunked_context_metadata( chunks.append(chunk) return MLACommonPrefillMetadata.ChunkedContextMetadata( - context_lens=context_lens_cpu.to(device, non_blocking=True), + context_lens=context_lens_gpu, workspace=chunked_prefill_workspace, chunks=chunks, context_lens_list=context_lens, @@ -2835,12 +2904,12 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]): output_lse = None for chunk in chunked_context.chunks: toks = chunk.num_context_tokens - block_table = prefill_metadata.block_table[chunk.request_slice] + chunk_block_table = prefill_metadata.block_table[chunk.request_slice] if self.kv_cache_dtype == "fp8_ds_mla": ops.cp_gather_and_upconvert_fp8_kv_cache( src_cache=kv_c_and_k_pe_cache, dst=workspace[:toks], - block_table=block_table, + block_table=chunk_block_table, workspace_starts=chunk.cu_seq_lens, batch_size=chunk.num_requests, seq_starts=chunk.starts, @@ -2853,7 +2922,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]): ops.gather_mla_context_cache_cpu( src_cache=kv_c_and_k_pe_cache, dst=workspace[:toks], - block_table=block_table, + block_table=chunk_block_table, starts=chunk.starts, cu_seq_lens=chunk.cu_seq_lens, ) @@ -2861,7 +2930,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]): ops.gather_and_maybe_dequant_cache( src_cache=kv_c_and_k_pe_cache, dst=workspace, - block_table=block_table, + block_table=chunk_block_table, cu_seq_lens=chunk.cu_seq_lens, token_to_seq=chunk.token_to_seq, num_tokens=toks, @@ -2874,7 +2943,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]): ops.cp_gather_cache( src_cache=kv_c_and_k_pe_cache, dst=workspace[:toks], - block_table=block_table, + block_table=chunk_block_table, cu_seq_lens=chunk.cu_seq_lens, batch_size=chunk.num_requests, seq_starts=chunk.starts, diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index fd29debbb7..5baddcd0a5 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -3,13 +3,15 @@ """Shared MHA implementation and metadata builder for sparse MLA backends.""" import math +from dataclasses import dataclass from shutil import which -from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar import numpy as np import torch from vllm import _custom_ops as ops +from vllm.config import get_current_vllm_config from vllm.distributed import ( get_dcp_group, get_tensor_model_parallel_world_size, @@ -17,6 +19,7 @@ from vllm.distributed import ( from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import ( MLACommonBaseImpl, + MLACommonDecodeMetadata, MLACommonMetadata, MLACommonPrefillMetadata, accumulate_mla_context_chunk, @@ -33,11 +36,19 @@ from vllm.utils.torch_utils import ( is_quantized_kv_cache, np_to_pinned_tensor, ) -from vllm.v1.attention.backend import AttentionMetadata, AttentionMetadataBuilder +from vllm.v1.attention.backend import AttentionMetadataBuilder from vllm.v1.attention.backends.fa_utils import get_flash_attn_version +from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroup, + SparseMLAIndexGroupBuilder, +) from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.attention.ops.dcp import MLADCPManager from vllm.v1.attention.ops.merge_attn_states import merge_attn_states +from vllm.v1.hisparse.runtime import ( + HiSparsePrefillStagingPlan, + build_hisparse_prefill_staging_plan, +) if TYPE_CHECKING: from vllm.config import VllmConfig @@ -47,8 +58,6 @@ if TYPE_CHECKING: logger = init_logger(__name__) -T = TypeVar("T", bound=AttentionMetadata) - GLOBAL_TOPK_MASK_MAX_BYTES = 128 * 1024 * 1024 # 128 MiB @@ -116,9 +125,45 @@ def _is_masked_mha_available( return fa_version == 4 and not is_quantized_kv_cache(kv_cache_dtype) +def _use_dense_mha_prefill( + vllm_config: "VllmConfig", prefill_max_seq_len: int, topk_tokens: int +) -> bool: + attention_config = vllm_config.attention_config + return ( + attention_config.hisparse_config is None + and not attention_config.sparse_mla_force_mqa + and prefill_max_seq_len <= topk_tokens + ) + + +@dataclass +class SparseMLAPrefillMetadata(MLACommonPrefillMetadata): + host_staging_plan: HiSparsePrefillStagingPlan | None = None + + +@dataclass(kw_only=True) +class SparseMLACommonMetadata(MLACommonMetadata[MLACommonDecodeMetadata]): + block_table: torch.Tensor + req_id_per_token: torch.Tensor + seq_lens: torch.Tensor | None = None + block_size: int = 64 + topk_tokens: int = 2048 + num_decodes: int = 0 + num_prefills: int = 0 + num_decode_tokens: int = 0 + decode_max_query_len: int = 0 + prefill_max_seq_len: int = 0 + prefill: SparseMLAPrefillMetadata | None = None + cp_kv_cache_interleave_size: int = 1 + + +T = TypeVar("T", bound=SparseMLACommonMetadata) + + class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): metadata_cls: type[T] require_uniform_decodes: ClassVar[bool] = False + hisparse_supports_multi_token_decode: ClassVar[bool] = False def __init__( self, @@ -138,6 +183,12 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): dtype=torch.int32, device=device, ) + self._context_lens_cpu = torch.empty( + vllm_config.scheduler_config.max_num_seqs, + dtype=torch.int32, + device="cpu", + pin_memory=PIN_MEMORY, + ) parallel_config = vllm_config.parallel_config self.use_pcp = parallel_config.prefill_context_parallel_size > 1 try: @@ -197,6 +248,22 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): layer_prefill_backend.clone() if layer_prefill_backend is not None else None ) + def _init_reorder_batch_threshold( + self, + reorder_batch_threshold: int | None = 1, + supports_spec_as_decode: bool = False, + supports_dcp_with_varlen: bool = False, + ) -> None: + if self.vllm_config.attention_config.hisparse_config is not None: + reorder_batch_threshold = 1 + if not self.hisparse_supports_multi_token_decode: + supports_spec_as_decode = False + super()._init_reorder_batch_threshold( + reorder_batch_threshold, + supports_spec_as_decode, + supports_dcp_with_varlen, + ) + @staticmethod def determine_chunked_prefill_workspace_size(vllm_config: "VllmConfig") -> int: scheduler_config = vllm_config.scheduler_config @@ -245,9 +312,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound assert seq_lens_cpu is not None - context_lens_cpu = torch.empty( - num_prefills, dtype=seq_lens_cpu.dtype, pin_memory=PIN_MEMORY - ) + context_lens_cpu = self._context_lens_cpu[:num_prefills] torch.subtract( seq_lens_cpu[num_decodes : num_decodes + num_prefills], prefill_query_lens_cpu, @@ -289,34 +354,66 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): prefill_max_query_len, prefill_query_lens_cpu, ) = self._build_prefill_fields(common_attn_metadata, num_decodes, num_prefills) + decode_max_query_len = 0 + if num_decodes: + if common_attn_metadata.max_logits_per_req is not None: + decode_max_query_len = common_attn_metadata.max_logits_per_req + elif num_prefills == 0: + decode_max_query_len = common_attn_metadata.max_query_len + else: + decode_query_lens_cpu = torch.diff( + common_attn_metadata.query_start_loc_cpu[: num_decodes + 1] + ) + decode_max_query_len = int(decode_query_lens_cpu.max().item()) prefill_max_seq_len = 0 - prefill: MLACommonPrefillMetadata | None = None + prefill: SparseMLAPrefillMetadata | None = None if num_prefills > 0 and self._prefill_backend is not None: seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound assert seq_lens_cpu is not None prefill_max_seq_len = int( seq_lens_cpu[num_decodes : num_decodes + num_prefills].max().item() ) - prefill = MLACommonPrefillMetadata( - block_table=common_attn_metadata.block_table_tensor[num_decodes:, ...], + block_table = common_attn_metadata.block_table_tensor[num_decodes:, ...] + chunked_context = self._build_chunked_context_fields( + common_attn_metadata, + num_decodes, + num_prefills, + prefill_query_lens_cpu, + ) + staging_plan = None + if self.vllm_config.attention_config.hisparse_config is not None: + prefill_seq_lens_cpu = seq_lens_cpu[ + num_decodes : num_decodes + num_prefills + ] + staging_block_capacity = int( + ( + (prefill_seq_lens_cpu + self.kv_cache_spec.block_size - 1) + // self.kv_cache_spec.block_size + ).sum() + ) + staging_plan = build_hisparse_prefill_staging_plan( + block_table, + common_attn_metadata.seq_lens[num_decodes:], + self.kv_cache_spec.block_size, + staging_block_capacity, + ) + if chunked_context is not None: + block_table = staging_plan.block_table + prefill = SparseMLAPrefillMetadata( + block_table=block_table, query_start_loc=prefill_query_start_loc, max_query_len=prefill_max_query_len, query_lens_cpu=prefill_query_lens_cpu, - chunked_context=self._build_chunked_context_fields( - common_attn_metadata, - num_decodes, - num_prefills, - prefill_query_lens_cpu, - ), + chunked_context=chunked_context, q_data_type=self.model_config.dtype, output_dtype=self.model_config.dtype, prefill_backend=self._prefill_backend, - use_dense_mha=( - prefill_max_seq_len <= self.topk_tokens - and not self.vllm_config.attention_config.sparse_mla_force_mqa + use_dense_mha=_use_dense_mha_prefill( + self.vllm_config, prefill_max_seq_len, self.topk_tokens ), topk_mask_workspace=self.topk_mask_workspace, + host_staging_plan=staging_plan, ) self._prefill_backend.prepare_metadata(prefill) @@ -335,6 +432,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): num_decodes=num_decodes, num_prefills=num_prefills, num_decode_tokens=num_decode_tokens, + decode_max_query_len=decode_max_query_len, prefill_max_seq_len=prefill_max_seq_len, prefill=prefill, cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, @@ -560,6 +658,7 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic kv_b_proj: "ColumnParallelLinear", indexer: object | None = None, topk_indices_buffer: torch.Tensor | None = None, + index_group_builder: SparseMLAIndexGroupBuilder | None = None, q_pad_num_heads: int | None = None, ) -> None: super().__init__( @@ -577,6 +676,21 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic ) self.init_topk_indices_buffer(indexer, topk_indices_buffer) + self.index_group: SparseMLAIndexGroup | None = None + self.index_group_index = 0 + if index_group_builder is None and self.topk_indices_buffer is not None: + index_group_builder = SparseMLAIndexGroupBuilder(self.topk_indices_buffer) + if index_group_builder is not None: + vllm_config = get_current_vllm_config() + self.index_group, self.index_group_index = ( + index_group_builder.register_layer( + indexer is not None, + vllm_config, + head_size=head_size, + kv_cache_dtype=kv_cache_dtype, + ) + ) + self._use_flashinfer_concat_mla_k = ( has_flashinfer() and which("ninja") is not None @@ -593,6 +707,33 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic kv_cache_dtype=kv_cache_dtype, ) + def record_logical_topk_ready(self) -> None: + if self.index_group is not None: + self.index_group.set_logical_topk_ready(self.index_group_index) + + def prepare_for_batch(self, attn_metadata: T | None) -> None: + if self.index_group is not None: + self.index_group.prepare_for_batch(self.index_group_index, attn_metadata) + + def _convert_logical_to_physical_topk( + self, + logical_topk_indices: torch.Tensor, + attn_metadata: Any, + *, + block_stride_rows: int | None, + return_valid_counts: bool, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + group = self.index_group + assert group is not None + assert self.dcp_world_size == 1 + return group.convert_logical_to_physical_topk( + self.index_group_index, + logical_topk_indices, + attn_metadata, + block_stride_rows=block_stride_rows, + return_valid_counts=return_valid_counts, + ) + @staticmethod def masked_mha_workspace_fits(prefill: MLACommonPrefillMetadata) -> bool: """Whether this prefill batch's top-k masks fit the workspace.""" @@ -853,8 +994,10 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic output: torch.Tensor, output_scale: torch.Tensor | None = None, ) -> None: - prefill_max_seq_len = attn_metadata.prefill_max_seq_len # type: ignore[attr-defined] - topk_tokens = attn_metadata.topk_tokens # type: ignore[attr-defined] + prefill_metadata = attn_metadata.prefill + assert prefill_metadata is not None + prefill_max_seq_len = attn_metadata.prefill_max_seq_len + topk_tokens = attn_metadata.topk_tokens force_dense = getattr(self, "_sparse_mla_force_dense_mha", False) force_masked = getattr(self, "_sparse_mla_force_masked_mha", False) if force_dense or (prefill_max_seq_len <= topk_tokens and not force_masked): @@ -863,7 +1006,7 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic kv_c_normed, k_pe, kv_c_and_k_pe_cache, - cast(MLACommonMetadata, attn_metadata), + attn_metadata, k_scale, output, output_scale, @@ -871,13 +1014,11 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic assert output_scale is None assert self.masked_mha_available - prefill_metadata = attn_metadata.prefill # type: ignore[attr-defined] - assert prefill_metadata is not None assert prefill_metadata.query_lens_cpu is not None assert self.topk_indices_buffer is not None q_lens = prefill_metadata.query_lens_cpu.tolist() - num_decode_tokens = attn_metadata.num_decode_tokens # type: ignore[attr-defined] + num_decode_tokens = attn_metadata.num_decode_tokens topk_all = self.topk_indices_buffer[ num_decode_tokens : num_decode_tokens + q.shape[0] ] diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py index 4211813062..9dd5ad6773 100644 --- a/vllm/model_executor/layers/mla.py +++ b/vllm/model_executor/layers/mla.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass +from typing import TYPE_CHECKING import torch @@ -11,6 +12,11 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.platforms import current_platform +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroupBuilder, + ) + @dataclass class MLAModules: @@ -30,6 +36,7 @@ class MLAModules: topk_indices_buffer: torch.Tensor | None indexer_rotary_emb: torch.nn.Module | None = None g_proj: torch.nn.Module | None = None + index_group_builder: "SparseMLAIndexGroupBuilder | None" = None # --8<-- [start:multi_head_latent_attention] @@ -128,6 +135,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer): use_sparse=self.is_sparse, indexer=self.indexer, topk_indices_buffer=mla_modules.topk_indices_buffer, + index_group_builder=mla_modules.index_group_builder, non_causal_multi_token_decode=non_causal_multi_token_decode, ) indexer_op = getattr(self.indexer, "indexer_op", None) @@ -223,6 +231,8 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer): if self.indexer and self.is_sparse and not self.skip_topk: self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb) + if self.is_sparse: + self.mla_attn.impl.record_logical_topk_ready() # type: ignore[attr-defined] if llama_4_scaling is not None: q *= llama_4_scaling diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index bb00a9cfe1..817f8ad361 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -425,7 +425,6 @@ def sparse_attn_indexer( quant_block_size, scale_fmt, ) - # The indexer and main MLA may classify the same short extend differently # because they use independent decode thresholds. Only the main MLA route # can determine whether the top-k indices will be consumed. @@ -955,6 +954,7 @@ class SparseAttnIndexer(CustomOp): self.dcp_rank, self.dcp_world_size, self.cp_kv_cache_interleave_size, + False, candidate_blocks=self.candidate_blocks, candidate_block_size=self.candidate_block_size, candidate_write=self.candidate_write, diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index bc99509c7f..bdc634102e 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -103,10 +103,14 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.backend import AttentionBackend, AttentionType +from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroupBuilder, + get_sparse_mla_index_group_max_rows, +) from vllm.v1.attention.backends.mla.indexer import ( DeepseekV32IndexerBackend, ) -from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec +from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec, SparseCacheRole from .interfaces import ( MixtureOfExperts, @@ -656,6 +660,7 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase): num_kv_heads=1, head_size=self.head_dim, dtype=self.dtype, + cache_role=SparseCacheRole.INDEXER, ) # Only has one vector instead of K + V def forward(self): ... @@ -1012,6 +1017,7 @@ class DeepseekV2MLAAttention(nn.Module): quant_config: QuantizationConfig | None = None, prefix: str = "", topk_indices_buffer: torch.Tensor | None = None, + index_group_builder: SparseMLAIndexGroupBuilder | None = None, input_size: int | None = None, reduce_results: bool = True, non_causal_multi_token_decode: bool = False, @@ -1199,6 +1205,7 @@ class DeepseekV2MLAAttention(nn.Module): indexer_rotary_emb=self.indexer_rope_emb, is_sparse=self.is_v32, topk_indices_buffer=topk_indices_buffer, + index_group_builder=index_group_builder, ) self.mla_attn = MultiHeadLatentAttentionWrapper( @@ -1242,6 +1249,7 @@ class DeepseekV2DecoderLayer(nn.Module): prefix: str, config: DeepseekV2Config | None = None, topk_indices_buffer: torch.Tensor | None = None, + index_group_builder: SparseMLAIndexGroupBuilder | None = None, ) -> None: super().__init__() @@ -1290,6 +1298,11 @@ class DeepseekV2DecoderLayer(nn.Module): and parallel_config.pipeline_parallel_size == 1 and is_moe_layer ) + attn_kwargs = ( + {"index_group_builder": index_group_builder} + if attn_cls is DeepseekV2MLAAttention + else {} + ) self.self_attn = attn_cls( vllm_config=vllm_config, config=config, @@ -1306,6 +1319,7 @@ class DeepseekV2DecoderLayer(nn.Module): prefix=f"{prefix}.self_attn", topk_indices_buffer=topk_indices_buffer, reduce_results=not self.use_sequence_parallel_moe, + **attn_kwargs, ) if is_moe_layer: @@ -1429,6 +1443,14 @@ class DeepseekV2Model(nn.Module): ) else: topk_indices_buffer = None + index_group_builder = ( + SparseMLAIndexGroupBuilder( + topk_indices_buffer, + get_sparse_mla_index_group_max_rows(vllm_config), + ) + if topk_indices_buffer is not None + else None + ) if get_pp_group().is_first_rank: self.embed_tokens = VocabParallelEmbedding( @@ -1445,6 +1467,7 @@ class DeepseekV2Model(nn.Module): vllm_config=vllm_config, prefix=prefix, topk_indices_buffer=topk_indices_buffer, + index_group_builder=index_group_builder, ), prefix=f"{prefix}.layers", ) diff --git a/vllm/model_executor/models/glm4_moe_lite.py b/vllm/model_executor/models/glm4_moe_lite.py index d0119719af..019380540d 100644 --- a/vllm/model_executor/models/glm4_moe_lite.py +++ b/vllm/model_executor/models/glm4_moe_lite.py @@ -67,6 +67,10 @@ from vllm.model_executor.models.glm4_moe import ( ) from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroupBuilder, + get_sparse_mla_index_group_max_rows, +) from .interfaces import SupportsLoRA, SupportsPP from .utils import ( @@ -108,6 +112,7 @@ class Glm4MoeLiteDecoderLayer(nn.Module): prefix: str, config: "Glm4MoeLiteConfig | None" = None, topk_indices_buffer: torch.Tensor | None = None, + index_group_builder: SparseMLAIndexGroupBuilder | None = None, ) -> None: super().__init__() @@ -136,6 +141,11 @@ class Glm4MoeLiteDecoderLayer(nn.Module): attn_cls = Glm4MoeLiteMLAAttention else: attn_cls = Glm4MoeLiteAttention + attn_kwargs: dict[str, typing.Any] = ( + {"index_group_builder": index_group_builder} + if attn_cls is Glm4MoeLiteMLAAttention + else {} + ) self.self_attn = attn_cls( vllm_config=vllm_config, @@ -152,6 +162,7 @@ class Glm4MoeLiteDecoderLayer(nn.Module): quant_config=quant_config, prefix=f"{prefix}.self_attn", topk_indices_buffer=topk_indices_buffer, + **attn_kwargs, ) if ( @@ -234,6 +245,14 @@ class Glm4MoeLiteModel(nn.Module): ) else: topk_indices_buffer = None + index_group_builder = ( + SparseMLAIndexGroupBuilder( + topk_indices_buffer, + get_sparse_mla_index_group_max_rows(vllm_config), + ) + if topk_indices_buffer is not None + else None + ) if get_pp_group().is_first_rank: self.embed_tokens = VocabParallelEmbedding( @@ -252,6 +271,7 @@ class Glm4MoeLiteModel(nn.Module): config=config, prefix=prefix, topk_indices_buffer=topk_indices_buffer, + index_group_builder=index_group_builder, ), prefix=f"{prefix}.layers", ) diff --git a/vllm/model_executor/models/glm4_moe_lite_mtp.py b/vllm/model_executor/models/glm4_moe_lite_mtp.py index 1d30a0eaf6..2ef9407fb7 100644 --- a/vllm/model_executor/models/glm4_moe_lite_mtp.py +++ b/vllm/model_executor/models/glm4_moe_lite_mtp.py @@ -51,6 +51,10 @@ from vllm.model_executor.model_loader.weight_utils import ( ) from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroupBuilder, + get_sparse_mla_index_group_max_rows, +) from .glm4_moe_lite import ( Glm4MixtureOfExperts, @@ -108,6 +112,14 @@ class Glm4MoeLiteMultiTokenPredictorLayer(nn.Module): ) else: topk_indices_buffer = None + index_group_builder = ( + SparseMLAIndexGroupBuilder( + topk_indices_buffer, + get_sparse_mla_index_group_max_rows(vllm_config), + ) + if topk_indices_buffer is not None + else None + ) self.shared_head = SharedHead( config=config, prefix=prefix, quant_config=quant_config @@ -117,6 +129,7 @@ class Glm4MoeLiteMultiTokenPredictorLayer(nn.Module): prefix=prefix, config=self.config, topk_indices_buffer=topk_indices_buffer, + index_group_builder=index_group_builder, ) def forward( diff --git a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py index 5209b94434..91ae7c44c2 100644 --- a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py +++ b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py @@ -47,6 +47,35 @@ _FLASHINFER_SM120_SPARSE_MLA_DECODE_LABELS = { _SPARSE_MLA_MIXED_WARMUP_TOKENS = 16 +def autotune_hisparse_flashinfer_attention(runner: "GPUModelRunner") -> None: + """Autotune each HiSparse FlashInfer sparse-MLA configuration.""" + from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( + FlashInferMLASparseImpl, + ) + + tuned: set[tuple[object, ...]] = set() + for layer in runner.vllm_config.compilation_config.static_forward_context.values(): + impl = getattr(layer, "impl", None) + if not isinstance(impl, FlashInferMLASparseImpl): + continue + if getattr(layer, "hisparse_cache", None) is None: + continue + if impl.topk_indices_buffer is None: + continue + key = ( + impl.kv_cache_dtype, + impl.num_heads, + impl.qk_nope_head_dim, + impl.qk_rope_head_dim, + impl.kv_lora_rank, + impl.topk_indices_buffer.shape[1], + ) + if key in tuned: + continue + impl.autotune_hisparse_decode(layer) + tuned.add(key) + + def _attention_backend_name(backend: object) -> str | None: get_name = getattr(backend, "get_name", None) if get_name is None: diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index c9c740f368..2991dc1027 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -25,6 +25,7 @@ from vllm.model_executor.warmup.flashinfer_autotune_cache import ( write_flashinfer_autotune_cache, ) from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( + autotune_hisparse_flashinfer_attention, deepseek_v4_sparse_mla_attention_warmup, flashinfer_sparse_mla_decode_autotune_warmup, ) @@ -353,9 +354,12 @@ def _flashinfer_autotune_token_counts(runner: "GPUModelRunner") -> tuple[int, .. return tuple(dict.fromkeys(token_counts)) -def _run_flashinfer_autotune_dummy_runs(runner: "GPUModelRunner") -> None: +def _run_flashinfer_autotune_dummy_runs( + runner: "GPUModelRunner", *, skip_attn: bool = False +) -> None: import vllm.utils.flashinfer as fi_utils + dummy_run_kwargs = {"skip_attn": True} if skip_attn else {} for num_tokens in _flashinfer_autotune_token_counts(runner): tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens) logger.info( @@ -369,6 +373,7 @@ def _run_flashinfer_autotune_dummy_runs(runner: "GPUModelRunner") -> None: skip_eplb=True, is_profile=True, randomize_inputs=True, + **dummy_run_kwargs, ) @@ -431,7 +436,14 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: torch.inference_mode(), fi_utils.autotune(tune_mode=True, **autotune_kwargs), ): - _run_flashinfer_autotune_dummy_runs(runner) + hisparse_enabled = ( + runner.vllm_config.attention_config.hisparse_config is not None + ) + if hisparse_enabled: + # HiSparse hot-buffer attention is bounded by decode batch + # size, not the prefill-sized batch used for the full model. + autotune_hisparse_flashinfer_attention(runner) + _run_flashinfer_autotune_dummy_runs(runner, skip_attn=hisparse_enabled) replayssm_autotune_warmup(runner) _autotune_kimi_k3_kda_qkvg(runner.get_model()) finally: diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py index 6a0d207ff5..b0fe554e7e 100644 --- a/vllm/models/deepseek_v32/attention.py +++ b/vllm/models/deepseek_v32/attention.py @@ -42,6 +42,9 @@ from vllm.v1.attention.ops.pcp import ( if TYPE_CHECKING: from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata + from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroupBuilder, + ) class DeepseekV32Indexer(nn.Module): @@ -126,6 +129,7 @@ class DeepseekV32Attention(MLAAttention): prefix: str, topk_indices_buffer: torch.Tensor | None = None, attn_backend: "type | None" = None, + index_group_builder: "SparseMLAIndexGroupBuilder | None" = None, ) -> None: quant_config = vllm_config.quant_config cache_config = vllm_config.cache_config @@ -209,6 +213,7 @@ class DeepseekV32Attention(MLAAttention): use_sparse=True, indexer=indexer, topk_indices_buffer=topk_indices_buffer, + index_group_builder=index_group_builder, attn_backend=attn_backend, ) @@ -304,6 +309,14 @@ class DeepseekV32Attention(MLAAttention): slot_mapping = forward_context.slot_mapping assert isinstance(slot_mapping, dict) mla_slot = slot_mapping.get(self.layer_name) + indexer_slot = ( + slot_mapping.get(self.indexer.k_cache.prefix) + if self.indexer is not None + else None + ) + hisparse_cache = self.hisparse_cache + layer_attn_metadata, _, _, _ = get_attention_context(self.layer_name) + self.impl.prepare_for_batch(layer_attn_metadata) if self.indexer is not None and not self.skip_topk: has_indexer = True @@ -331,12 +344,17 @@ class DeepseekV32Attention(MLAAttention): mla_k_scale = None indexer_k_cache = None mla_slot = None + indexer_slot = None else: - mla_kv_cache = self.kv_cache + mla_kv_cache = None if hisparse_cache is not None else self.kv_cache mla_k_scale = self._k_scale - kv_c_out = torch.empty_like(kv_c) - k_pe_out = torch.empty_like(k_pe) + if self.use_pcp or hisparse_cache is None: + kv_c_out = torch.empty_like(kv_c) + k_pe_out = torch.empty_like(k_pe) + else: + kv_c_out = torch.empty_like(kv_c) + k_pe_out = torch.empty_like(k_pe) q_c = fused_norm_rope( positions, q_c, @@ -354,6 +372,7 @@ class DeepseekV32Attention(MLAAttention): indexer_k_rope_cos_sin_cache, self.topk_indices_buffer, slot_mapping=mla_slot, + indexer_slot_mapping=indexer_slot, indexer_k_cache=indexer_k_cache, mla_kv_cache=mla_kv_cache, mla_kv_cache_dtype=self.kv_cache_dtype, @@ -365,6 +384,18 @@ class DeepseekV32Attention(MLAAttention): index_k_out=index_k_out, ) + if hisparse_cache is not None and mla_slot is not None: + assert kv_c_out is not None and k_pe_out is not None + self.update_kv_cache( + kv_c_out, + k_pe_out, + self.kv_cache, + mla_slot, + layer_attn_metadata, + self.kv_cache_dtype, + self._k_scale, + ) + q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) ql_nope = torch.bmm(q_nope.transpose(0, 1), self.W_UK_T).transpose(0, 1) @@ -459,6 +490,7 @@ class DeepseekV32Attention(MLAAttention): ), skip_topk_buffer_clear=True, ) + self.impl.record_logical_topk_ready() # type: ignore[attr-defined] attn_metadata, _, kv_cache, layer_slot_mapping = get_attention_context( self.layer_name @@ -467,6 +499,7 @@ class DeepseekV32Attention(MLAAttention): output.zero_() return attn_metadata = cast("MLACommonMetadata", attn_metadata) + self.impl.prepare_for_batch(attn_metadata) if self.use_pcp: assert kv_c is not None and k_pe is not None diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index 1c91d30157..ef7fcde5d3 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -68,7 +68,7 @@ def _fp8_quant_and_cache_write( kv_cache_ptr, kv_cache_scale_ptr, cache_block_size, - cache_stride, + cache_block_stride, offsets, HEAD_DIM: tl.constexpr, ): @@ -76,7 +76,7 @@ def _fp8_quant_and_cache_write( block_idx = slot_idx // cache_block_size block_offset = slot_idx % cache_block_size - block_start = block_idx * cache_block_size * cache_stride + block_start = block_idx * cache_block_stride tl.store( kv_cache_ptr + block_start + block_offset * HEAD_DIM + offsets, @@ -129,8 +129,9 @@ def _fused_norm_rope_kernel( index_k_out_ptr, index_k_out_stride, INDEX_K_HALF_ROT_DIM: tl.constexpr, - # Cache params (shared by indexer K and MLA) + # Cache params slot_mapping_ptr, + indexer_slot_mapping_ptr, # Index K FP8 cache indexer_cache_ptr, indexer_cache_scale_ptr, @@ -400,8 +401,8 @@ def _fused_norm_rope_kernel( ) # PCP inserts index K after gathering; other paths write it directly. - if indexer_cache_ptr is not None and slot_mapping_ptr is not None: - slot_idx = tl.load(slot_mapping_ptr + tok_idx) + if indexer_cache_ptr is not None and indexer_slot_mapping_ptr is not None: + slot_idx = tl.load(indexer_slot_mapping_ptr + tok_idx) _fp8_quant_and_cache_write( result, index_k_mask, @@ -431,8 +432,9 @@ def fused_norm_rope( index_k_layer_norm_eps: float, index_k_rope_cos_sin_cache: torch.Tensor | None, topk_indices_buffer: torch.Tensor, - # Cache params for fused writes (single slot_mapping for both caches) + # Cache params for fused writes slot_mapping: torch.Tensor | None = None, + indexer_slot_mapping: torch.Tensor | None = None, indexer_k_cache: torch.Tensor | None = None, mla_kv_cache: torch.Tensor | None = None, mla_kv_cache_dtype: str = "auto", @@ -476,9 +478,11 @@ def fused_norm_rope( # --- Indexer K cache setup --- if indexer_k_cache is not None: assert slot_mapping is not None + if indexer_slot_mapping is None: + indexer_slot_mapping = slot_mapping idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32) idx_cache_block_size = indexer_k_cache.shape[1] - idx_cache_stride = indexer_k_cache.shape[2] + idx_cache_stride = indexer_k_cache.stride(0) if indexer_k_cache.dtype == torch.uint8: indexer_k_cache = indexer_k_cache.view(torch.float8_e4m3fn) else: @@ -578,6 +582,7 @@ def fused_norm_rope( index_k_rope_cos_sin_cache.shape[-1] // 2, # Cache params slot_mapping, + indexer_slot_mapping, indexer_k_cache, idx_cache_scale_view, idx_cache_block_size, diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index 92b74e536f..b42ba28c8d 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -46,6 +46,10 @@ from vllm.models.common.ops.sequence_parallel import ( ) from vllm.models.deepseek_v32.attention import DeepseekV32Attention from vllm.sequence import IntermediateTensors +from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroupBuilder, + get_sparse_mla_index_group_max_rows, +) from .glm52_low_latency_gemm import enable_glm52_low_latency_gemm @@ -57,6 +61,7 @@ class DeepseekV32DecoderLayer(torch.nn.Module): prefix: str, config=None, topk_indices_buffer: torch.Tensor | None = None, + index_group_builder: SparseMLAIndexGroupBuilder | None = None, ) -> None: super().__init__() @@ -80,6 +85,7 @@ class DeepseekV32DecoderLayer(torch.nn.Module): config=config, prefix=f"{prefix}.self_attn", topk_indices_buffer=topk_indices_buffer, + index_group_builder=index_group_builder, ) if ( @@ -189,6 +195,10 @@ class DeepseekV32Model(torch.nn.Module): dtype=torch.int32, device=self.device, ) + index_group_builder = SparseMLAIndexGroupBuilder( + self.topk_indices_buffer, + get_sparse_mla_index_group_max_rows(vllm_config), + ) if get_pp_group().is_first_rank: self.embed_tokens = make_input_embedding( @@ -209,6 +219,7 @@ class DeepseekV32Model(torch.nn.Module): vllm_config=vllm_config, prefix=prefix, topk_indices_buffer=self.topk_indices_buffer, + index_group_builder=index_group_builder, ), prefix=f"{prefix}.layers", ) diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index a5348ed220..bdc6b95c15 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -48,6 +48,10 @@ from vllm.models.common.ops.sequence_parallel import ( from vllm.models.deepseek_v32.common.kernels import fused_eh_norm from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from vllm.v1.attention.backends.mla.index_group import ( + SparseMLAIndexGroupBuilder, + get_sparse_mla_index_group_max_rows, +) from .glm52_low_latency_gemm import ( build_glm52_plan, @@ -80,6 +84,10 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module): dtype=torch.int32, device=current_platform.device_type, ) + index_group_builder = SparseMLAIndexGroupBuilder( + topk_indices_buffer, + get_sparse_mla_index_group_max_rows(vllm_config), + ) self.shared_head = SharedHead( config=config, prefix=prefix, quant_config=quant_config ) @@ -88,6 +96,7 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module): prefix, config=config, topk_indices_buffer=topk_indices_buffer, + index_group_builder=index_group_builder, ) def forward( diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index f5a8db0c24..ac02bca1fc 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/vllm/models/minimax_m3/common/indexer.py @@ -55,6 +55,7 @@ from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheSpec, MLAAttentionSpec, + SparseCacheRole, ) logger = init_logger(__name__) @@ -148,6 +149,7 @@ class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase): num_kv_heads=1, head_size=self.head_dim, dtype=self.dtype, + cache_role=SparseCacheRole.INDEXER, ) def forward(self) -> None: ... diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index d75c79fe0f..f4031a7ee0 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -22,6 +22,7 @@ if TYPE_CHECKING: from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey from vllm.platforms.interface import DeviceCapability + from vllm.v1.hisparse.runtime import HiSparseCacheHandle from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheLayout, @@ -838,6 +839,9 @@ class AttentionImplBase(ABC, Generic[T]): def process_weights_after_loading(self, act_dtype: torch.dtype): pass + def prepare_for_batch(self, attn_metadata: T | None) -> None: + """Prepare implementation-specific state for the current batch.""" + class AttentionImpl(AttentionImplBase[T], Generic[T]): """Standard attention implementation with forward method.""" @@ -956,6 +960,7 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]): class MLAAttentionImpl(AttentionImplBase[T], Generic[T]): """MLA attention implementation with forward_mqa and forward_mha methods.""" + hisparse_cache: "HiSparseCacheHandle | None" = None supports_pcp: bool = True @abstractmethod diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index 2882a81b11..43720fb3ed 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -8,9 +8,9 @@ import torch from vllm.config import VllmConfig from vllm.config.cache import CacheDType -from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata from vllm.model_executor.layers.attention.sparse_mla_attention import ( SparseMLACommonImpl, + SparseMLACommonMetadata, SparseMLACommonMetadataBuilder, ) from vllm.platforms.interface import DeviceCapability @@ -18,11 +18,11 @@ from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, AttentionLayer, - AttentionMetadata, MLAAttentionImpl, MultipleOf, ) from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla +from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( flat_kv_row_view, triton_convert_req_index_to_global_index, @@ -106,26 +106,8 @@ class FlashAttnMLASparseBackend(AttentionBackend): @dataclass -class FlashAttnMLASparseMetadata(AttentionMetadata): - num_reqs: int - max_query_len: int - max_seq_len: int - - num_actual_tokens: int - query_start_loc: torch.Tensor - slot_mapping: torch.Tensor - - block_table: torch.Tensor - req_id_per_token: torch.Tensor - seq_lens: torch.Tensor - block_size: int = 64 - topk_tokens: int = 2048 - num_decodes: int = 0 - num_prefills: int = 0 - num_decode_tokens: int = 0 - prefill_max_seq_len: int = 0 - prefill: MLACommonPrefillMetadata | None = None - cp_kv_cache_interleave_size: int = 1 +class FlashAttnMLASparseMetadata(SparseMLACommonMetadata): + pass class FlashAttnMLASparseMetadataBuilder( @@ -214,6 +196,76 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]): assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_actual_toks] + index_group = self.index_group + if isinstance(index_group, HiSparseMLAIndexGroup): + num_decode_tokens = attn_metadata.num_decode_tokens + outputs = [] + if num_decode_tokens: + physical_topk, valid_counts = ( + index_group.convert_decode_logical_to_physical_topk( + self.index_group_index, + topk_indices[:num_decode_tokens], + attn_metadata, + return_valid_counts=True, + ) + ) + outputs.append( + self._run_mqa_kernel( + q_nope[:num_decode_tokens], + q_rope[:num_decode_tokens], + index_group.physical_kv_cache(self.index_group_index).view( + kv_c_and_k_pe_cache.dtype + ), + physical_topk, + valid_counts, + attn_metadata.block_size, + ) + ) + if num_decode_tokens < num_actual_toks: + cache = index_group.cache(self.index_group_index) + if num_decode_tokens == 0 and cache.all_context_pages_resident: + physical_topk, valid_counts = ( + index_group.convert_logical_to_physical_topk( + self.index_group_index, + topk_indices, + attn_metadata, + block_stride_rows=None, + return_valid_counts=True, + ) + ) + prefill_cache = index_group.physical_kv_cache( + self.index_group_index + ).view(kv_c_and_k_pe_cache.dtype) + else: + prefill_cache, block_table, req_ids = ( + index_group.stage_prefill_rows( + self.index_group_index, + kv_c_and_k_pe_cache, + attn_metadata, + ) + ) + physical_topk, valid_counts = ( + triton_convert_req_index_to_global_index( + req_ids, + block_table, + topk_indices[num_decode_tokens:], + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + return_valid_counts=True, + ) + ) + outputs.append( + self._run_mqa_kernel( + q_nope[num_decode_tokens:], + q_rope[num_decode_tokens:], + prefill_cache, + physical_topk, + valid_counts, + attn_metadata.block_size, + ) + ) + return torch.cat(outputs) if len(outputs) > 1 else outputs[0], None + kv_rows, block_stride_rows = flat_kv_row_view( kv_c_and_k_pe_cache, attn_metadata.block_size ) @@ -226,9 +278,36 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]): NUM_TOPK_TOKENS=topk_indices.shape[1], return_valid_counts=True, ) + return ( + self._run_mqa_kernel( + q_nope, + q_rope, + kv_rows, + topk_indices, + valid_counts, + attn_metadata.block_size, + cache_is_flat=True, + ), + None, + ) + + def _run_mqa_kernel( + self, + q_nope: torch.Tensor, + q_rope: torch.Tensor, + kv_cache: torch.Tensor, + topk_indices: torch.Tensor, + valid_counts: torch.Tensor, + block_size: int, + *, + cache_is_flat: bool = False, + ) -> torch.Tensor: + kv_rows = ( + kv_cache if cache_is_flat else flat_kv_row_view(kv_cache, block_size)[0] + ) cu_seqlens_q = torch.arange( - 0, num_actual_toks + 1, dtype=torch.int32, device=q_rope.device + 0, q_rope.shape[0] + 1, dtype=torch.int32, device=q_rope.device ) k_cache = kv_rows[:, self.kv_lora_rank :].unsqueeze(1).unsqueeze(1) v_cache = kv_rows[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1) @@ -247,4 +326,4 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]): causal=True, fa_version=3, ) - return out, None + return out diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index b34c983752..ae2e512840 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -11,22 +11,23 @@ from vllm import envs from vllm.config import VllmConfig from vllm.config.cache import CacheDType from vllm.logger import init_logger -from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata from vllm.model_executor.layers.attention.sparse_mla_attention import ( SparseMLACommonImpl, + SparseMLACommonMetadata, SparseMLACommonMetadataBuilder, ) +from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.utils.torch_utils import is_quantized_kv_cache from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, AttentionLayer, - AttentionMetadata, AttentionType, MLAAttentionImpl, MultipleOf, ) +from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( flat_kv_row_view, triton_convert_req_index_to_global_index, @@ -224,33 +225,9 @@ class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase): @dataclass -class FlashInferMLASparseMetadata(AttentionMetadata): +class FlashInferMLASparseMetadata(SparseMLACommonMetadata): """Attention metadata for FlashInfer MLA Sparse backend.""" - num_reqs: int - max_query_len: int - max_seq_len: int - num_actual_tokens: int - - # Query start locations - query_start_loc: torch.Tensor - slot_mapping: torch.Tensor - block_table: torch.Tensor - req_id_per_token: torch.Tensor - - # Sequence lengths for all requests (context + query) - seq_lens: torch.Tensor - num_decodes: int - num_prefills: int - num_decode_tokens: int - prefill_max_seq_len: int = 0 - prefill: MLACommonPrefillMetadata | None = None - - # Sparse-specific - block_size: int = 64 - topk_tokens: int = 2048 - cp_kv_cache_interleave_size: int = 1 - class FlashInferMLASparseMetadataBuilder( SparseMLACommonMetadataBuilder[FlashInferMLASparseMetadata] @@ -286,6 +263,7 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde """Metadata builder for the SM100 TRT-LLM sparse MLA kernel.""" _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS + hisparse_supports_multi_token_decode: ClassVar[bool] = True def _build_req_id_per_token( self, @@ -402,6 +380,78 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]): assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_actual_toks] + self._prepare_mqa_kernel(layer, q.device) + + index_group = self.index_group + if isinstance(index_group, HiSparseMLAIndexGroup): + num_decode_tokens = attn_metadata.num_decode_tokens + decode_out: torch.Tensor | None = None + decode_lse: torch.Tensor | None = None + if num_decode_tokens > 0: + physical_topk, valid_counts = ( + index_group.convert_decode_logical_to_physical_topk( + self.index_group_index, + topk_indices[:num_decode_tokens], + attn_metadata, + return_valid_counts=True, + ) + ) + decode_out, decode_lse = self._run_mqa_kernel( + q[:num_decode_tokens], + index_group.physical_kv_cache(self.index_group_index).view( + kv_c_and_k_pe_cache.dtype + ), + physical_topk, + valid_counts, + ) + if num_decode_tokens == num_actual_toks: + return decode_out, decode_lse + + cache = index_group.cache(self.index_group_index) + if num_decode_tokens == 0 and cache.all_context_pages_resident: + physical_topk, valid_counts = ( + index_group.convert_logical_to_physical_topk( + self.index_group_index, + topk_indices, + attn_metadata, + block_stride_rows=None, + return_valid_counts=True, + ) + ) + return self._run_mqa_kernel( + q, + index_group.physical_kv_cache(self.index_group_index).view( + kv_c_and_k_pe_cache.dtype + ), + physical_topk, + valid_counts, + ) + + prefill_cache, block_table, req_ids = index_group.stage_prefill_rows( + self.index_group_index, kv_c_and_k_pe_cache, attn_metadata + ) + prefill_indices, prefill_lens = triton_convert_req_index_to_global_index( + req_ids, + block_table, + topk_indices[num_decode_tokens:], + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + return_valid_counts=True, + ) + prefill_out, prefill_lse = self._run_mqa_kernel( + q[num_decode_tokens:], + prefill_cache, + prefill_indices, + prefill_lens, + ) + if decode_out is None: + return prefill_out, prefill_lse + output = torch.cat((decode_out, prefill_out)) + if decode_lse is None: + return output, None + assert prefill_lse is not None + return output, torch.cat((decode_lse, prefill_lse)) + _, block_stride_rows = flat_kv_row_view( kv_c_and_k_pe_cache, attn_metadata.block_size ) @@ -420,18 +470,27 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]): return_valid_counts=True, ) else: - topk_indices_physical, seq_lens = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[:num_actual_toks], - attn_metadata.block_table, + topk_indices_physical, seq_lens = self._convert_logical_to_physical_topk( topk_indices, - BLOCK_SIZE=attn_metadata.block_size, - BLOCK_STRIDE_ROWS=block_stride_rows, - NUM_TOPK_TOKENS=topk_indices.shape[1], + attn_metadata, + block_stride_rows=block_stride_rows, return_valid_counts=True, ) + return self._run_mqa_kernel( + q, + kv_c_and_k_pe_cache, + topk_indices_physical, + seq_lens, + ) + + def _prepare_mqa_kernel( + self, + layer: AttentionLayer, + device: torch.device, + ) -> None: if self._workspace_buffer is None: - self._workspace_buffer = _get_workspace_buffer(q.device) + self._workspace_buffer = _get_workspace_buffer(device) if self.bmm1_scale is None: self.bmm1_scale = self.scale @@ -442,15 +501,71 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]): if is_quantized_kv_cache(self.kv_cache_dtype): self.bmm2_scale *= layer._k_scale_float + def autotune_hisparse_decode(self, layer: AttentionLayer) -> None: + """Autotune the largest legal HiSparse decode batch.""" + assert isinstance(self.index_group, HiSparseMLAIndexGroup) + cache = self.index_group.cache(self.index_group_index) + assert self.topk_indices_buffer is not None + + runtime = cache.runtime + kv_cache = runtime.hot.attention_cache + num_tokens = runtime.max_num_reqs + topk_tokens = self.topk_indices_buffer.shape[1] + self._prepare_mqa_kernel(layer, kv_cache.device) + + q_dtype = ( + current_platform.fp8_dtype() + if is_quantized_kv_cache(self.kv_cache_dtype) + else kv_cache.dtype + ) + q = torch.zeros( + ( + num_tokens, + self.num_heads, + self.kv_lora_rank + self.qk_rope_head_dim, + ), + dtype=q_dtype, + device=kv_cache.device, + ) + topk_indices = ( + torch.arange( + topk_tokens, + dtype=torch.int32, + device=kv_cache.device, + ) + .expand(num_tokens, -1) + .contiguous() + ) + seq_lens = torch.full( + (num_tokens,), + topk_tokens, + dtype=torch.int32, + device=kv_cache.device, + ) + self._run_mqa_kernel(q, kv_cache, topk_indices, seq_lens) + + def _run_mqa_kernel( + self, + q: torch.Tensor, + kv_cache: torch.Tensor, + topk_indices: torch.Tensor, + seq_lens: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + assert self._workspace_buffer is not None + assert self.bmm1_scale is not None + assert self.bmm2_scale is not None + from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla + kv_cache = kv_cache.view(q.dtype) + # Single-token sparse decode. trtllm-gen requires the q_len_per_request # dim, but the sparse attention mask is fully per-token (each query token # carries its own top-k index row), so unsqueeze is sufficient and # correct. The MTP/multi-token q_len grouping is a perf-only layout and is # deferred until MTP is validated end-to-end for this backend. query = q.unsqueeze(1) - block_tables = topk_indices_physical.unsqueeze(1) + block_tables = topk_indices.unsqueeze(1) seq_lens_arg = seq_lens # page_table width = topk buffer width, which kpool widens past @@ -460,7 +575,7 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]): # count), so the -1 padding slots past seq_lens are never attended to. # Use the actual buffer width instead of the fixed topk_tokens, which # mismatches the page_table when index_kpool > 1. - sparse_topk_capacity = topk_indices_physical.shape[1] + sparse_topk_capacity = topk_indices.shape[1] extra_kwargs: dict[str, torch.Tensor] = {} empty_rows: torch.Tensor | None = None @@ -473,14 +588,12 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]): # valid indices into a contiguous prefix, which is what the kernel # requires of the page table. empty_rows = seq_lens == 0 - topk_indices_physical[:, 0] = topk_indices_physical[:, 0].masked_fill( - empty_rows, 0 - ) + topk_indices[:, 0] = topk_indices[:, 0].masked_fill(empty_rows, 0) extra_kwargs["sparse_mla_top_k_lens"] = seq_lens.clamp(min=1) kernel_out = trtllm_batch_decode_with_kv_cache_mla( query=query, - kv_cache=kv_c_and_k_pe_cache.unsqueeze(1), + kv_cache=kv_cache.unsqueeze(1), workspace_buffer=self._workspace_buffer, qk_nope_head_dim=self.qk_nope_head_dim, kv_lora_rank=self.kv_lora_rank, @@ -506,7 +619,7 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]): if lse is not None: lse = self._normalize_lse(lse, out.shape[0], out.shape[1]) if empty_rows is None and lse is not None: - empty_rows = (topk_indices_physical == -1).all(dim=-1) + empty_rows = (topk_indices == -1).all(dim=-1) if empty_rows is not None: out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0) if lse is not None: diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py index b7459203e0..e59c9ce46d 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py @@ -7,17 +7,14 @@ from typing import TYPE_CHECKING, cast import torch from vllm.model_executor.layers.attention.sparse_mla_attention import ( - SharedTopkIndicesBuffer, -) -from vllm.v1.attention.backend import ( - AttentionLayer, - AttentionType, - MLAAttentionImpl, + SparseMLACommonImpl, ) +from vllm.v1.attention.backend import AttentionLayer, AttentionType from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseMetadata, _get_workspace_buffer, ) +from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( triton_convert_req_index_to_global_index, ) @@ -32,9 +29,7 @@ def _kv_scale_format_for_model(model_type: str | None) -> str: return "pow2_fp32" -class FlashInferMLASparseSM120Impl( - MLAAttentionImpl[FlashInferMLASparseMetadata], SharedTopkIndicesBuffer -): +class FlashInferMLASparseSM120Impl(SparseMLACommonImpl[FlashInferMLASparseMetadata]): """SM120 FlashInfer sparse-MLA implementation.""" is_sparse = True @@ -65,20 +60,28 @@ class FlashInferMLASparseSM120Impl( "FLASHINFER_MLA_SPARSE_SM120 only supports decoder self-attention" ) - self.num_heads = num_heads - self.head_size = head_size - self.scale = float(scale) - self.num_kv_heads = num_kv_heads - self.kv_cache_dtype = kv_cache_dtype - if self.kv_cache_dtype != "fp8_ds_mla": + if kv_cache_dtype != "fp8_ds_mla": raise NotImplementedError( "FLASHINFER_MLA_SPARSE_SM120 requires the packed fp8_ds_mla " f"KV cache layout; got kv_cache_dtype={kv_cache_dtype!r}." ) - self.kv_lora_rank: int = mla_args["kv_lora_rank"] - self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"] - self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"] + topk_indices_buffer = mla_args.pop("topk_indices_buffer", None) + super().__init__( + num_heads, + head_size, + scale, + num_kv_heads, + alibi_slopes, + sliding_window, + kv_cache_dtype, + logits_soft_cap, + attn_type, + kv_sharing_target_layer_name, + indexer=indexer, + topk_indices_buffer=topk_indices_buffer, + **mla_args, + ) from vllm.config import get_current_vllm_config vllm_config = get_current_vllm_config() @@ -89,9 +92,6 @@ class FlashInferMLASparseSM120Impl( ) self.kv_scale_format = _kv_scale_format_for_model(model_type) - # Skip-topk layers are built with indexer=None and get the shared - # buffer via mla_args instead (cf. FLASHMLA_SPARSE). - self.init_topk_indices_buffer(indexer, mla_args.get("topk_indices_buffer")) from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120 if not has_flashinfer_sparse_mla_sm120(): @@ -118,6 +118,70 @@ class FlashInferMLASparseSM120Impl( assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_actual_toks] + index_group = self.index_group + if isinstance(index_group, HiSparseMLAIndexGroup): + num_decode_tokens = attn_metadata.num_decode_tokens + outputs = [] + if num_decode_tokens: + topk_indices_physical = cast( + torch.Tensor, + index_group.convert_decode_logical_to_physical_topk( + self.index_group_index, + topk_indices[:num_decode_tokens], + attn_metadata, + return_valid_counts=False, + ), + ) + outputs.append( + self._run_mqa_kernel( + q[:num_decode_tokens], + index_group.physical_kv_cache(self.index_group_index), + topk_indices_physical, + ) + ) + if num_decode_tokens < num_actual_toks: + cache = index_group.cache(self.index_group_index) + if num_decode_tokens == 0 and cache.all_context_pages_resident: + topk_indices_physical = cast( + torch.Tensor, + index_group.convert_logical_to_physical_topk( + self.index_group_index, + topk_indices, + attn_metadata, + block_stride_rows=None, + return_valid_counts=False, + ), + ) + prefill_cache = index_group.physical_kv_cache( + self.index_group_index + ) + else: + prefill_cache, block_table, req_ids = ( + index_group.stage_prefill_rows( + self.index_group_index, + kv_c_and_k_pe_cache, + attn_metadata, + ) + ) + topk_indices_physical = cast( + torch.Tensor, + triton_convert_req_index_to_global_index( + req_ids, + block_table, + topk_indices[num_decode_tokens:], + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + ), + ) + outputs.append( + self._run_mqa_kernel( + q[num_decode_tokens:], + prefill_cache, + topk_indices_physical, + ) + ) + output = torch.cat(outputs) if len(outputs) > 1 else outputs[0] + return output, None topk_indices_physical = cast( torch.Tensor, @@ -129,6 +193,22 @@ class FlashInferMLASparseSM120Impl( NUM_TOPK_TOKENS=topk_indices.shape[1], ), ) + return ( + self._run_mqa_kernel( + q, + kv_c_and_k_pe_cache, + topk_indices_physical, + ), + None, + ) + + def _run_mqa_kernel( + self, + q: torch.Tensor, + kv_cache: torch.Tensor, + topk_indices_physical: torch.Tensor, + ) -> torch.Tensor: + num_actual_toks = q.shape[0] output = q.new_empty( (num_actual_toks, self.num_heads, self.kv_lora_rank), @@ -142,20 +222,21 @@ class FlashInferMLASparseSM120Impl( flashinfer_trtllm_batch_decode_with_kv_cache_mla, ) + sparse_capacity = topk_indices_physical.shape[1] out = flashinfer_trtllm_batch_decode_with_kv_cache_mla( query=q.unsqueeze(1), - kv_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(1), + kv_cache=kv_cache.view(torch.uint8).unsqueeze(1), workspace_buffer=self._workspace_buffer, qk_nope_head_dim=self.qk_nope_head_dim, kv_lora_rank=self.kv_lora_rank, qk_rope_head_dim=self.qk_rope_head_dim, block_tables=topk_indices_physical.unsqueeze(1), seq_lens=None, - max_seq_len=attn_metadata.topk_tokens, + max_seq_len=sparse_capacity, out=output.unsqueeze(1), bmm1_scale=self.scale, bmm2_scale=1.0, - sparse_mla_top_k=attn_metadata.topk_tokens, + sparse_mla_top_k=sparse_capacity, kv_scale_format=self.kv_scale_format, ) - return out.squeeze(1), None + return out.squeeze(1) diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index b837cc366b..4ecf559683 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -9,23 +9,23 @@ from vllm import _custom_ops as ops from vllm.config import VllmConfig, get_current_vllm_config from vllm.config.cache import CacheDType from vllm.logger import init_logger -from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata from vllm.model_executor.layers.attention.sparse_mla_attention import ( SparseMLACommonImpl, + SparseMLACommonMetadata, SparseMLACommonMetadataBuilder, ) from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.utils.platform_utils import num_compute_units -from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.utils.torch_utils import current_stream, is_quantized_kv_cache from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, AttentionLayer, - AttentionMetadata, CommonAttentionMetadata, MultipleOf, ) +from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( flat_kv_row_view, triton_convert_req_index_to_global_index, @@ -172,28 +172,7 @@ class FlashMLASparseBackend(AttentionBackend): @dataclass -class FlashMLASparseMetadata(AttentionMetadata): - num_reqs: int - max_query_len: int - max_seq_len: int - - num_actual_tokens: int # Number of tokens excluding padding. - query_start_loc: torch.Tensor - slot_mapping: torch.Tensor - - block_table: torch.Tensor - req_id_per_token: torch.Tensor - block_size: int = 64 - topk_tokens: int = 2048 - - num_decodes: int = 0 - num_prefills: int = 0 - num_decode_tokens: int = 0 - seq_lens: torch.Tensor | None = None - prefill_max_seq_len: int = 0 - prefill: MLACommonPrefillMetadata | None = None - cp_kv_cache_interleave_size: int = 1 - +class FlashMLASparseMetadata(SparseMLACommonMetadata): @dataclass class FP8KernelMetadata: scheduler_metadata: FlashMLASchedMeta @@ -233,6 +212,7 @@ class FlashMLASparseMetadata(AttentionMetadata): req_start_idx: int workspace_starts: torch.Tensor chunk_tot_seqlen: int + seq_lens: torch.Tensor | None = None chunks: list[Chunk] @@ -263,6 +243,7 @@ class FlashMLASparseMetadataBuilder( ): _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH require_uniform_decodes: ClassVar[bool] = True + hisparse_supports_multi_token_decode: ClassVar[bool] = True metadata_cls = FlashMLASparseMetadata def __init__( @@ -283,6 +264,9 @@ class FlashMLASparseMetadataBuilder( ) else: threshold = {16: 128, 32: 128, 64: 256, 128: 256}.get(num_q_heads, 256) + self.use_hisparse = vllm_config.attention_config.hisparse_config is not None + if self.use_hisparse: + threshold = 1 # Varlen decodes are safe under DCP: causality comes from the # indexer's top-k indices, not from the kernel metadata. self._init_reorder_batch_threshold( @@ -347,7 +331,9 @@ class FlashMLASparseMetadataBuilder( device=device, ) - self.fp8_use_mixed_batch = self.num_heads < MIN_HEADS_FOR_BF16_PREFILL + self.fp8_use_mixed_batch = ( + self.num_heads < MIN_HEADS_FOR_BF16_PREFILL and not self.use_hisparse + ) if parallel_config.decode_context_parallel_size > 1: if parallel_config.dcp_comm_backend != "ag_rs": @@ -385,37 +371,21 @@ class FlashMLASparseMetadataBuilder( def _build_fp8_mixed_decode_prefill( self, - common_attn_metadata: CommonAttentionMetadata, - ) -> "FlashMLASparseMetadata.FP8KernelMetadata": + ) -> FlashMLASparseMetadata.FP8KernelMetadata: """Build FP8 metadata treating MQA tokens as one batch. The scheduler initializes lazily from the runtime query shape, which may be the full batch or only decodes when prefills use dense MHA. This avoids the BF16 prefill kernel's head-padding overhead at high TP. """ - num_tokens = common_attn_metadata.num_actual_tokens - # Use padded head count since that's what the kernel will see - padded_heads = self.fp8_decode_padded_heads - - # Build metadata for all tokens as a single batch - scheduler_metadata, _ = get_mla_metadata( - cache_seqlens=self.topk_tokens_tensor[:1], # Single batch - num_q_tokens_per_head_k=num_tokens * padded_heads, - topk=self.topk_tokens, - num_heads_q=padded_heads, - num_heads_k=1, - is_fp8_kvcache=True, - ) - - fp8_metadata = FlashMLASparseMetadata.FP8KernelMetadata( + scheduler_metadata, _ = get_mla_metadata() + return FlashMLASparseMetadata.FP8KernelMetadata( scheduler_metadata=scheduler_metadata, cache_lens=self.max_model_len_tensor[:1], dummy_block_table=self.dummy_block_table[:1], ) - return fp8_metadata - def _build_fp8_separate_prefill_decode( self, common_attn_metadata: CommonAttentionMetadata, @@ -520,6 +490,9 @@ class FlashMLASparseMetadataBuilder( chunk_block_table = common_attn_metadata.block_table_tensor[ num_decodes + chunk_start : num_decodes + chunk_end ] + chunk_seq_lens = common_attn_metadata.seq_lens[ + num_decodes + chunk_start : num_decodes + chunk_end + ] prefill_chunks.append( FP8Meta.Prefill.Chunk( @@ -528,6 +501,7 @@ class FlashMLASparseMetadataBuilder( req_start_idx=chunk_start, workspace_starts=chunk_workspace_starts, chunk_tot_seqlen=chunk_tot_seqlen, + seq_lens=chunk_seq_lens, ) ) @@ -569,9 +543,7 @@ class FlashMLASparseMetadataBuilder( metadata.fp8_use_mixed_batch = self.fp8_use_mixed_batch if self.use_fp8_kv_cache: if self.fp8_use_mixed_batch: - metadata.fp8_extra_metadata = self._build_fp8_mixed_decode_prefill( - common_attn_metadata - ) + metadata.fp8_extra_metadata = self._build_fp8_mixed_decode_prefill() else: metadata.fp8_extra_metadata = self._build_fp8_separate_prefill_decode( common_attn_metadata, metadata @@ -679,30 +651,83 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): topk_indices: torch.Tensor, attn_metadata: FlashMLASparseMetadata, actual_num_heads: int, - ) -> tuple[torch.Tensor, torch.Tensor]: - # Convert per-request indices to global slots (decode) or workspace - # offsets (prefill). req_id_per_token covers the whole batch; slice it - # to the MQA tokens (q may exclude prefill tokens routed to dense MHA). + ) -> tuple[torch.Tensor, torch.Tensor | None]: + index_group = self.index_group + if isinstance(index_group, HiSparseMLAIndexGroup): + cache = index_group.cache(self.index_group_index) + else: + cache = None + block_table = attn_metadata.block_table + # req_id_per_token covers the whole batch; slice it to the MQA tokens + # (q may exclude prefill tokens routed to dense MHA). + req_id_per_token = attn_metadata.req_id_per_token[: topk_indices.shape[0]] + decode_out: torch.Tensor | None = None + if cache is not None: + assert isinstance(index_group, HiSparseMLAIndexGroup) + num_decode_tokens = attn_metadata.num_decode_tokens + if num_decode_tokens > 0: + decode_topk, decode_lengths = ( + index_group.convert_decode_logical_to_physical_topk( + self.index_group_index, + topk_indices[:num_decode_tokens], + attn_metadata, + return_valid_counts=True, + ) + ) + decode_out, _ = self._bf16_flash_mla_kernel( + q[:num_decode_tokens], + index_group.physical_kv_cache(self.index_group_index), + decode_topk, + decode_lengths, + actual_num_heads, + ) + if num_decode_tokens == q.shape[0]: + return decode_out, None + q = q[num_decode_tokens:] + topk_indices = topk_indices[num_decode_tokens:] + kv_c_and_k_pe_cache, block_table, req_id_per_token = ( + index_group.stage_prefill_rows( + self.index_group_index, kv_c_and_k_pe_cache, attn_metadata + ) + ) + # Convert per-request indices to global slots (decode) or workspace offsets. kv_rows, block_stride_rows = flat_kv_row_view( kv_c_and_k_pe_cache, attn_metadata.block_size ) - topk_indices, topk_length = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[: topk_indices.shape[0]], - attn_metadata.block_table, - topk_indices, - BLOCK_SIZE=attn_metadata.block_size, - BLOCK_STRIDE_ROWS=block_stride_rows, - NUM_TOPK_TOKENS=topk_indices.shape[1], - return_valid_counts=True, + decode_only = ( + attn_metadata.num_decode_tokens + == attn_metadata.num_actual_tokens + == topk_indices.shape[0] ) + uses_host_cache = isinstance(index_group, HiSparseMLAIndexGroup) + if not uses_host_cache and decode_only: + topk_indices, topk_length = self._convert_logical_to_physical_topk( + topk_indices, + attn_metadata, + block_stride_rows=block_stride_rows, + return_valid_counts=True, + ) + else: + topk_indices, topk_length = triton_convert_req_index_to_global_index( + req_id_per_token, + block_table, + topk_indices, + BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + NUM_TOPK_TOKENS=topk_indices.shape[1], + return_valid_counts=True, + ) - return self._bf16_flash_mla_kernel( + attn_out, lse = self._bf16_flash_mla_kernel( q, kv_rows, topk_indices, topk_length, actual_num_heads, ) + if decode_out is None: + return attn_out, lse + return torch.cat([decode_out, attn_out], dim=0), None def _forward_fp8_kv_separate_prefill_decode( self, @@ -721,6 +746,28 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): "FP8 sparse MLA expects either the decode subset or the full batch" ) + decode_topk: torch.Tensor | None = None + index_group = self.index_group + uses_host_cache = isinstance(index_group, HiSparseMLAIndexGroup) + if uses_host_cache and num_decode_tokens > 0: + decode_topk = topk_indices[:num_decode_tokens] + + prefill_ready = None + if num_prefill_tokens > 0 and uses_host_cache: + assert fp8_metadata.prefill is not None + first_chunk = fp8_metadata.prefill.chunks[0] + assert isinstance(index_group, HiSparseMLAIndexGroup) + prefill_ready = index_group.gather_fp8_prefill( + self.index_group_index, + kv_c_and_k_pe_cache, + self.prefill_bf16_workspace[: first_chunk.chunk_tot_seqlen], + first_chunk.block_table, + first_chunk.workspace_starts, + len(first_chunk.block_table), + attn_metadata, + first_chunk.req_start_idx, + ) + prefill_request_ids = None prefill_workspace_starts = None has_prefill_workspace = False @@ -736,17 +783,26 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): # For BF16 cache: always use global cache slots (no workspace) # prefill_workspace_starts has been adjusted in-place per chunk so # prefill indices automatically come out chunk-local - topk_indices, topk_length = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[: topk_indices.shape[0]], - attn_metadata.block_table, - topk_indices, - BLOCK_SIZE=attn_metadata.block_size, - NUM_TOPK_TOKENS=topk_indices.shape[1], - HAS_PREFILL_WORKSPACE=has_prefill_workspace, - prefill_workspace_request_ids=prefill_request_ids, - prefill_workspace_starts=prefill_workspace_starts, - return_valid_counts=True, - ) + topk_length = None + if num_prefill_tokens == 0 and not uses_host_cache: + topk_indices, topk_length = self._convert_logical_to_physical_topk( + topk_indices, + attn_metadata, + block_stride_rows=None, + return_valid_counts=True, + ) + elif num_prefill_tokens > 0: + topk_indices, topk_length = triton_convert_req_index_to_global_index( + attn_metadata.req_id_per_token[: topk_indices.shape[0]], + attn_metadata.block_table, + topk_indices, + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + HAS_PREFILL_WORKSPACE=has_prefill_workspace, + prefill_workspace_request_ids=prefill_request_ids, + prefill_workspace_starts=prefill_workspace_starts, + return_valid_counts=True, + ) fp8_metadata = attn_metadata.fp8_extra_metadata assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8SeparatePrefillDecode) @@ -755,6 +811,16 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): q: torch.Tensor, topk_indices: torch.Tensor, ) -> torch.Tensor: + assert fp8_metadata.decode is not None + if uses_host_cache: + return self._host_backed_fp8_decode( + q, + topk_indices, + attn_metadata, + fp8_metadata.decode.kernel_metadata, + num_decodes, + fp8_metadata.decode.decode_query_len, + ) # Reshape q: (num_decode_tokens, num_heads, head_dim) # -> (num_decodes, seq_len, num_heads, head_dim) q = reshape_query_for_spec_decode(q, num_decodes) @@ -762,7 +828,6 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): # Reshape topk_indices: (num_decode_tokens, topk) # -> (num_decodes, seq_len, topk) topk_indices = topk_indices.view(num_decodes, seq_len, -1) - assert fp8_metadata.decode is not None attn_out, _ = self._fp8_flash_mla_kernel( q=q, kv_c_and_k_pe_cache=kv_c_and_k_pe_cache, @@ -776,7 +841,9 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): # Pure decode: direct call without allocation if num_decode_tokens > 0 and num_prefill_tokens == 0: assert fp8_metadata.decode is not None - attn_out = _fp8_decode(q, topk_indices) + attn_out = _fp8_decode( + q, decode_topk if decode_topk is not None else topk_indices + ) else: # Mixed or pure prefill: allocate output tensor attn_out = q.new_empty( @@ -788,13 +855,30 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): if num_decode_tokens > 0: attn_out[:num_decode_tokens] = _fp8_decode( q[:num_decode_tokens], - topk_indices[:num_decode_tokens], + decode_topk + if decode_topk is not None + else topk_indices[:num_decode_tokens], ) assert fp8_metadata.prefill is not None - for chunk in fp8_metadata.prefill.chunks: + for chunk_index, chunk in enumerate(fp8_metadata.prefill.chunks): chunk_workspace = self.prefill_bf16_workspace[: chunk.chunk_tot_seqlen] - if self.kv_cache_dtype == "fp8_ds_mla": + if uses_host_cache and chunk_index > 0: + assert isinstance(index_group, HiSparseMLAIndexGroup) + prefill_ready = index_group.gather_fp8_prefill( + self.index_group_index, + kv_c_and_k_pe_cache, + chunk_workspace, + chunk.block_table, + chunk.workspace_starts, + len(chunk.block_table), + attn_metadata, + chunk.req_start_idx, + ) + if uses_host_cache: + assert prefill_ready is not None + current_stream().wait_event(prefill_ready) + elif self.kv_cache_dtype == "fp8_ds_mla": ops.cp_gather_and_upconvert_fp8_kv_cache( kv_c_and_k_pe_cache, chunk_workspace, @@ -813,6 +897,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): chunk_q = q[chunk.tokens_slice] chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice] + assert topk_length is not None chunk_topk_length = topk_length[chunk.tokens_slice] attn_out[chunk.tokens_slice], _ = self._bf16_flash_mla_kernel( @@ -839,6 +924,17 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): The lse is only returned when DCP needs it, otherwise None. """ + assert attn_metadata.fp8_extra_metadata is not None + assert isinstance( + attn_metadata.fp8_extra_metadata, + FlashMLASparseMetadata.FP8KernelMetadata, + ) + fp8_metadata = attn_metadata.fp8_extra_metadata + + block_table = attn_metadata.block_table + # req_id_per_token covers the whole batch; slice it to the MQA tokens + # (q may exclude prefill tokens routed to dense MHA). + req_id_per_token = attn_metadata.req_id_per_token[: topk_indices.shape[0]] if self.dcp_world_size > 1: # The indexer emits global token ids; keep this rank's shard and # convert to local slots. compact_valid_to_front=False keeps the @@ -846,8 +942,8 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): # empty-row neutralization below relies on. req_id is sliced to # topk_indices rows (the converter grids from req_id). topk_indices = triton_filter_and_convert_dcp_index( - attn_metadata.req_id_per_token[: topk_indices.shape[0]], - attn_metadata.block_table, + req_id_per_token, + block_table, topk_indices, dcp_size=self.dcp_world_size, dcp_rank=self.dcp_rank, @@ -859,19 +955,26 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): else: # Convert per-request indices to global slots (decode) or workspace # offsets (prefill). - topk_indices = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[: topk_indices.shape[0]], - attn_metadata.block_table, - topk_indices, - BLOCK_SIZE=attn_metadata.block_size, - NUM_TOPK_TOKENS=topk_indices.shape[1], + decode_only = ( + attn_metadata.num_decode_tokens + == attn_metadata.num_actual_tokens + == topk_indices.shape[0] ) - - assert attn_metadata.fp8_extra_metadata is not None - assert isinstance( - attn_metadata.fp8_extra_metadata, FlashMLASparseMetadata.FP8KernelMetadata - ) - fp8_metadata = attn_metadata.fp8_extra_metadata + if decode_only: + topk_indices = self._convert_logical_to_physical_topk( + topk_indices, + attn_metadata, + block_stride_rows=None, + return_valid_counts=False, + ) + else: + topk_indices = triton_convert_req_index_to_global_index( + req_id_per_token, + block_table, + topk_indices, + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + ) _attn_out, _lse = self._fp8_flash_mla_kernel( q=q.unsqueeze(0), # unsqueeze to add batch_dim: (T, H, D) -> (1, T, H, D) @@ -937,6 +1040,37 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): return out, lse + def _host_backed_fp8_decode( + self, + q: torch.Tensor, + topk_indices: torch.Tensor, + attn_metadata: FlashMLASparseMetadata, + kernel_metadata: FlashMLASparseMetadata.FP8KernelMetadata, + num_decodes: int, + decode_query_len: int, + ) -> torch.Tensor: + assert isinstance(self.index_group, HiSparseMLAIndexGroup) + physical_topk = self.index_group.convert_decode_logical_to_physical_topk( + self.index_group_index, + topk_indices, + attn_metadata, + return_valid_counts=False, + num_decodes=num_decodes, + decode_query_len=decode_query_len, + ) + assert isinstance(physical_topk, torch.Tensor) + q = reshape_query_for_spec_decode(q, num_decodes) + physical_topk = physical_topk.view(num_decodes, q.shape[1], -1) + output, _ = self._fp8_flash_mla_kernel( + q=q, + kv_c_and_k_pe_cache=self.index_group.physical_kv_cache( + self.index_group_index + ), + topk_indices=physical_topk, + kernel_metadata=kernel_metadata, + ) + return reshape_attn_output_for_spec_decode(output) + def _bf16_flash_mla_kernel( self, q: torch.Tensor, diff --git a/vllm/v1/attention/backends/mla/index_group.py b/vllm/v1/attention/backends/mla/index_group.py new file mode 100644 index 0000000000..8ba0171835 --- /dev/null +++ b/vllm/v1/attention/backends/mla/index_group.py @@ -0,0 +1,503 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + +from vllm import _custom_ops as ops +from vllm.config import VllmConfig +from vllm.utils.torch_utils import current_stream +from vllm.v1.attention.backends.mla.sparse_utils import ( + triton_convert_req_index_to_global_index, +) +from vllm.v1.hisparse.runtime import ( + FP8_DS_MLA_ROW_BYTES, + HiSparseCacheHandle, + HiSparsePrefillStagingPlan, + create_hisparse_cache_handle, +) + + +def _create_side_stream(device: torch.device) -> torch.Stream: + return torch.Stream(device=device) + + +def _create_event() -> torch.Event: + return torch.Event() + + +@dataclass +class SparseMLAIndexGroup: + """Layers that consume one sparse-indexer result.""" + + logical_topk_indices: torch.Tensor + physical_topk_indices: torch.Tensor + valid_topk_counts: torch.Tensor + request_ids: torch.Tensor + side_stream: torch.Stream + logical_topk_ready: torch.Event + physical_topk_ready: torch.Event + has_indexer: bool + num_layers: int = 0 + + def register_layer( + self, + vllm_config: VllmConfig | None = None, + *, + head_size: int | None = None, + kv_cache_dtype: str | None = None, + ) -> int: + layer_index = self.num_layers + self.num_layers += 1 + return layer_index + + def set_logical_topk_ready(self, layer_index: int) -> None: + if layer_index == 0 and self.has_indexer: + self.logical_topk_ready.record(current_stream()) + + def prepare_for_batch(self, layer_index: int, attn_metadata: Any | None) -> None: + pass + + def gather_fp8_prefill( + self, + layer_index: int, + source_cache: torch.Tensor, + dst: torch.Tensor, + block_table: torch.Tensor, + workspace_starts: torch.Tensor, + batch_size: int, + attn_metadata: Any, + request_start: int, + ) -> torch.Event | None: + ops.cp_gather_and_upconvert_fp8_kv_cache( + source_cache, + dst, + block_table, + workspace_starts, + batch_size, + ) + return None + + def convert_logical_to_physical_topk( + self, + layer_index: int, + logical_topk_indices: torch.Tensor, + attn_metadata: Any, + *, + block_stride_rows: int | None, + return_valid_counts: bool, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Convert one index result once for all layers in this group.""" + return self._convert_once( + layer_index, + logical_topk_indices, + attn_metadata.req_id_per_token[: logical_topk_indices.shape[0]], + attn_metadata.block_table, + attn_metadata.block_size, + block_stride_rows=block_stride_rows, + return_valid_counts=return_valid_counts, + ) + + def _convert_once( + self, + layer_index: int, + logical_topk_indices: torch.Tensor, + req_id_per_token: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + *, + block_stride_rows: int | None, + return_valid_counts: bool, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + num_tokens = logical_topk_indices.shape[0] + if num_tokens > self.physical_topk_indices.shape[0]: + return triton_convert_req_index_to_global_index( + req_id_per_token, + block_table, + logical_topk_indices, + BLOCK_SIZE=block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + NUM_TOPK_TOKENS=logical_topk_indices.shape[1], + return_valid_counts=return_valid_counts, + ) + + physical_topk_indices = self.physical_topk_indices[:num_tokens] + valid_topk_counts = self.valid_topk_counts[:num_tokens] + compute_stream = current_stream() + if layer_index == 0: + if self.has_indexer: + self.side_stream.wait_event(self.logical_topk_ready) + else: + self.side_stream.wait_stream(compute_stream) + with self.side_stream: + triton_convert_req_index_to_global_index( + req_id_per_token, + block_table, + logical_topk_indices, + BLOCK_SIZE=block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + NUM_TOPK_TOKENS=logical_topk_indices.shape[1], + return_valid_counts=return_valid_counts, + out=physical_topk_indices, + valid_counts_out=( + valid_topk_counts if return_valid_counts else None + ), + ) + self.physical_topk_ready.record(self.side_stream) + compute_stream.wait_event(self.physical_topk_ready) + + if return_valid_counts: + return physical_topk_indices, valid_topk_counts + return physical_topk_indices + + +@dataclass +class HiSparseMLAIndexGroup(SparseMLAIndexGroup): + """Sparse index group with per-layer host-backed KV caches.""" + + caches: list[HiSparseCacheHandle] = field(default_factory=list) + hisparse_group: Any | None = None + prefill_stream: torch.Stream | None = None + prefill_ready_events: list[torch.Event] = field(default_factory=list) + + def register_layer( + self, + vllm_config: VllmConfig | None = None, + *, + head_size: int | None = None, + kv_cache_dtype: str | None = None, + ) -> int: + assert vllm_config is not None + assert head_size is not None + assert kv_cache_dtype is not None + layer_index = super().register_layer() + if self.prefill_stream is None: + self.prefill_stream = _create_side_stream(self.logical_topk_indices.device) + if kv_cache_dtype == "fp8_ds_mla": + row_width = FP8_DS_MLA_ROW_BYTES + kv_dtype = torch.uint8 + else: + from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype + + row_width = head_size + kv_dtype = kv_cache_dtype_str_to_dtype( + kv_cache_dtype, vllm_config.model_config + ) + cache = create_hisparse_cache_handle( + vllm_config, + self.logical_topk_indices.shape[1], + is_index_group_leader=layer_index == 0, + row_width=row_width, + kv_dtype=kv_dtype, + index_group=self, + ) + assert cache is not None + self.caches.append(cache) + cache.index_group_caches = self.caches + self.prefill_ready_events.append(_create_event()) + return layer_index + + def cache(self, layer_index: int) -> HiSparseCacheHandle: + return self.caches[layer_index] + + def physical_kv_cache(self, layer_index: int) -> torch.Tensor: + cache = self.cache(layer_index) + return cache.runtime.hot.attention_cache + + def prepare_for_batch(self, layer_index: int, attn_metadata: Any | None) -> None: + if layer_index == 0: + self.cache(0).prepare_group_for_batch(attn_metadata) + + def convert_logical_to_physical_topk( + self, + layer_index: int, + logical_topk_indices: torch.Tensor, + attn_metadata: Any, + *, + block_stride_rows: int | None, + return_valid_counts: bool, + req_id_per_token: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + cache = self.cache(layer_index) + num_tokens = logical_topk_indices.shape[0] + if req_id_per_token is None: + req_id_per_token = attn_metadata.req_id_per_token[:num_tokens] + if num_tokens > self.physical_topk_indices.shape[0]: + # Prefill-sized batches do not fit the decode residency workspace. + # Non-resident prefills are staged before reaching this path. + assert cache.all_context_pages_resident + leader = self.cache(0) + assert leader.view is not None and leader.block_table is not None + return self._convert_once( + layer_index, + logical_topk_indices, + req_id_per_token, + leader.block_table, + leader.view.block_size, + block_stride_rows=leader.view.attention_block_stride, + return_valid_counts=return_valid_counts, + ) + source_block_table = cache.source_block_table + assert source_block_table is not None + return cache.swap_in( + req_id_per_token, + block_table=source_block_table, + logical_topk_indices=logical_topk_indices, + block_size=attn_metadata.block_size, + return_valid_counts=return_valid_counts, + ) + + def convert_decode_logical_to_physical_topk( + self, + layer_index: int, + logical_topk_indices: torch.Tensor, + attn_metadata: Any, + *, + return_valid_counts: bool, + num_decodes: int | None = None, + decode_query_len: int | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + num_tokens = logical_topk_indices.shape[0] + if num_decodes is None: + num_decodes = attn_metadata.num_decodes + if decode_query_len is None: + decode_query_len = attn_metadata.decode_max_query_len + if decode_query_len == 1: + return self.convert_logical_to_physical_topk( + layer_index, + logical_topk_indices, + attn_metadata, + block_stride_rows=None, + return_valid_counts=return_valid_counts, + req_id_per_token=self.request_ids[:num_decodes], + ) + + if num_tokens != num_decodes * decode_query_len: + query_start_loc = attn_metadata.query_start_loc[: num_decodes + 1] + physical_topk_indices = self.physical_topk_indices[: num_tokens + 1] + valid_topk_counts = self.valid_topk_counts[: num_tokens + 1] + request_ids = self.request_ids[:num_decodes] + cache = self.cache(layer_index) + source_block_table = cache.source_block_table + assert source_block_table is not None + for step in range(decode_query_len): + token_indices = query_start_loc[:-1] + step + active = (token_indices < query_start_loc[1:]) & ( + token_indices < num_tokens + ) + output_indices = torch.where( + active, token_indices, torch.full_like(token_indices, num_tokens) + ).long() + token_indices = token_indices.clamp(0, num_tokens - 1).long() + step_topk = logical_topk_indices.index_select( + 0, token_indices + ).masked_fill(~active.unsqueeze(1), -1) + step_result = cache.swap_in( + request_ids, + block_table=source_block_table, + logical_topk_indices=step_topk, + block_size=attn_metadata.block_size, + return_valid_counts=return_valid_counts, + ) + if layer_index == 0: + if return_valid_counts: + step_indices, step_counts = step_result + valid_topk_counts.index_copy_(0, output_indices, step_counts) + else: + step_indices = step_result + physical_topk_indices.index_copy_(0, output_indices, step_indices) + physical_topk_indices = physical_topk_indices[:num_tokens] + if return_valid_counts: + return physical_topk_indices, valid_topk_counts[:num_tokens] + return physical_topk_indices + + assert num_tokens == num_decodes * decode_query_len + logical_topk_by_request = logical_topk_indices.view( + num_decodes, decode_query_len, -1 + ) + physical_topk_by_request = self.physical_topk_indices[:num_tokens].view( + num_decodes, decode_query_len, -1 + ) + valid_topk_by_request = self.valid_topk_counts[:num_tokens].view( + num_decodes, decode_query_len + ) + request_ids = self.request_ids[:num_decodes] + for step in range(decode_query_len): + cache = self.cache(layer_index) + source_block_table = cache.source_block_table + assert source_block_table is not None + cache.swap_in( + request_ids, + block_table=source_block_table, + logical_topk_indices=logical_topk_by_request[:, step], + block_size=attn_metadata.block_size, + return_valid_counts=return_valid_counts, + attention_indices_out=( + physical_topk_by_request[:, step] if layer_index == 0 else None + ), + valid_counts_out=( + valid_topk_by_request[:, step] + if layer_index == 0 and return_valid_counts + else None + ), + ) + physical_topk_indices = physical_topk_by_request.view(num_tokens, -1) + if return_valid_counts: + return physical_topk_indices, valid_topk_by_request.view(num_tokens) + return physical_topk_indices + + def stage_prefill_rows( + self, layer_index: int, kv_cache: torch.Tensor, attn_metadata: Any + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + cache = self.cache(layer_index) + prefill = attn_metadata.prefill + staging_plan = prefill.host_staging_plan if prefill is not None else None + assert staging_plan is not None + resident_cache = None + if cache.view is not None and cache.block_table is not None: + staging_plan.ensure_gpu_sources( + cache.block_table[attn_metadata.num_decodes :], + cache.view.block_size, + ) + resident_cache = cache.view.cache + staged_cache = cache.runtime.gather_prefill_cache( + kv_cache, + staging_plan, + resident_cache=resident_cache, + ) + req_ids = attn_metadata.req_id_per_token[attn_metadata.num_decode_tokens :] + if attn_metadata.num_decodes > 0: + req_ids = req_ids - attn_metadata.num_decodes + return staged_cache, staging_plan.block_table, req_ids + + def _prefill_gather_plan( + self, layer_index: int, attn_metadata: Any + ) -> HiSparsePrefillStagingPlan: + cache = self.cache(layer_index) + prefill = attn_metadata.prefill + plan = prefill.host_staging_plan if prefill is not None else None + assert plan is not None + assert cache.view is not None and cache.block_table is not None + plan.ensure_gpu_sources( + cache.block_table[attn_metadata.num_decodes :], + cache.view.block_size, + ) + assert plan.gpu_row_ids is not None + return plan + + def gather_fp8_prefill( + self, + layer_index: int, + source_cache: torch.Tensor, + dst: torch.Tensor, + block_table: torch.Tensor, + workspace_starts: torch.Tensor, + batch_size: int, + attn_metadata: Any, + request_start: int, + ) -> torch.Event | None: + cache = self.cache(layer_index) + plan = self._prefill_gather_plan(layer_index, attn_metadata) + assert cache.view is not None and plan.gpu_row_ids is not None + row_width = source_cache.shape[-1] + prefill_stream = self.prefill_stream + assert prefill_stream is not None + block_table = plan.block_table[request_start : request_start + batch_size] + prefill_stream.wait_stream(current_stream()) + with prefill_stream: + ops.cp_gather_and_upconvert_fp8_kv_cache( + cache.view.cache, + dst, + block_table, + workspace_starts, + batch_size, + host_cache=source_cache.view(-1, row_width), + host_row_ids=plan.row_ids, + device_row_ids=plan.gpu_row_ids, + ) + ready = self.prefill_ready_events[layer_index] + ready.record(prefill_stream) + return ready + + +class SparseMLAIndexGroupBuilder: + """Assign sparse MLA layers to their index-producing layer.""" + + def __init__( + self, logical_topk_indices: torch.Tensor, max_decode_rows: int | None = None + ) -> None: + self.logical_topk_indices = logical_topk_indices + self.max_decode_rows = ( + logical_topk_indices.shape[0] + if max_decode_rows is None + else max_decode_rows + ) + self.current_group: SparseMLAIndexGroup | None = None + + def register_layer( + self, + is_index_producing_layer: bool, + vllm_config: VllmConfig | None = None, + *, + head_size: int | None = None, + kv_cache_dtype: str | None = None, + ) -> tuple[SparseMLAIndexGroup, int]: + if is_index_producing_layer or self.current_group is None: + group_cls = ( + HiSparseMLAIndexGroup + if vllm_config is not None + and vllm_config.attention_config.hisparse_config is not None + else SparseMLAIndexGroup + ) + workspace_rows = self.max_decode_rows + physical_topk_indices = torch.empty( + (workspace_rows + 1, self.logical_topk_indices.shape[1]), + dtype=self.logical_topk_indices.dtype, + device=self.logical_topk_indices.device, + ) + self.current_group = group_cls( + logical_topk_indices=self.logical_topk_indices, + physical_topk_indices=physical_topk_indices, + valid_topk_counts=torch.empty( + workspace_rows + 1, + dtype=torch.int32, + device=self.logical_topk_indices.device, + ), + request_ids=torch.arange( + workspace_rows, + dtype=torch.int32, + device=self.logical_topk_indices.device, + ), + side_stream=_create_side_stream(self.logical_topk_indices.device), + logical_topk_ready=_create_event(), + physical_topk_ready=_create_event(), + has_indexer=is_index_producing_layer, + ) + group = self.current_group + group_index = group.register_layer( + vllm_config, + head_size=head_size, + kv_cache_dtype=kv_cache_dtype, + ) + return group, group_index + + +def get_sparse_mla_index_group_max_rows(vllm_config: VllmConfig) -> int: + max_query_len = 1 + speculative_config = vllm_config.speculative_config + if ( + speculative_config is not None + and speculative_config.num_speculative_tokens is not None + ): + max_query_len += speculative_config.num_speculative_tokens * ( + 2 if speculative_config.parallel_drafting else 1 + ) + scheduler_config = vllm_config.scheduler_config + return min( + scheduler_config.max_num_batched_tokens, + scheduler_config.max_num_seqs * max_query_len, + ) diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index 8d196f9630..d9b45892a6 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -407,6 +407,8 @@ def triton_convert_req_index_to_global_index( prefill_workspace_request_ids: torch.Tensor | None = None, prefill_workspace_starts: torch.Tensor | None = None, return_valid_counts: bool = False, + out: torch.Tensor | None = None, + valid_counts_out: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """ out[token_id, indice_id] = @@ -462,16 +464,36 @@ def triton_convert_req_index_to_global_index( # contiguous prefix [0, valid_count) and leaves the tail unwritten, so # pre-fill -1 there. flash_mla_sparse_fwd then bounds attention to # [:topk_length] == exactly the valid set (no dropped tokens). - if return_valid_counts: - out = torch.full_like(token_indices_c, -1) + if out is None: + out = ( + torch.full_like(token_indices_c, -1) + if return_valid_counts + else torch.empty_like(token_indices_c) + ) else: - out = torch.empty_like(token_indices_c) + assert out.dtype == token_indices_c.dtype + assert out.device == token_indices_c.device + assert out.shape == token_indices_c.shape + assert out.is_contiguous() + if return_valid_counts: + out.fill_(-1) valid_counts: torch.Tensor | None = None if return_valid_counts: - # Zero-init only matters for the atomic accumulation path. - alloc = torch.empty if single_tile else torch.zeros - valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device) + if valid_counts_out is None: + # Zero-init only matters for the atomic accumulation path. + alloc = torch.empty if single_tile else torch.zeros + valid_counts = alloc( + num_tokens, dtype=torch.int32, device=token_indices.device + ) + else: + assert valid_counts_out.dtype == torch.int32 + assert valid_counts_out.device == token_indices.device + assert valid_counts_out.shape == (num_tokens,) + assert valid_counts_out.is_contiguous() + valid_counts = valid_counts_out + if not single_tile: + valid_counts.zero_() # Prepare prefill pointers if HAS_PREFILL_WORKSPACE: diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index 1efa0e0cea..e2a06e7e4e 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Iterable, Sequence +from collections.abc import Callable, Iterable, Sequence from typing import Any from vllm.distributed.kv_events import ( @@ -157,6 +157,7 @@ class BlockPool: actual block size can be a multiple of hash_block_size. enable_kv_cache_events: Whether to enable kv cache events. metrics_collector: Optional metrics collector for tracking block residency. + medium: Storage medium reported in KV cache events. """ def __init__( @@ -166,14 +167,16 @@ class BlockPool: hash_block_size: int, enable_kv_cache_events: bool = False, metrics_collector: KVCacheMetricsCollector | None = None, + medium: str = MEDIUM_GPU, ): assert isinstance(num_gpu_blocks, int) and num_gpu_blocks > 0 self.num_gpu_blocks = num_gpu_blocks + self.medium = medium self.enable_caching = enable_caching self.hash_block_size = hash_block_size # All kv-cache blocks. self.blocks: list[KVCacheBlock] = [ - KVCacheBlock(idx) for idx in range(num_gpu_blocks) + KVCacheBlock(idx, pool=self) for idx in range(num_gpu_blocks) ] # Free block queue that constructs and manipulates a doubly linked # list of free blocks (including eviction candidates when caching is @@ -194,6 +197,9 @@ class BlockPool: self.kv_event_queue: list[KVCacheEvent] = [] self.metrics_collector = metrics_collector + # Callbacks for blocks released with ``unpin_blocks`` whose contents + # are still being read until the pool reuses them. + self._reuse_watchers: dict[int, Callable[[KVCacheBlock], None]] = {} def get_cached_block( self, block_hash: BlockHash, kv_cache_group_ids: list[int] @@ -364,7 +370,7 @@ class BlockPool: token_ids=request.all_token_ids[start_token_idx:end_token_idx], block_size=block_size, lora_id=request.lora_request.adapter_id if request.lora_request else None, - medium=MEDIUM_GPU, + medium=self.medium, lora_name=request.lora_request.name if request.lora_request else None, extra_keys=extra_keys_list if extra_keys_list else None, group_idx=kv_cache_group_id, @@ -543,7 +549,7 @@ class BlockPool: lora_id=request.lora_request.adapter_id if request.lora_request else None, - medium=MEDIUM_GPU, + medium=self.medium, lora_name=request.lora_request.name if request.lora_request else None, @@ -610,7 +616,7 @@ class BlockPool: self.kv_event_queue.append( BlockRemoved( block_hashes=[maybe_convert_block_hash(get_block_hash(block_hash))], - medium=MEDIUM_GPU, + medium=self.medium, group_idx=get_group_id(block_hash), ) ) @@ -671,6 +677,9 @@ class BlockPool: ret: list[KVCacheBlock] = self.free_block_queue.popleft_n(num_blocks) + if self._reuse_watchers: + self._notify_reuse(ret) + # In order to only iterate the list once, we duplicated code a bit if self.enable_caching: for block in ret: @@ -687,6 +696,33 @@ class BlockPool: self.metrics_collector.on_block_allocated(block) return ret + def _notify_reuse(self, blocks: Iterable[KVCacheBlock]) -> None: + for block in blocks: + watcher = self._reuse_watchers.pop(block.block_id, None) + if watcher is not None: + watcher(block) + + def unpin_blocks( + self, + blocks: Iterable[KVCacheBlock], + on_reuse: Callable[[KVCacheBlock], None], + ) -> None: + """Release references to blocks that stay readable until reused. + + The blocks become last-resort eviction candidates regardless of prefix + caching: they join the tail of the free queue and count as free. + ``on_reuse`` fires when ``get_new_blocks`` hands a block out (or the + cache is reset), so the caller can stop reading it. + """ + released: list[KVCacheBlock] = [] + for block in blocks: + assert block.ref_cnt > 0 and not block.is_null + block.ref_cnt -= 1 + self._reuse_watchers[block.block_id] = on_reuse + if block.ref_cnt == 0: + released.append(block) + self.free_block_queue.append_n(released) + def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool: """ If a block is cached in `cached_block_hash_to_block`, we reset its hash @@ -742,7 +778,11 @@ class BlockPool: # Identify blocks with hash (LRU cache) and without it (never match APC) blocks_to_evict_last = [] blocks_to_evict_first = [] + other_pools: dict[BlockPool, list[KVCacheBlock]] = {} for block in ordered_blocks: + if block.pool is not None and block.pool is not self: + other_pools.setdefault(block.pool, []).append(block) + continue block.ref_cnt -= 1 if block.ref_cnt == 0 and not block.is_null: if block.block_hash is None or not self.enable_caching: @@ -756,6 +796,8 @@ class BlockPool: self.free_block_queue.prepend_n(blocks_to_evict_first) # Blocks to reuse last are appended to the end of the free queue. self.free_block_queue.append_n(blocks_to_evict_last) + for pool, blocks in other_pools.items(): + pool.free_blocks(blocks) def evict_blocks(self, block_ids: set[int]) -> None: """evict blocks from the prefix cache by their block IDs. @@ -797,6 +839,10 @@ class BlockPool: # Remove all hashes so that no new blocks will hit. self.cached_block_hash_to_block = BlockHashToBlockMap() self.cached_block_hashes_by_block.clear() + if self._reuse_watchers: + self._notify_reuse( + [self.blocks[block_id] for block_id in list(self._reuse_watchers)] + ) # Remove all hashes from all blocks. for block in self.blocks: diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 5fa8812924..afb03ac3ea 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -139,6 +139,7 @@ class KVCacheCoordinator(ABC): max_in_flight_tokens=max_in_flight_tokens, max_model_len=max_model_len, block_pool=self.block_pool, + role=kv_cache_group.role, enable_caching=enable_caching, kv_cache_group_id=i, dcp_world_size=dcp_world_size_for_kv_cache_spec( @@ -179,7 +180,7 @@ class KVCacheCoordinator(ABC): apply_admission_cap: bool = False, ) -> int: """ - Get the number of blocks needed to be allocated for the request. + Get the number of device blocks needed to be allocated for the request. Args: request_id: The request ID. @@ -343,6 +344,8 @@ class KVCacheCoordinator(ABC): """ boundaries = self.get_replay_boundaries(request) for manager in self.single_type_managers: + if not manager.enable_caching: + continue # Only cache tokens with finalized KV. The last num_reprefillable_tokens # tokens can be re-prefilled during multi-module MTP. num_tokens_to_cache = max( @@ -355,6 +358,27 @@ class KVCacheCoordinator(ABC): replay_boundaries=boundaries, ) + def emit_cached_block_events( + self, request: Request, computed_blocks: tuple[list[KVCacheBlock], ...] + ) -> None: + for group_idx, group_blocks in enumerate(computed_blocks): + if group_blocks: + manager = self.single_type_managers[group_idx] + group = self.kv_cache_config.kv_cache_groups[group_idx] + manager.block_pool.emit_cached_block_events( + request, + len(group_blocks), + group.kv_cache_spec.block_size, + group_idx, + ) + + def reset_prefix_cache(self) -> bool: + """Reset each manager's pool once.""" + pools = dict.fromkeys( + manager.block_pool for manager in self.single_type_managers + ) + return all([pool.reset_prefix_cache() for pool in pools]) + def free(self, request_id: str) -> None: """ Free the blocks for the request. @@ -728,10 +752,12 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): # Try to find an existing group with the same spec for idx, group in enumerate(self.attention_groups): - if group.spec == spec: - assert manager_cls is group.manager_cls, ( - "Expected same manager class for identical KV cache specs." - ) + if ( + group.spec == spec + and group.manager_cls is manager_cls + and self.single_type_managers[group.group_ids[0]].block_pool + is self.single_type_managers[i].block_pool + ): group.group_ids.append(i) if use_eagle and not group.use_eagle: self.attention_groups[idx] = group._replace(use_eagle=True) @@ -782,6 +808,8 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): cached_num_computed_tokens = self._align_cacheable(num_computed_tokens) boundaries = self.get_replay_boundaries(request) for manager in self.single_type_managers: + if not manager.enable_caching: + continue num_tokens_to_cache = cached_num_computed_tokens # EAGLE groups match one block past each aligned boundary and drop # it, so make that lookahead block eligible to be cached. @@ -895,7 +923,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): block_hashes=block_hashes, max_length=_max_length, kv_cache_group_ids=group_ids, - block_pool=self.block_pool, + block_pool=self.single_type_managers[first_group_id].block_pool, kv_cache_spec=spec, drop_eagle_block=drop_eagle_block, alignment_tokens=self._cache_hit_alignment_tokens, @@ -929,7 +957,10 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): for group in self.attention_groups: if not isinstance(group.spec, FullAttentionSpec): continue - group_block_size = self.single_type_managers[group.group_ids[0]].block_size + manager = self.single_type_managers[group.group_ids[0]] + if manager.retains_longer_hit: + continue + group_block_size = manager.block_size num_blocks = cdiv(hit_length, group_block_size) for group_id in group.group_ids: if (blks := hit_blocks_by_group[group_id]) is not None: @@ -966,7 +997,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): block_hashes=block_hashes, max_length=max_cache_hit_length, kv_cache_group_ids=group_ids, - block_pool=self.block_pool, + block_pool=manager.block_pool, kv_cache_spec=spec, drop_eagle_block=use_eagle, alignment_tokens=self._cache_hit_alignment_tokens, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 084ba24311..20e037f60e 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -183,6 +183,11 @@ class KVCacheManager: manager.fine_grained_prefix_cache = True self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) self.block_pool = self.coordinator.block_pool + self.retained_hit_group_ids = tuple( + manager.kv_cache_group_id + for manager in self.coordinator.single_type_managers + if manager.retains_longer_hit + ) self.kv_cache_config = kv_cache_config # Watermark: minimum number of KV cache blocks to keep free when @@ -287,17 +292,7 @@ class KVCacheManager: and self.enable_kv_cache_events and getattr(request, "kv_cache_report_mode", "incremental") == "full" ): - for group_idx, group_blocks in enumerate(computed_blocks): - num_blocks = len(group_blocks) - if num_blocks > 0: - group = self.kv_cache_config.kv_cache_groups[group_idx] - block_size = group.kv_cache_spec.block_size - self.block_pool.emit_cached_block_events( - request, - num_blocks, - block_size, - group_idx, - ) + self.coordinator.emit_cached_block_events(request, computed_blocks) # The junction to pin is where the lagging sparse-retention group stops # (``num_new_computed_tokens``) plus the uncached shared prefix -- i.e. @@ -463,6 +458,16 @@ class KVCacheManager: if new_computed_blocks is not None: new_computed_block_list = new_computed_blocks.blocks + if self.retained_hit_group_ids: + # Only retain the prefix covered by the local and external hits. + reused = num_new_computed_tokens + num_external_computed_tokens + groups = list(new_computed_block_list) + for group_id in self.retained_hit_group_ids: + manager = self.coordinator.single_type_managers[group_id] + groups[group_id] = groups[group_id][ + : cdiv(reused, manager.block_size) + ] + new_computed_block_list = tuple(groups) else: new_computed_block_list = self.empty_kv_cache_blocks.blocks @@ -564,7 +569,7 @@ class KVCacheManager: # P/D: delay caching blocks if we have to recv from # remote. Update state for locally cached blocks. - if not self.enable_caching or delay_cache_blocks: + if delay_cache_blocks: return self.create_kv_cache_blocks(new_blocks) # NOTE(woosuk): We want to commit (cache) up to num_local_computed_tokens @@ -639,7 +644,7 @@ class KVCacheManager: bool: True if the prefix cache is successfully reset, False otherwise. """ - if not self.block_pool.reset_prefix_cache(): + if not self.coordinator.reset_prefix_cache(): return False if self.log_stats: assert self.prefix_cache_stats is not None @@ -741,6 +746,8 @@ class KVCacheManager: self.kv_cache_config.kv_cache_groups, self.get_blocks(request.request_id).blocks, ): + if not group.kv_cache_spec.prefix_cacheable: + continue if isinstance( group.kv_cache_spec, (CrossAttentionSpec, EncoderOnlyAttentionSpec), diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index e2e38ec5d5..996cce8c43 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -8,23 +8,30 @@ import math import os from collections import defaultdict from collections.abc import Callable, Iterable, Iterator, Sequence -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace from functools import partial -from typing import Any, NamedTuple, NewType, TypeAlias, cast, overload +from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, cast, overload from vllm import envs from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.utils.hashing import xxhash, xxhash_cbor -from vllm.utils.math_utils import cdiv +from vllm.utils.math_utils import cdiv, round_up from vllm.utils.mem_utils import format_gib from vllm.utils.torch_utils import get_dtype_size +from vllm.v1.hisparse.layout import ( + get_hisparse_gpu_memory_usage, + get_hisparse_host_pool_bytes, + get_hisparse_kv_cache_config, + get_hisparse_kv_cache_groups, +) from vllm.v1.kv_cache_interface import ( AttentionSpec, ChunkedLocalAttentionSpec, CircularBufferSpec, FullAttentionSpec, HiddenStateCacheSpec, + HiSparseHotSpec, KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, @@ -44,6 +51,10 @@ from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.request import Request from vllm.v1.utils import tensor_data +if TYPE_CHECKING: + from vllm.v1.core.block_pool import BlockPool + + # BlockHash represents the hash of a single KV-cache block used for # prefix caching. Treating it as a distinct type from `bytes` helps # catch accidental misuse when passing around raw byte strings. @@ -169,6 +180,10 @@ class KVCacheBlock: block_id: int # Reference count. ref_cnt: int = 0 + # Deferred frees may contain blocks from multiple pools. + pool: "BlockPool | None" = field( + default=None, repr=False, compare=False, kw_only=True + ) # The hash key (block hash + group id) of the block, only available # when the block is full and cached. _block_hash: BlockHashWithGroupId | None = None @@ -1057,21 +1072,32 @@ def get_max_concurrency_for_kv_cache_config( A request at max_model_len consumes whole blocks from each group's block table — cdiv(per-request bytes, page bytes) of the group's spec — and all - groups draw those block ids from one shared pool, so the per-request + device groups draw those block ids from one shared pool, so the per-request total is the sum over groups. The memory/page ratio is identical whether a group carries an aggregated UniformTypeKVCacheSpecs (worker config) or a representative per-layer spec (scheduler config), so both capacity call sites agree. + + Host groups use a separate pool; the smaller concurrency limit applies. """ - num_blocks_per_request = sum( - cdiv( + num_blocks_per_request = 0 + host_blocks_per_request = 0 + for group in kv_cache_config.kv_cache_groups: + required = cdiv( group.kv_cache_spec.max_memory_usage_bytes(vllm_config), group.kv_cache_spec.page_size_bytes, ) - for group in kv_cache_config.kv_cache_groups - ) - max_concurrency = kv_cache_config.num_blocks / num_blocks_per_request - return max_concurrency + if group.host_resident: + host_blocks_per_request += required + else: + num_blocks_per_request += required + limits = [kv_cache_config.num_blocks / num_blocks_per_request] + if host_blocks_per_request: + assert kv_cache_config.hisparse_host_num_blocks is not None + limits.append( + kv_cache_config.hisparse_host_num_blocks / host_blocks_per_request + ) + return min(limits) def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int: @@ -1575,6 +1601,13 @@ def _get_kv_cache_bytes_per_block( for group in kv_cache_groups ) assert bytes_per_block > 0 + hot_page_sizes = [ + group.kv_cache_spec.page_size_bytes + for group in kv_cache_groups + if isinstance(group.kv_cache_spec, HiSparseHotSpec) + ] + if hot_page_sizes: + bytes_per_block = round_up(bytes_per_block, math.lcm(*hot_page_sizes)) return bytes_per_block @@ -1640,6 +1673,12 @@ def get_kv_cache_config_from_groups( ), ) + if vllm_config.attention_config.hisparse_config is not None: + host_budget = get_hisparse_host_pool_bytes(vllm_config) + return get_hisparse_kv_cache_config( + vllm_config, kv_cache_groups, available_memory, host_budget + ) + if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None: ( attn_group, @@ -2210,6 +2249,9 @@ def get_kv_cache_groups( # attention free models. return [] + if hisparse_groups := get_hisparse_kv_cache_groups(vllm_config, kv_cache_spec): + return hisparse_groups + if is_kv_cache_spec_uniform(kv_cache_spec): # KV cache of all layers are the same, which is true for # most models. Allocate the same amount of memory for @@ -2287,6 +2329,10 @@ def generate_scheduler_kv_cache_config( assert all( [cfg.num_blocks == kv_cache_configs[0].num_blocks for cfg in kv_cache_configs] ) + assert all( + cfg.hisparse_host_num_blocks == kv_cache_configs[0].hisparse_host_num_blocks + for cfg in kv_cache_configs + ) # All workers have the same kv_cache_config except layer names, so use # an arbitrary one to initialize the scheduler. cfg = copy.deepcopy(kv_cache_configs[0]) @@ -2347,6 +2393,9 @@ def _max_memory_usage_bytes_from_groups( if not kv_cache_groups: return 0 + if vllm_config.attention_config.hisparse_config is not None: + return get_hisparse_gpu_memory_usage(vllm_config, kv_cache_groups) + if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None: ( attn_group, @@ -2396,9 +2445,21 @@ def _estimate_max_model_len_from_groups( Returns 0 if even 1 token doesn't fit. """ original_max = vllm_config.model_config.max_model_len + hisparse_enabled = ( + vllm_config.attention_config.hisparse_config is not None + and bool(kv_cache_groups) + ) def fits(model_len: int) -> bool: vllm_config.model_config.max_model_len = model_len + if hisparse_enabled: + try: + config = get_kv_cache_config_from_groups( + vllm_config, kv_cache_groups, available_memory + ) + except ValueError: + return False + return get_max_concurrency_for_kv_cache_config(vllm_config, config) >= 1 return ( _max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups) <= available_memory @@ -2629,6 +2690,9 @@ def get_kv_cache_configs( adjusted_memory.append(override * bytes_per_block) available_memory = adjusted_memory + if vllm_config.attention_config.hisparse_config is not None: + available_memory = [min(available_memory)] * len(available_memory) + # Reserve the null block BlockPool permanently holds back, so auto-fit and # the capacity check both plan against usable blocks. Allocation below # still uses the full memory. diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 93fc2d83c1..3a46fa40a7 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -294,6 +294,9 @@ class SchedulerOutput: # CoW copies to apply after zeroing new blocks and before forward. kv_cache_block_copies: list[KVCacheBlockCopy] | None = None + # Complete block-table rows that replace incrementally appended block IDs. + block_table_updates: dict[str, tuple[list[int], ...]] | None = None + # Scheduler-local; always None by the time this reaches a worker. kv_connector_block_state: KVConnectorBlockState | None = None diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 8b2a551b55..8220373845 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -316,10 +316,9 @@ class Scheduler(SchedulerInterface): self.cache_config.enable_mamba_fine_grained_prefix_cache ), ) - # Bind GPU block pool to the KV connector. This must happen after - # kv_cache_manager is constructed so block_pool is available. + # Bind after construction so connectors can access the cache manager. if self.connector is not None: - self.connector.bind_gpu_block_pool(self.kv_cache_manager.block_pool) + self.connector.bind_kv_cache_manager(self.kv_cache_manager) self.use_pp = self.parallel_config.pipeline_parallel_size > 1 self.use_v2_model_runner = vllm_config.use_v2_model_runner @@ -738,6 +737,12 @@ class Scheduler(SchedulerInterface): # The request can be scheduled. break + if ( + self.connector is not None + and self.connector.has_pending_block_frees() + ): + break + # The request cannot be scheduled. # Preempt the lowest-priority request. if self.policy == SchedulingPolicy.PRIORITY: @@ -1895,14 +1900,20 @@ class Scheduler(SchedulerInterface): outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) spec_decoding_stats: SpecDecodingStats | None = None - failed_kv_load_req_ids = None + failed_kv_load_req_ids: set[str] = set() + if kv_connector_output and kv_connector_output.failed_recving: + failed_kv_load_req_ids.update( + self._handle_failed_recving(kv_connector_output.failed_recving) + ) if kv_connector_output and kv_connector_output.invalid_block_ids: # These blocks contain externally computed tokens that failed to # load. Identify affected requests and adjust their computed token # count to trigger recomputation of the invalid blocks. - failed_kv_load_req_ids = self._handle_invalid_blocks( - kv_connector_output.invalid_block_ids, - num_scheduled_tokens, + failed_kv_load_req_ids.update( + self._handle_invalid_blocks( + kv_connector_output.invalid_block_ids, + num_scheduled_tokens, + ) ) # Persist per-step routed experts into the scheduler-side slot @@ -3135,6 +3146,21 @@ class Scheduler(SchedulerInterface): return affected_req_ids, total_affected_tokens, blocks_to_evict + def _handle_failed_recving(self, failed_req_ids: set[str]) -> set[str]: + """Fail closed for layouts whose block IDs are not globally unique.""" + affected_req_ids: set[str] = set() + for req_id in failed_req_ids: + request = self.requests.get(req_id) + if ( + request is not None + and request.status == RequestStatus.WAITING_FOR_REMOTE_KVS + ): + affected_req_ids.add(req_id) + if self.recompute_kv_load_failures: + request.num_computed_tokens = 0 + self.failed_recving_kv_req_ids.add(req_id) + return set() if self.recompute_kv_load_failures else affected_req_ids + def _handle_invalid_blocks( self, invalid_block_ids: set[int], num_scheduled_tokens: dict[str, int] ) -> set[str]: diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 2015967dd2..d4e68d8adc 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -4,8 +4,9 @@ import itertools from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import Sequence -from typing import ClassVar +from typing import TYPE_CHECKING, ClassVar +from vllm.distributed.kv_events import MEDIUM_CPU from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.v1.core.block_pool import BlockPool @@ -16,6 +17,7 @@ from vllm.v1.core.kv_cache_utils import ( KVCacheBlock, resolve_block_hashes, ) +from vllm.v1.hisparse.block_pool import SharedEventQueueBlockPool from vllm.v1.kv_cache_interface import ( AttentionSpec, ChunkedLocalAttentionSpec, @@ -23,7 +25,10 @@ from vllm.v1.kv_cache_interface import ( CrossAttentionSpec, FullAttentionSpec, HiddenStateCacheSpec, + HiSparseHotSpec, + HiSparseResidentSpec, KpoolTailSpec, + KVCacheGroupRole, KVCacheSpec, MambaSpec, MLAAttentionSpec, @@ -37,6 +42,9 @@ from vllm.v1.kv_cache_interface import ( from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.request import Request +if TYPE_CHECKING: + from vllm.v1.hisparse.coordinator import HiSparseCoordinator + logger = init_logger(__name__) @@ -48,6 +56,9 @@ class SingleTypeKVCacheManager(ABC): supports_fine_grained_hash_lookup: ClassVar[bool] = False + # Keep this group's longer prefix until external cache lookup completes. + retains_longer_hit: bool = False + @property def has_positionally_stable_blocks(self) -> bool: """Whether positional offload scans can follow this block table. @@ -2191,10 +2202,388 @@ class SinkFullAttentionManager(FullAttentionManager): assert sink_len is not None and sink_len > 0 and sink_len % self.block_size == 0 +class HiSparseSourceManager(FullAttentionManager): + """Host-tier manager with a private pool; publishes hashes once durable. + + Host capacity is best effort: a page that cannot get a host block keeps a + null host entry, so its GPU copy is never written back and stays pinned. + """ + + coordinator: "HiSparseCoordinator | None" = None + # The host tier keeps prefixes the device groups have already lost. + retains_longer_hit = True + + @property + def records_new_block_ids(self) -> bool: + return False + + def take_new_block_ids(self) -> list[int]: + self.new_block_ids = [] + return [] + + def bind_host_pool(self, num_blocks: int) -> None: + """Replace the device pool this group was built with by a host one.""" + device_pool = self.block_pool + self.block_pool = SharedEventQueueBlockPool( + num_gpu_blocks=num_blocks, + enable_caching=self.enable_caching and self.kv_cache_spec.prefix_cacheable, + hash_block_size=device_pool.hash_block_size, + enable_kv_cache_events=device_pool.enable_kv_cache_events, + metrics_collector=device_pool.metrics_collector, + medium=MEDIUM_CPU, + event_owner=device_pool, + ) + self._null_block = self.block_pool.null_block + + def take_pending_cow_copies(self) -> list[tuple[KVCacheBlock, KVCacheBlock]]: + """Host copies never reach the worker's generic block-copy path.""" + return [] + + def take_host_cow_copies(self) -> list[tuple[KVCacheBlock, KVCacheBlock]]: + """Drain host copies for the coordinator to hand to the connector.""" + copies = self._pending_cow_copies + self._pending_cow_copies = [] + return copies + + def get_num_blocks_to_allocate( + self, + request_id: str, + num_tokens: int, + new_computed_blocks: Sequence[KVCacheBlock], + total_computed_tokens: int, + num_local_computed_tokens: int, + num_tokens_main_model: int, + apply_admission_cap: bool = False, + ) -> int: + if ( + total_computed_tokens > num_local_computed_tokens + or self._has_partial_local_hit( + new_computed_blocks, num_local_computed_tokens + ) + ): + # External loads need real destinations; future GPU-computed pages + # remain best effort. Use the same admission sentinel as Mamba. + required = super().get_num_blocks_to_allocate( + request_id, + total_computed_tokens, + new_computed_blocks, + total_computed_tokens, + num_local_computed_tokens, + total_computed_tokens, + ) + if required > self.block_pool.get_num_free_blocks(): + assert self.coordinator is not None + assert self.coordinator.gpu_pool is not None + return self.coordinator.gpu_pool.num_gpu_blocks + 1 + return 0 + + def allocate_new_blocks( + self, request_id: str, num_tokens: int, num_tokens_main_model: int + ) -> list[KVCacheBlock]: + req_blocks = self.req_to_blocks[request_id] + num_required = cdiv(num_tokens, self.block_size) + num_free = self.block_pool.get_num_free_blocks() + if request_id in self._partial_hit_reqs: + num_free -= 1 + fit_blocks = min(num_required, len(req_blocks) + max(0, num_free)) + new_blocks = super().allocate_new_blocks( + request_id, + fit_blocks * self.block_size, + min(num_tokens_main_model, fit_blocks * self.block_size), + ) + req_blocks.extend([self._null_block] * (num_required - len(req_blocks))) + return new_blocks + + def allocate_external_computed_blocks( + self, + request_id: str, + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + if num_external_computed_tokens <= 0: + return + # The connector writes these pages; only a successful receive makes + # them readable, not advancing the request's computed-token count. + assert self.coordinator is not None + self.coordinator.record_pending_host_import( + request_id, num_local_computed_tokens + num_external_computed_tokens + ) + super().allocate_external_computed_blocks( + request_id, num_local_computed_tokens, num_external_computed_tokens + ) + + def cache_blocks( + self, + request: Request, + num_tokens: int, + retention_interval: int | None = None, + *, + replay_boundaries: Sequence[int], + ) -> None: + assert self.coordinator is not None + self.coordinator.publish_when_ready( + request, + num_tokens, + retention_interval, + replay_boundaries=replay_boundaries, + ) + + def publish_blocks( + self, + request: Request, + num_tokens: int, + retention_interval: int | None = None, + *, + replay_boundaries: Sequence[int], + ) -> None: + super().cache_blocks( + request, + num_tokens, + retention_interval=retention_interval, + replay_boundaries=replay_boundaries, + ) + + def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: + assert self.coordinator is not None + self.coordinator.free(request_id) + return super().pop_blocks_for_free(request_id) + + +class _HiSparseAuxiliaryManager(SingleTypeKVCacheManager): + """Base for ephemeral groups whose host source owns prefix caching.""" + + coordinator: "HiSparseCoordinator | None" = None + + def __init__(self, kv_cache_spec: KVCacheSpec, **kwargs) -> None: + # Never prefix-cached, but the per-step ``cache_blocks`` hook is where + # residency work runs, so stay opted in regardless of prefix caching. + kwargs["enable_caching"] = True + super().__init__(kv_cache_spec, **kwargs) + + def cache_blocks( + self, + request: Request, + num_tokens: int, + retention_interval: int | None = None, + *, + replay_boundaries: Sequence[int], + ) -> None: + return None + + def get_num_common_prefix_blocks(self, running_request_id: str) -> int: + return 0 + + @classmethod + def find_longest_cache_hit( + cls, + block_hashes: BlockHashList, + max_length: int, + kv_cache_group_ids: list[int], + block_pool: BlockPool, + kv_cache_spec: KVCacheSpec, + drop_eagle_block: bool, + alignment_tokens: int, + dcp_world_size: int = 1, + pcp_world_size: int = 1, + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: + return tuple([] for _ in kv_cache_group_ids), 0 + + +class HiSparseHotManager(_HiSparseAuxiliaryManager): + """Allocate a hot region only after a request acquires CPU-only history.""" + + def __init__(self, kv_cache_spec: HiSparseHotSpec, **kwargs) -> None: + super().__init__(kv_cache_spec, **kwargs) + self.blocks_per_request = kv_cache_spec.blocks_per_request + self.hot_required: set[str] = set() + + def require_hot(self, request_id: str) -> None: + self.hot_required.add(request_id) + + def has_hot(self, request_id: str) -> bool: + return len(self.req_to_blocks.get(request_id, ())) == self.blocks_per_request + + def get_num_blocks_to_allocate( + self, + request_id: str, + num_tokens: int, + new_computed_blocks: Sequence[KVCacheBlock], + total_computed_tokens: int, + num_local_computed_tokens: int, + num_tokens_main_model: int, + apply_admission_cap: bool = False, + ) -> int: + assert not new_computed_blocks + # A hot region is needed to read host-backed history: an external + # import, a new request resuming a host prefix, or one already asked + # to transition. Running requests keep their earlier answer. + host_import = total_computed_tokens > num_local_computed_tokens + resumes_host_prefix = ( + num_local_computed_tokens > 0 and request_id not in self.num_cached_block + ) + if host_import or resumes_host_prefix or request_id in self.hot_required: + return self.get_num_required_blocks(request_id) + return 0 + + def get_num_required_blocks(self, request_id: str) -> int: + return max( + self.blocks_per_request - len(self.req_to_blocks.get(request_id, ())), + 0, + ) + + def add_local_computed_blocks( + self, + request_id: str, + new_computed_blocks: Sequence[KVCacheBlock], + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + assert not new_computed_blocks + self.num_cached_block[request_id] = 0 + assert self.coordinator is not None + if num_local_computed_tokens > 0 or not self.coordinator.resident_managers: + self.require_hot(request_id) + + def allocate_external_computed_blocks( + self, + request_id: str, + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + self.require_hot(request_id) + + def allocate_new_blocks( + self, request_id: str, num_tokens: int, num_tokens_main_model: int + ) -> list[KVCacheBlock]: + # Cold admissions can bypass add_local_computed_blocks. + self.num_cached_block[request_id] = 0 + if request_id not in self.hot_required: + return [] + req_blocks = self.req_to_blocks[request_id] + num_new_blocks = self.blocks_per_request - len(req_blocks) + if num_new_blocks <= 0: + return [] + new_blocks = self.block_pool.get_new_blocks(num_new_blocks) + req_blocks.extend(new_blocks) + return new_blocks + + def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: + self.hot_required.discard(request_id) + return super().pop_blocks_for_free(request_id) + + +class HiSparseResidentManager(_HiSparseAuxiliaryManager): + """Track GPU-resident pages for otherwise host-backed KV.""" + + def get_num_blocks_to_allocate( + self, + request_id: str, + num_tokens: int, + new_computed_blocks: Sequence[KVCacheBlock], + total_computed_tokens: int, + num_local_computed_tokens: int, + num_tokens_main_model: int, + apply_admission_cap: bool = False, + ) -> int: + del num_tokens_main_model + assert not new_computed_blocks + if total_computed_tokens > num_local_computed_tokens: + if total_computed_tokens % self.block_size != 0: + raise ValueError( + "A host-only HiSparse import must end on a cache-block boundary." + ) + imported_pages = total_computed_tokens // self.block_size + return max(cdiv(num_tokens, self.block_size) - imported_pages, 0) + existing = len(self.req_to_blocks.get(request_id, ())) + host_pages = cdiv(num_local_computed_tokens, self.block_size) + required = cdiv(num_tokens, self.block_size) + if apply_admission_cap: + assert self._max_admission_blocks_per_request is not None + required = min(required, self._max_admission_blocks_per_request) + return max(required - max(existing, host_pages), 0) + + def allocate_external_computed_blocks( + self, + request_id: str, + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + """Represent imported host history with null resident pages.""" + assert num_external_computed_tokens > 0 + num_tokens = num_local_computed_tokens + num_external_computed_tokens + blocks = self.req_to_blocks[request_id] + tail_page = cdiv(num_tokens, self.block_size) - 1 + if len(blocks) <= tail_page: + blocks.extend([self._null_block] * (tail_page + 1 - len(blocks))) + + def add_local_computed_blocks( + self, + request_id: str, + new_computed_blocks: Sequence[KVCacheBlock], + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + assert not new_computed_blocks + req_blocks = self.req_to_blocks[request_id] + assert not req_blocks + num_host_pages = cdiv(num_local_computed_tokens, self.block_size) + req_blocks.extend([self._null_block] * num_host_pages) + self.num_cached_block[request_id] = 0 + assert self.coordinator is not None + self.coordinator.commit_computed_blocks(request_id, num_host_pages) + + def cache_blocks( + self, + request: Request, + num_tokens: int, + retention_interval: int | None = None, + *, + replay_boundaries: Sequence[int], + ) -> None: + assert self.coordinator is not None + self.coordinator.plan_prefix_materialization(request.request_id, num_tokens) + self.coordinator.update_residency(request.request_id) + + def allocate_new_blocks( + self, request_id: str, num_tokens: int, num_tokens_main_model: int + ) -> list[KVCacheBlock]: + del num_tokens_main_model + req_blocks = self.req_to_blocks[request_id] + num_new_blocks = cdiv(num_tokens, self.block_size) - len(req_blocks) + if num_new_blocks <= 0: + return [] + new_blocks = self.block_pool.get_new_blocks(num_new_blocks) + req_blocks.extend(new_blocks) + return new_blocks + + def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: + assert self.coordinator is not None + self.coordinator.free(request_id) + return super().pop_blocks_for_free(request_id) + + def adopt_resident_page( + self, request_id: str, block_idx: int, block: KVCacheBlock + ) -> bool: + """Point a null prefix page at a pinned GPU copy of its contents.""" + blocks = self.req_to_blocks.get(request_id) + if blocks is None or block_idx >= len(blocks) or not blocks[block_idx].is_null: + return False + blocks[block_idx] = block + return True + + def get_resident_page(self, request_id: str, block_idx: int) -> KVCacheBlock | None: + blocks = self.req_to_blocks.get(request_id) + if blocks is None or block_idx >= len(blocks): + return None + block = blocks[block_idx] + return None if block.is_null else block + + def get_manager_for_kv_cache_spec( kv_cache_spec: KVCacheSpec, max_in_flight_tokens: int, max_model_len: int, + role: str | None = None, **kwargs, ) -> SingleTypeKVCacheManager: """ @@ -2212,7 +2601,7 @@ def get_manager_for_kv_cache_spec( Returns: An instance of the appropriate SingleTypeKVCacheManager subclass """ - manager_class = KVCacheSpecRegistry.get_manager_class(kv_cache_spec) + manager_class = KVCacheSpecRegistry.get_manager_class(kv_cache_spec, role) assert manager_class is not None, ( f"No manager registered for KVCacheSpec {type(kv_cache_spec)}" ) @@ -2224,7 +2613,7 @@ def get_manager_for_kv_cache_spec( # FullAttentionSpec sizing without a separate admission cap. if isinstance( kv_cache_spec, - (SlidingWindowSpec, ChunkedLocalAttentionSpec), + (SlidingWindowSpec, ChunkedLocalAttentionSpec, HiSparseResidentSpec), ): kwargs["max_admission_blocks_per_request"] = ( kv_cache_spec.max_admission_blocks_per_request( @@ -2238,11 +2627,24 @@ def get_manager_for_kv_cache_spec( def register_all_kvcache_specs(vllm_config): """Built-in spec registration""" + KVCacheSpecRegistry.register_role_manager( + KVCacheGroupRole.HISPARSE_SOURCE, HiSparseSourceManager + ) KVCacheSpecRegistry.register( FullAttentionSpec, FullAttentionManager, uniform_type_base_spec=FullAttentionSpec, ) + KVCacheSpecRegistry.register( + HiSparseHotSpec, + HiSparseHotManager, + uniform_type_base_spec=HiSparseHotSpec, + ) + KVCacheSpecRegistry.register( + HiSparseResidentSpec, + HiSparseResidentManager, + uniform_type_base_spec=HiSparseResidentSpec, + ) KVCacheSpecRegistry.register( SlidingWindowSpec, diff --git a/vllm/v1/hisparse/__init__.py b/vllm/v1/hisparse/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/v1/hisparse/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py new file mode 100644 index 0000000000..263e553d32 --- /dev/null +++ b/vllm/v1/hisparse/binding.py @@ -0,0 +1,277 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Worker-side wiring of HiSparse caches to attention layers.""" + +from collections.abc import Mapping +from copy import copy +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.config import VllmConfig +from vllm.v1.hisparse.layout import ( + HISPARSE_HOT_SUFFIX, + HISPARSE_RESIDENT_SUFFIX, +) +from vllm.v1.hisparse.runtime import ( + HiSparseCacheHandle, + HiSparseHostPool, + initialize_hisparse_runtime_buffers, + release_pinned_state, +) +from vllm.v1.kv_cache_interface import ( + HiSparseHotSpec, + HiSparseResidentSpec, + KVCacheConfig, + KVCacheLayout, + KVCacheSpec, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, + create_kv_cache_views, +) +from vllm.v1.worker.utils import allocate_kv_cache, select_common_block_size + +if TYPE_CHECKING: + from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase + from vllm.v1.worker.gpu.block_table import BlockTables + + +def resolve_hisparse_block_size( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], + attn_layers: Mapping[str, "AttentionLayerBase"], +) -> None: + """Resolve a common kernel block size in-place for HiSparse MLA specs.""" + if vllm_config.attention_config.hisparse_config is None: + return + mla_specs = { + name: spec + for name, spec in kv_cache_spec.items() + if isinstance(spec, MLAAttentionSpec) + } + if not mla_specs: + return + block_sizes = {spec.block_size for spec in mla_specs.values()} + if len(block_sizes) != 1: + raise ValueError("HiSparse requires one scheduler block size.") + backends = [attn_layers[name].get_attn_backend() for name in mla_specs] + try: + block_size = select_common_block_size(block_sizes.pop(), backends) + except ValueError as error: + raise ValueError( + "HiSparse requires a GPU block size supported by every sparse " + f"attention and indexer backend: {error}" + ) from error + kv_cache_spec.update( + (name, spec.copy_with_new_block_size(block_size)) + for name, spec in mla_specs.items() + ) + + +def allocate_hisparse_kv_caches( + kv_cache_config: KVCacheConfig, + device: torch.device, + layout: KVCacheLayout, + kernel_block_sizes: list[int], + host_pool: HiSparseHostPool, +) -> dict[str, torch.Tensor]: + """Allocate the host pool separately from the shared device backing.""" + device_config = copy(kv_cache_config) + device_config.kv_cache_tensors = [ + tensor + for tensor in kv_cache_config.kv_cache_tensors + if not tensor.host_resident + ] + kv_caches = allocate_kv_cache(device_config, device, layout, kernel_block_sizes) + host_tensors = [ + tensor for tensor in kv_cache_config.kv_cache_tensors if tensor.host_resident + ] + (host_size,) = {tensor.size for tensor in host_tensors} + backing = host_pool.allocate(host_size) + (host_group_id,) = kv_cache_config.host_group_ids + host_spec = kv_cache_config.kv_cache_groups[host_group_id].kv_cache_spec + for tensor in host_tensors: + spec = ( + host_spec.kv_cache_specs[tensor.layers[0]] + if isinstance(host_spec, UniformTypeKVCacheSpecs) + else host_spec + ) + kernel_block_size = kernel_block_sizes[host_group_id] + if isinstance(spec, MLAAttentionSpec) and spec.storage_block_size is not None: + kernel_block_size = spec.storage_block_size + views = create_kv_cache_views( + backing, + spec, + kv_cache_config.num_blocks_of(tensor), + layout, + tensor, + kernel_block_size=kernel_block_size, + ) + kv_caches.update(zip(tensor.layers, views)) + return kv_caches + + +def init_hisparse_kv_cache( + kv_cache_config: KVCacheConfig, + device: torch.device, + kernel_block_sizes: list[int], + vllm_config: VllmConfig, + forward_context: dict[str, Any], + block_tables: "BlockTables", +) -> dict[str, torch.Tensor]: + """Allocate and bind HiSparse caches within the caller's allocation context.""" + host_pool = HiSparseHostPool() + kv_caches = allocate_hisparse_kv_caches( + kv_cache_config, + device, + vllm_config.cache_config.get_resolved_kv_cache_layout(), + kernel_block_sizes, + host_pool, + ) + cache_handles = bind_hisparse_kv_caches( + forward_context=forward_context, + kv_cache_config=kv_cache_config, + kv_caches=kv_caches, + block_tables=block_tables, + host_pool=host_pool, + ) + initialize_hisparse_runtime_buffers( + cache_handles, + max_num_reqs=vllm_config.scheduler_config.max_num_seqs, + max_num_batched_tokens=vllm_config.scheduler_config.max_num_batched_tokens, + ) + return kv_caches + + +def _get_hisparse_cache( + forward_context: dict[str, Any], layer_name: str +) -> HiSparseCacheHandle: + attention_layer = forward_context[layer_name] + hisparse_cache = attention_layer.hisparse_cache + assert hisparse_cache is not None + return hisparse_cache + + +def release_hisparse_profiling_cache(forward_context: dict[str, Any]) -> None: + cache_handles = [ + cache + for layer in forward_context.values() + if (cache := getattr(layer, "hisparse_cache", None)) is not None + ] + runtimes = { + id(cache.runtime): cache.runtime + for cache in cache_handles + if hasattr(cache.runtime, "_host_cache") + } + if not runtimes: + return + + registered_pools = list( + { + runtime.registered_host_pool.data_ptr(): runtime.registered_host_pool + for runtime in runtimes.values() + }.values() + ) + release_pinned_state(list(runtimes.values()), registered_pools) + for cache in cache_handles: + cache.mirror_staging_cache = None + cache.mirror_staging_slots = None + + +def bind_hisparse_kv_caches( + *, + forward_context: dict[str, Any], + kv_cache_config: KVCacheConfig, + kv_caches: dict[str, torch.Tensor], + block_tables: "BlockTables", + host_pool: HiSparseHostPool, +) -> list[HiSparseCacheHandle]: + """Bind existing cache storage and block tables; return the bound handles.""" + assert host_pool.registered is not None + tensor_configs = { + name: tensor_config + for tensor_config in kv_cache_config.kv_cache_tensors + for name in tensor_config.layers + } + resident_source_index = 0 + for group_id, group in enumerate(kv_cache_config.kv_cache_groups): + if not isinstance(group.kv_cache_spec, HiSparseResidentSpec): + continue + for cache_name in group.layer_names: + assert cache_name.endswith(HISPARSE_RESIDENT_SUFFIX) + layer_name = cache_name[: -len(HISPARSE_RESIDENT_SUFFIX)] + tensor_config = tensor_configs[cache_name] + assert not tensor_config.host_resident + cache_handle = _get_hisparse_cache(forward_context, layer_name) + cache_handle.bind_cache( + kv_caches[cache_name], + byte_offset=tensor_config.offset, + block_stride=tensor_config.block_stride, + num_blocks=kv_cache_config.num_blocks, + block_size=group.kv_cache_spec.block_size, + block_table=block_tables.input_block_tables[group_id], + slot_mapping=block_tables.slot_mappings[group_id], + ) + assert cache_handle.view is not None + kv_caches[cache_name] = cache_handle.view.cache + cache_handle.runtime.resident_source_index = resident_source_index + resident_source_index += 1 + + hot_backing: torch.Tensor | None = None + cache_handles: list[HiSparseCacheHandle] = [] + for group_id, group in enumerate(kv_cache_config.kv_cache_groups): + if not isinstance(group.kv_cache_spec, HiSparseHotSpec): + continue + for cache_name in group.layer_names: + assert cache_name.endswith(HISPARSE_HOT_SUFFIX) + raw_tensor = kv_caches[cache_name] + if hot_backing is None: + hot_backing = raw_tensor + elif hot_backing.untyped_storage().data_ptr() != ( + raw_tensor.untyped_storage().data_ptr() + ): + raise RuntimeError("HiSparse hot tensors must share one GPU backing.") + layer_name = cache_name[: -len(HISPARSE_HOT_SUFFIX)] + cache_handle = _get_hisparse_cache(forward_context, layer_name) + tensor_config = tensor_configs[cache_name] + assert not tensor_config.host_resident + cache_handle.runtime.bind_hot_cache( + raw_tensor, + byte_offset=tensor_config.offset, + block_stride=tensor_config.block_stride, + num_blocks=kv_cache_config.num_blocks, + block_size=group.kv_cache_spec.block_size, + block_table=block_tables.input_block_tables[group_id], + ) + resident = cache_handle.view + hot = cache_handle.runtime.hot + assert resident is not None + if ( + resident.cache.untyped_storage().data_ptr() + != hot.cache.untyped_storage().data_ptr() + or resident.cache.data_ptr() != hot.cache.data_ptr() + or resident.cache.stride() != hot.cache.stride() + ): + raise RuntimeError("HiSparse resident and hot layouts must match.") + source_cache = kv_caches[layer_name] + assert source_cache.untyped_storage().data_ptr() == ( + host_pool.registered.untyped_storage().data_ptr() + ) + cache_handle.runtime.bind_source_cache( + source_cache, + registered_host_pool=host_pool.registered, + ) + cache_handles.append(cache_handle) + # The hot slab is raw storage owned by the runtime, not a layer + # cache: nothing downstream binds or registers it. + del kv_caches[cache_name] + + if hot_backing is None or not cache_handles: + raise RuntimeError("HiSparse found no hot-cache handles.") + (source_group_id,) = kv_cache_config.host_group_ids + for cache_handle in cache_handles: + cache_handle.source_block_table = block_tables.input_block_tables[ + source_group_id + ] + cache_handle.mirror_slot_mapping = block_tables.slot_mappings[source_group_id] + return cache_handles diff --git a/vllm/v1/hisparse/block_pool.py b/vllm/v1/hisparse/block_pool.py new file mode 100644 index 0000000000..cbb620d647 --- /dev/null +++ b/vllm/v1/hisparse/block_pool.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.distributed.kv_events import KVCacheEvent +from vllm.v1.core.block_pool import BlockPool + + +class SharedEventQueueBlockPool(BlockPool): + """A pool that publishes into another pool's live KV event queue. + + The owner rebinds its queue on every drain, so this reads it through the + owner rather than holding a reference. + """ + + def __init__(self, *args, event_owner: BlockPool, **kwargs) -> None: + self._event_owner = event_owner + super().__init__(*args, **kwargs) + + @property + def kv_event_queue(self) -> list[KVCacheEvent]: + return self._event_owner.kv_event_queue + + @kv_event_queue.setter + def kv_event_queue(self, events: list[KVCacheEvent]) -> None: + # ``BlockPool.__init__`` seeds an empty queue and ``take_events`` + # swaps in a fresh one; both belong to the owner, which drains it. + assert not events diff --git a/vllm/v1/hisparse/coordinator.py b/vllm/v1/hisparse/coordinator.py new file mode 100644 index 0000000000..ecbc0a270a --- /dev/null +++ b/vllm/v1/hisparse/coordinator.py @@ -0,0 +1,752 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from vllm.utils.math_utils import cdiv +from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.kv_cache_utils import ( + BlockHashWithGroupId, + KVCacheBlock, + KVCacheBlockCopy, +) +from vllm.v1.core.single_type_kv_cache_manager import ( + HiSparseHotManager, + HiSparseResidentManager, + HiSparseSourceManager, + SingleTypeKVCacheManager, +) +from vllm.v1.hisparse.types import ( + SparseKVOffloadCommand, + SparseKVPageTransfer, + SparseKVRowMirror, +) +from vllm.v1.kv_cache_interface import ( + HiSparseResidentSpec, + KVCacheConfig, +) +from vllm.v1.request import Request + +if TYPE_CHECKING: + from vllm.v1.core.kv_cache_manager import KVCacheManager + +# Sealed pages this many positions behind the block-table tail stay pinned so +# a page written by an in-flight step is never handed out under it. +_ACTIVE_TAIL_PAGES = 2 + + +@dataclass +class _PendingPublication: + request: Request + num_computed_tokens: int + num_pages: int + retention_interval: int | None + replay_boundaries: Sequence[int] + + +@dataclass +class _HiSparseRequestState: + """Residency bookkeeping for one request. + + A resident page moves through three states: dirty (only the GPU copy + exists), clean (``valid_pages``: the host copy is complete) and unpinned + (``unpinned_pages``: clean, and its allocation reference was released so + the pool may reuse it). ``pinned_clean`` tracks clean pages still holding + their reference, which only happens before the request can read from host. + """ + + valid_pages: set[int] = field(default_factory=set) + ready_prefix_pages: int = 0 + pending_pages: dict[int, int] = field(default_factory=dict) + publication: _PendingPublication | None = None + copies_recorded_blocks: int = 0 + pinned_clean: set[int] = field(default_factory=set) + unpinned_pages: set[int] = field(default_factory=set) + + +@dataclass +class _PendingSpill: + transfer_id: int + page: tuple[str, int] + request_state: _HiSparseRequestState + host_block: KVCacheBlock + resident_blocks: tuple[KVCacheBlock, ...] + expected_worker_completions: int = 0 + worker_completions: int = 0 + enqueue_applied: bool = False + + +class HiSparseCoordinator: + """Own HiSparse host allocation, host publication, spills, and GPU residency. + + GPU-resident pages are a write-back cache of the host tier. Once a page's + host copy is durable and its owner can read from host (it has, or has + asked for, a hot region), the page's allocation reference is released with + ``BlockPool.unpin_blocks``: the block stays in the block table and keeps + being read, but the pool counts it as free and may hand it out, at which + point the owner's page is nulled and its block table republished. + """ + + def __init__( + self, + kv_cache_config: KVCacheConfig, + managers: tuple[SingleTypeKVCacheManager, ...], + max_model_len: int, + ) -> None: + self.managers = managers + self.max_model_len = max_model_len + groups = kv_cache_config.kv_cache_groups + + resident_managers: list[HiSparseResidentManager] = [] + hot_managers: list[HiSparseHotManager] = [] + resident_group_ids: list[int] = [] + for group_id, (group, manager) in enumerate(zip(groups, managers)): + if isinstance(group.kv_cache_spec, HiSparseResidentSpec): + if not isinstance(manager, HiSparseResidentManager): + raise TypeError( + "HiSparse resident specs require resident cache managers." + ) + resident_managers.append(manager) + resident_group_ids.append(group_id) + assert not group.host_resident + if isinstance(manager, HiSparseHotManager): + hot_managers.append(manager) + assert not group.host_resident + if isinstance(manager, (HiSparseHotManager, HiSparseResidentManager)): + manager.coordinator = self + self.resident_managers = tuple(resident_managers) + self.hot_managers = tuple(hot_managers) + self.host_manager: HiSparseSourceManager | None = None + self.host_group_id: int | None = None + self.max_spill_pages = 0 + for group_id, manager in enumerate(managers): + if isinstance(manager, HiSparseSourceManager): + if self.host_manager is not None: + raise ValueError("Only one HiSparse host group is supported.") + num_host_blocks = kv_cache_config.hisparse_host_num_blocks + if num_host_blocks is None: + raise ValueError("HiSparse host group needs host capacity.") + self.host_group_id = group_id + self.host_manager = manager + manager.coordinator = self + manager.bind_host_pool(num_host_blocks) + self.has_host_cache = self.host_manager is not None + self.gpu_pool: BlockPool | None = None + self.transition_watermark = 0 + if self.resident_managers: + assert self.host_manager is not None + assert self.host_group_id is not None + if self.host_group_id > min(resident_group_ids): + raise ValueError("HiSparse host group must precede resident groups.") + resident_block_sizes = { + manager.block_size for manager in self.resident_managers + } + if len(resident_block_sizes) != 1: + raise ValueError( + "HiSparse resident cache groups must use one block size." + ) + resident_block_size = resident_block_sizes.pop() + if self.host_manager.block_size != resident_block_size: + raise ValueError("HiSparse host and resident block sizes must match.") + self.max_spill_pages = max_model_len // resident_block_size + self.gpu_pool = self.resident_managers[0].block_pool + # Requests start reading from host once the shared pool runs this + # low, so admissions find evictable pages rather than pinned ones. + hot_cost = sum(manager.blocks_per_request for manager in hot_managers) + self.transition_watermark = max(hot_cost, kv_cache_config.num_blocks // 10) + + # Host copies handed to the worker, by destination block id, until it + # reports having run them. + self._retained_copies: dict[int, tuple[KVCacheBlock, KVCacheBlock]] = {} + # request -> prefix length an external load is still filling in. + self._pending_imports: dict[str, int] = {} + self.block_table_updates: set[str] = set() + self.spills_to_send: list[SparseKVPageTransfer] = [] + self.pending_spills: dict[int, _PendingSpill] = {} + self.request_states: dict[str, _HiSparseRequestState] = {} + self.next_spill_id = 0 + # Published host block hash -> GPU copies of that page, readable until + # the pool reuses them. Lets a host prefix hit come back GPU-resident. + self.copies: dict[BlockHashWithGroupId, tuple[KVCacheBlock, ...]] = {} + self._copy_by_block: dict[int, BlockHashWithGroupId] = {} + # Unpinned GPU block id -> (request, page) pairs still reading it. + self._owners: dict[int, set[tuple[str, int]]] = {} + + def _get_request_state(self, request_id: str) -> _HiSparseRequestState: + state = self.request_states.get(request_id) + if state is None: + state = _HiSparseRequestState() + self.request_states[request_id] = state + return state + + def commit_computed_blocks(self, request_id: str, num_host_pages: int) -> None: + """Account for a prefix hit on host pages the request now references.""" + if not self.resident_managers or self.host_manager is None: + return + host_blocks = self.host_manager.req_to_blocks.get(request_id, ()) + num_host_pages = min(num_host_pages, len(host_blocks)) + if num_host_pages == 0: + return + state = self._get_request_state(request_id) + state.valid_pages.update(range(num_host_pages)) + state.ready_prefix_pages = max(num_host_pages, state.ready_prefix_pages) + self._adopt_copies(request_id, state, host_blocks[:num_host_pages]) + state.copies_recorded_blocks = max(state.copies_recorded_blocks, num_host_pages) + + def take_host_block_copies(self) -> tuple[KVCacheBlockCopy, ...]: + """Drain host copy-on-write work, retaining both endpoints. + + The worker runs the copies in the step that carries them, so the + endpoints are released when that step's output comes back. + """ + if self.host_manager is None: + return () + pairs = self.host_manager.take_host_cow_copies() + for source_block, cow_block in pairs: + self._retained_copies[cow_block.block_id] = (source_block, cow_block) + return tuple( + KVCacheBlockCopy( + src_block_id=source_block.block_id, + dst_block_id=cow_block.block_id, + ) + for source_block, cow_block in pairs + ) + + def release_completed_host_copies(self, dst_block_ids: Iterable[int]) -> None: + """Release the endpoints of host copies the worker reports having run.""" + blocks: list[KVCacheBlock] = [] + for dst_block_id in dst_block_ids: + pair = self._retained_copies.pop(dst_block_id, None) + if pair is not None: + blocks.extend(pair) + if not blocks: + return + assert self.host_manager is not None + self.host_manager.block_pool.free_blocks(reversed(blocks)) + + def get_host_block_pool(self) -> BlockPool | None: + manager = self.host_manager + return manager.block_pool if manager is not None else None + + # ------------------------------------------------------------------ + # GPU copies of published host pages + # ------------------------------------------------------------------ + + def _adopt_copies( + self, + request_id: str, + state: _HiSparseRequestState, + host_blocks: Sequence[KVCacheBlock], + ) -> None: + """Point null prefix pages at readable GPU copies of their contents. + + Runs after the hit length is fully reconciled, so copy hits and misses + can be scattered per page without affecting the prefix hit. + """ + if not self.copies: + return + for host_idx, host_block in enumerate(host_blocks): + if host_block.is_null or host_block.block_hash is None: + continue + blocks = self.copies.get(host_block.block_hash) + if blocks is None: + continue + for manager, block in zip(self.resident_managers, blocks): + if manager.adopt_resident_page(request_id, host_idx, block): + manager.block_pool.touch([block]) + state.pinned_clean.add(host_idx) + + def _record_copies(self, request_id: str, num_computed_tokens: int) -> None: + """Index the GPU copies of just-published host blocks for later hits.""" + if not self.resident_managers: + return + assert self.host_manager is not None + host_blocks = self.host_manager.req_to_blocks.get(request_id) + if not host_blocks: + return + state = self._get_request_state(request_id) + num_host_blocks = min( + num_computed_tokens // self.host_manager.block_size, len(host_blocks) + ) + for host_idx in range(state.copies_recorded_blocks, num_host_blocks): + host_block = host_blocks[host_idx] + if host_block.is_null or host_block.block_hash is None: + continue + blocks: list[KVCacheBlock] = [] + for manager in self.resident_managers: + block = manager.get_resident_page(request_id, host_idx) + if block is None: + break + blocks.append(block) + if len(blocks) != len(self.resident_managers): + continue + self._drop_copy(host_block.block_hash) + self.copies[host_block.block_hash] = tuple(blocks) + for block in blocks: + self._copy_by_block[block.block_id] = host_block.block_hash + state.copies_recorded_blocks = max( + state.copies_recorded_blocks, num_host_blocks + ) + + def _drop_copy(self, host_hash: BlockHashWithGroupId) -> None: + blocks = self.copies.pop(host_hash, None) + if blocks is None: + return + for block in blocks: + self._copy_by_block.pop(block.block_id, None) + + # ------------------------------------------------------------------ + # Residency: pinning, unpinning and losing pages + # ------------------------------------------------------------------ + + def _can_read_from_host(self, request_id: str) -> bool: + """Whether resident pages may be dropped under the request.""" + return bool(self.hot_managers) and all( + manager.has_hot(request_id) for manager in self.hot_managers + ) + + def _resident_page_blocks( + self, request_id: str, page_idx: int + ) -> list[KVCacheBlock] | None: + blocks: list[KVCacheBlock] = [] + for manager in self.resident_managers: + req_blocks = manager.req_to_blocks.get(request_id) + if ( + req_blocks is None + or page_idx >= len(req_blocks) - _ACTIVE_TAIL_PAGES + or req_blocks[page_idx].is_null + ): + return None + blocks.append(req_blocks[page_idx]) + return blocks + + def _unpin_page( + self, request_id: str, state: _HiSparseRequestState, page_idx: int + ) -> bool: + if page_idx in state.pending_pages or page_idx not in state.valid_pages: + return False + blocks = self._resident_page_blocks(request_id, page_idx) + if blocks is None: + return False + for manager, block in zip(self.resident_managers, blocks): + self._owners.setdefault(block.block_id, set()).add((request_id, page_idx)) + manager.block_pool.unpin_blocks([block], self._on_block_reused) + state.pinned_clean.discard(page_idx) + state.unpinned_pages.add(page_idx) + return True + + def _unpin_clean_pages(self, request_id: str, state: _HiSparseRequestState) -> None: + for page_idx in sorted(state.pinned_clean): + self._unpin_page(request_id, state, page_idx) + + def _page_became_clean( + self, request_id: str, state: _HiSparseRequestState, page_idx: int + ) -> None: + state.pinned_clean.add(page_idx) + if self._can_read_from_host(request_id): + self._unpin_page(request_id, state, page_idx) + + def _on_block_reused(self, block: KVCacheBlock) -> None: + host_hash = self._copy_by_block.get(block.block_id) + if host_hash is not None: + self._drop_copy(host_hash) + for request_id, page_idx in self._owners.pop(block.block_id, ()): + self._lose_page(request_id, page_idx) + + def _lose_page(self, request_id: str, page_idx: int) -> None: + state = self.request_states.get(request_id) + if state is None: + return + for manager in self.resident_managers: + blocks = manager.req_to_blocks.get(request_id) + if blocks is None or page_idx >= len(blocks) or blocks[page_idx].is_null: + continue + block = blocks[page_idx] + blocks[page_idx] = manager._null_block + owners = self._owners.get(block.block_id) + if owners is not None: + owners.discard((request_id, page_idx)) + state.unpinned_pages.discard(page_idx) + self.block_table_updates.add(request_id) + for hot_manager in self.hot_managers: + hot_manager.require_hot(request_id) + + def update_residency(self, request_id: str) -> None: + """Per-step residency policy for a scheduled request. + + A request that can read from host releases every clean sealed page to + the pool. One that cannot keeps its pages pinned until the shared pool + runs low, then asks for a hot region. Pages remain pinned until that region + is allocated on a subsequent scheduling pass. + """ + if not self.resident_managers: + return + state = self._get_request_state(request_id) + if not self._can_read_from_host(request_id): + assert self.gpu_pool is not None + if self.gpu_pool.get_num_free_blocks() >= self.transition_watermark: + return + for manager in self.hot_managers: + manager.require_hot(request_id) + return + self._unpin_clean_pages(request_id, state) + + # ------------------------------------------------------------------ + # Host publication and spills + # ------------------------------------------------------------------ + + def plan_prefix_materialization( + self, request_id: str, num_computed_tokens: int + ) -> None: + """Eagerly write back every sealed page that has no host copy yet.""" + if not self.resident_managers: + return + assert self.host_manager is not None + host_block_size = self.host_manager.block_size + num_pages = num_computed_tokens // host_block_size + state = self._get_request_state(request_id) + budget = max(self.max_spill_pages - len(self.spills_to_send), 0) + for page_idx in range(num_pages): + if budget == 0: + break + if page_idx in state.pending_pages: + continue + if page_idx in state.valid_pages: + continue + if not all( + manager.get_resident_page(request_id, page_idx) is not None + for manager in self.resident_managers + ): + continue + if self._plan_spill(request_id, page_idx, after_forward=True): + budget -= 1 + + def publish_when_ready( + self, + request: Request, + num_computed_tokens: int, + retention_interval: int | None, + *, + replay_boundaries: Sequence[int], + ) -> None: + """Publish host-source hashes only after their pages are durable.""" + manager = self.host_manager + if manager is None: + return + num_pages = num_computed_tokens // manager.block_size + request_id = request.request_id + state = self._get_request_state(request_id) + if state.ready_prefix_pages >= num_pages: + manager.publish_blocks( + request, + num_computed_tokens, + retention_interval=retention_interval, + replay_boundaries=replay_boundaries, + ) + self._record_copies(request_id, num_computed_tokens) + state.publication = None + return + state.publication = _PendingPublication( + request=request, + num_computed_tokens=num_computed_tokens, + num_pages=num_pages, + retention_interval=retention_interval, + replay_boundaries=replay_boundaries, + ) + + def _publish_host_blocks_if_ready(self, request_id: str) -> None: + state = self.request_states.get(request_id) + if state is None: + return + publication = state.publication + if publication is None or state.ready_prefix_pages < publication.num_pages: + return + assert self.host_manager is not None + self.host_manager.publish_blocks( + publication.request, + publication.num_computed_tokens, + retention_interval=publication.retention_interval, + replay_boundaries=publication.replay_boundaries, + ) + self._record_copies(request_id, publication.num_computed_tokens) + state.publication = None + + def record_pending_host_import(self, request_id: str, num_tokens: int) -> None: + """Note a prefix an external load is populating in host pages.""" + self._pending_imports[request_id] = num_tokens + + def finish_host_import(self, request_id: str, *, failed: bool) -> None: + num_tokens = self._pending_imports.pop(request_id, None) + if num_tokens is not None and not failed: + self.complete_host_import(request_id, num_tokens) + + def complete_host_import(self, request_id: str, num_computed_tokens: int) -> None: + """Publish externally populated host pages after connector completion.""" + if not self.resident_managers: + return + block_size = self.resident_managers[0].block_size + num_pages = num_computed_tokens // block_size + state = self._get_request_state(request_id) + state.valid_pages = set(range(num_pages)) + state.ready_prefix_pages = num_pages + self._publish_host_blocks_if_ready(request_id) + + def _plan_spill( + self, + request_id: str, + page_idx: int, + *, + after_forward: bool, + ) -> bool: + assert self.host_manager is not None + state = self._get_request_state(request_id) + host_blocks = self.host_manager.req_to_blocks.get(request_id) + if host_blocks is None or page_idx >= len(host_blocks): + return False + host_block = host_blocks[page_idx] + if host_block.is_null: + return False + + blocks: list[KVCacheBlock] = [] + for manager in self.resident_managers: + block = manager.get_resident_page(request_id, page_idx) + if block is None: + return False + blocks.append(block) + + self.host_manager.block_pool.touch([host_block]) + for manager, block in zip(self.resident_managers, blocks): + manager.block_pool.touch([block]) + spill_id = self.next_spill_id + self.next_spill_id += 1 + plan = SparseKVPageTransfer( + transfer_id=spill_id, + destination_block_id=host_block.block_id, + source_block_ids=tuple(block.block_id for block in blocks), + after_forward=after_forward, + ) + self.pending_spills[spill_id] = _PendingSpill( + transfer_id=spill_id, + page=(request_id, page_idx), + request_state=state, + host_block=host_block, + resident_blocks=tuple(blocks), + ) + state.pending_pages[page_idx] = spill_id + self.spills_to_send.append(plan) + return True + + # ------------------------------------------------------------------ + # Worker-facing views + # ------------------------------------------------------------------ + + def build_row_mirrors( + self, + requests: Iterable[tuple[str, int, int]], + ) -> tuple[SparseKVRowMirror, ...]: + """Map scheduled rows to scheduler-owned resident and host blocks.""" + if not self.resident_managers or self.host_manager is None: + return () + block_size = self.resident_managers[0].block_size + mirrors: list[SparseKVRowMirror] = [] + for request_id, num_computed_tokens, num_scheduled_tokens in requests: + host_blocks = self.host_manager.req_to_blocks.get(request_id) + resident_blocks = [ + manager.req_to_blocks.get(request_id) + for manager in self.resident_managers + ] + if host_blocks is None or any(blocks is None for blocks in resident_blocks): + continue + token_position = num_computed_tokens + end_position = token_position + num_scheduled_tokens + while token_position < end_position: + page_idx, row_offset = divmod(token_position, block_size) + num_rows = min(block_size - row_offset, end_position - token_position) + # A page without a GPU copy has no mirror source, and a page + # without a host block has no destination. Skip just that page: + # its neighbours in the window are still mirrorable, and a + # window only ever moves forward, so dropping them here would + # leave their host rows permanently stale. + source_starts = [] + for blocks in resident_blocks: + assert blocks is not None + if page_idx >= len(blocks) or blocks[page_idx].is_null: + break + source_starts.append( + blocks[page_idx].block_id * block_size + row_offset + ) + if len(source_starts) != len(resident_blocks): + token_position += num_rows + continue + host_block_idx = page_idx + if host_block_idx >= len(host_blocks): + token_position += num_rows + continue + host_block = host_blocks[host_block_idx] + if host_block.is_null: + token_position += num_rows + continue + mirrors.append( + SparseKVRowMirror( + source_starts=tuple(source_starts), + destination_start=( + host_block.block_id * block_size + row_offset + ), + num_rows=num_rows, + ) + ) + token_position += num_rows + return tuple(mirrors) + + def all_context_pages_resident( + self, + requests: Iterable[tuple[str, int, int]], + ) -> bool: + """Return whether every scheduled request can read only resident KV.""" + if not self.resident_managers: + return False + block_size = self.resident_managers[0].block_size + for request_id, num_computed_tokens, num_scheduled_tokens in requests: + num_pages = cdiv(num_computed_tokens + num_scheduled_tokens, block_size) + for page_idx in range(num_pages): + if any( + manager.get_resident_page(request_id, page_idx) is None + for manager in self.resident_managers + ): + return False + return True + + def take_block_table_updates(self) -> dict[str, tuple[list[int], ...]]: + updates = { + request_id: tuple( + [block.block_id for block in manager.req_to_blocks.get(request_id, [])] + for manager in self.managers + ) + for request_id in self.block_table_updates + } + self.block_table_updates.clear() + return updates + + def build_offload_command(self) -> SparseKVOffloadCommand | None: + """Package residency decisions for the worker connector.""" + if not self.resident_managers: + return None + command = SparseKVOffloadCommand(page_transfers=self.spills_to_send) + self.spills_to_send = [] + return command + + def has_pending_reclamation(self) -> bool: + """Whether an in-flight spill will free GPU blocks on completion.""" + return any( + pending.page[0] in self.request_states + and self._can_read_from_host(pending.page[0]) + for pending in self.pending_spills.values() + ) + + def has_pending_work(self) -> bool: + return bool(self.spills_to_send or self.pending_spills or self._retained_copies) + + def update_spills( + self, + enqueued_counts: Mapping[int, int], + completed_counts: Mapping[int, int], + ) -> None: + for spill_id, count in enqueued_counts.items(): + pending = self.pending_spills.get(spill_id) + if pending is not None and pending.expected_worker_completions == 0: + pending.expected_worker_completions = count + for spill_id, count in completed_counts.items(): + pending = self.pending_spills.get(spill_id) + if pending is not None: + pending.worker_completions += count + + self._apply_enqueued_spills() + self._complete_host_writes() + + def _apply_enqueued_spills(self) -> None: + """Drop the DMA pins once the worker has enqueued a spill.""" + for pending in self.pending_spills.values(): + if ( + pending.enqueue_applied + or not pending.expected_worker_completions + or pending.worker_completions < pending.expected_worker_completions + ): + continue + for manager, block in zip(self.resident_managers, pending.resident_blocks): + manager.block_pool.free_blocks([block]) + pending.enqueue_applied = True + + def _complete_host_writes(self) -> None: + completed = [ + pending + for pending in self.pending_spills.values() + if pending.enqueue_applied + and pending.expected_worker_completions + and pending.worker_completions >= pending.expected_worker_completions + ] + completed_request_ids: set[str] = set() + for pending in completed: + self.pending_spills.pop(pending.transfer_id, None) + request_id, page_idx = pending.page + pending.request_state.pending_pages.pop(page_idx, None) + if self.request_states.get(request_id) is pending.request_state: + pending.request_state.valid_pages.add(page_idx) + self._page_became_clean(request_id, pending.request_state, page_idx) + completed_request_ids.add(request_id) + assert self.host_manager is not None + self.host_manager.block_pool.free_blocks([pending.host_block]) + for request_id in completed_request_ids: + state = self.request_states[request_id] + while state.ready_prefix_pages in state.valid_pages: + state.ready_prefix_pages += 1 + self._publish_host_blocks_if_ready(request_id) + + def free(self, request_id: str) -> None: + """Detach the request; its clean pages stay readable copies in the pool.""" + self._pending_imports.pop(request_id, None) + state = self.request_states.pop(request_id, None) + if state is None: + return + for manager in self.resident_managers: + blocks = manager.req_to_blocks.get(request_id) + if not blocks: + continue + for page_idx, block in enumerate(blocks): + if block.is_null: + continue + if page_idx in state.unpinned_pages: + owners = self._owners.get(block.block_id) + if owners is not None: + owners.discard((request_id, page_idx)) + elif ( + page_idx in state.valid_pages + and page_idx not in state.pending_pages + ): + manager.block_pool.unpin_blocks([block], self._on_block_reused) + else: + continue + blocks[page_idx] = manager._null_block + + +def get_hisparse_coordinator( + kv_cache_manager: "KVCacheManager", +) -> HiSparseCoordinator: + """Return the coordinator shared by a KV cache manager's HiSparse groups. + + Built on first call and memoised on the HiSparse managers it binds itself + to. Must run before the first request is admitted. + """ + managers = tuple(kv_cache_manager.coordinator.single_type_managers) + for manager in managers: + coordinator = getattr(manager, "coordinator", None) + if coordinator is not None: + assert isinstance(coordinator, HiSparseCoordinator) + return coordinator + coordinator = HiSparseCoordinator( + kv_cache_manager.kv_cache_config, managers, kv_cache_manager.max_model_len + ) + if coordinator.host_manager is None: + raise ValueError("No HiSparse cache group is configured.") + return coordinator diff --git a/vllm/v1/hisparse/layout.py b/vllm/v1/hisparse/layout.py new file mode 100644 index 0000000000..d74b0e8cfd --- /dev/null +++ b/vllm/v1/hisparse/layout.py @@ -0,0 +1,345 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections import defaultdict +from dataclasses import dataclass +from typing import cast + +from vllm.config import VllmConfig +from vllm.config.kv_transfer import hisparse_host_pool_gib +from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv +from vllm.v1.hisparse.runtime import ResolvedHiSparseConfig +from vllm.v1.kv_cache_interface import ( + HiSparseHotSpec, + HiSparseResidentSpec, + KVCacheConfig, + KVCacheGroupRole, + KVCacheGroupSpec, + KVCacheSpec, + KVCacheTensor, + MLAAttentionSpec, + SparseCacheRole, + UniformTypeKVCacheSpecs, + compute_layout_strides, +) +from vllm.v1.kv_cache_layout import KVCacheLayout + +logger = init_logger(__name__) + +HISPARSE_HOT_SUFFIX = ".hisparse_hot" +HISPARSE_RESIDENT_SUFFIX = ".hisparse_resident" + + +@dataclass(frozen=True) +class HiSparseLayout: + source_group: KVCacheGroupSpec + device_groups: list[KVCacheGroupSpec] + host_num_blocks: int + + +def get_hisparse_kv_cache_groups( + vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec] +) -> list[KVCacheGroupSpec] | None: + attention_config = getattr(vllm_config, "attention_config", None) + if attention_config is None or attention_config.hisparse_config is None: + return None + + mla_specs: dict[str, KVCacheSpec] = { + name: spec + for name, spec in kv_cache_spec.items() + if isinstance(spec, MLAAttentionSpec) + } + other_specs = { + name: spec + for name, spec in kv_cache_spec.items() + if not isinstance(spec, MLAAttentionSpec) + } + if not mla_specs or not other_specs: + return None + + from vllm.v1.core.kv_cache_utils import get_kv_cache_groups + + mla_group_spec = UniformTypeKVCacheSpecs.from_specs(mla_specs) + assert mla_group_spec is not None + mla_group = KVCacheGroupSpec(list(mla_specs), mla_group_spec) + return [mla_group, *get_kv_cache_groups(vllm_config, other_specs)] + + +def get_hisparse_host_pool_bytes(vllm_config: VllmConfig) -> int: + host_pool_gib = hisparse_host_pool_gib(vllm_config.kv_transfer_config) + if host_pool_gib is None: + raise ValueError("HiSparse requires HiSparseConnector with host_pool_gib") + return int(host_pool_gib * 2**30) + + +def _partition_hisparse_specs( + groups: list[KVCacheGroupSpec], +) -> tuple[dict[str, MLAAttentionSpec], dict[str, MLAAttentionSpec]]: + group_spec = groups[0].kv_cache_spec + if not isinstance(group_spec, UniformTypeKVCacheSpecs): + raise ValueError("HiSparse requires uniform sparse-MLA cache specs.") + + specs = group_spec.kv_cache_specs + if any( + isinstance(spec, MLAAttentionSpec) and spec.model_version == "deepseek_v4" + for spec in specs.values() + ): + raise ValueError("HiSparse does not support DeepSeek V4.") + if not all(isinstance(spec, MLAAttentionSpec) for spec in specs.values()): + raise ValueError("HiSparse requires its first cache group to contain MLA only.") + + source_specs = { + name: spec + for name, spec in specs.items() + if isinstance(spec, MLAAttentionSpec) + and spec.cache_role is SparseCacheRole.SPARSE + } + indexer_specs = { + name: spec + for name, spec in specs.items() + if isinstance(spec, MLAAttentionSpec) + and spec.cache_role is SparseCacheRole.INDEXER + } + if not source_specs or not indexer_specs: + raise ValueError("HiSparse requires sparse-MLA and indexer cache specs.") + return source_specs, indexer_specs + + +def get_hisparse_gpu_memory_usage( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], +) -> int: + _, indexer_specs = _partition_hisparse_specs(kv_cache_groups) + return sum( + spec.max_memory_usage_bytes(vllm_config) for spec in indexer_specs.values() + ) + sum( + group.kv_cache_spec.max_memory_usage_bytes(vllm_config) + for group in kv_cache_groups[1:] + ) + + +def create_hisparse_layout( + vllm_config: VllmConfig, + groups: list[KVCacheGroupSpec], + host_budget: int, +) -> HiSparseLayout: + source_specs, indexer_specs = _partition_hisparse_specs(groups) + block_sizes = { + spec.block_size for spec in (*source_specs.values(), *indexer_specs.values()) + } + if len(block_sizes) != 1: + raise ValueError("HiSparse requires one resolved GPU block size.") + gpu_block_size = block_sizes.pop() + + config = ResolvedHiSparseConfig.from_vllm_config( + vllm_config, + vllm_config.model_config.hf_config.index_topk, + gpu_block_size, + ) + assert config is not None + indexer_group_spec = UniformTypeKVCacheSpecs.from_specs( + cast(dict[str, KVCacheSpec], indexer_specs) + ) + assert indexer_group_spec is not None + indexer_group = KVCacheGroupSpec( + list(indexer_specs), + indexer_group_spec, + enable_kv_transfer=True, + role=KVCacheGroupRole.HISPARSE_INDEXER, + ) + + indexer_page = sum(spec.page_size_bytes for spec in indexer_specs.values()) + hot_blocks_per_request = cdiv(config.device_buffer_size, gpu_block_size) + hot_units: list[list[tuple[str, MLAAttentionSpec]]] = [] + for layer_name, layer_spec in source_specs.items(): + if layer_spec.is_index_group_leader or not hot_units: + hot_units.append([]) + hot_units[-1].append((f"{layer_name}{HISPARSE_HOT_SUFFIX}", layer_spec)) + + resident_groups: list[KVCacheGroupSpec] = [] + hot_groups: list[KVCacheGroupSpec] = [] + + def append_hot_group(layers: list[tuple[str, MLAAttentionSpec]]) -> None: + page_sizes = {spec.page_size_bytes for _, spec in layers} + if len(page_sizes) != 1: + raise ValueError( + "HiSparse hot-cache groups require one page size, got " + f"{sorted(page_sizes)}." + ) + page_size = page_sizes.pop() + names = [name for name, _ in layers] + resident_groups.append( + KVCacheGroupSpec( + [ + name[: -len(HISPARSE_HOT_SUFFIX)] + HISPARSE_RESIDENT_SUFFIX + for name in names + ], + HiSparseResidentSpec( + block_size=gpu_block_size, + page_size=page_size, + ), + enable_kv_transfer=False, + ) + ) + hot_groups.append( + KVCacheGroupSpec( + names, + HiSparseHotSpec( + block_size=gpu_block_size, + page_size=page_size, + blocks_per_request=hot_blocks_per_request, + ), + enable_kv_transfer=False, + ) + ) + + current: list[tuple[str, MLAAttentionSpec]] = [] + current_page = 0 + for unit in hot_units: + unit_page = sum(spec.page_size_bytes for _, spec in unit) + if current and current_page + unit_page > indexer_page: + append_hot_group(current) + current = [] + current_page = 0 + current.extend(unit) + current_page += unit_page + if current: + append_hot_group(current) + + source_group_spec = UniformTypeKVCacheSpecs.from_specs( + cast(dict[str, KVCacheSpec], source_specs) + ) + assert source_group_spec is not None + source_group = KVCacheGroupSpec( + list(source_specs), + source_group_spec, + role=KVCacheGroupRole.HISPARSE_SOURCE, + host_resident=True, + enable_kv_transfer=True, + ) + regular_groups = groups[1:] + gpu_groups = [indexer_group, *resident_groups, *hot_groups, *regular_groups] + + host_num_blocks = host_budget // sum( + spec.page_size_bytes for spec in source_specs.values() + ) + if host_num_blocks <= 0: + raise ValueError("HiSparse has no allocatable host blocks.") + + return HiSparseLayout( + source_group=source_group, + device_groups=gpu_groups, + host_num_blocks=host_num_blocks, + ) + + +def _build_hisparse_kv_cache_tensors( + kv_cache_groups: list[KVCacheGroupSpec], + num_blocks: int, + size: int, + layout: KVCacheLayout, + bytes_per_block: int, + *, + host_resident: bool = False, +) -> list[KVCacheTensor]: + interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None + tensors: list[KVCacheTensor] = [] + for group in kv_cache_groups: + group_spec = group.kv_cache_spec + layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list) + if isinstance(group_spec, UniformTypeKVCacheSpecs): + for layer_name, spec in group_spec.kv_cache_specs.items(): + layers_by_spec[spec].append(layer_name) + elif group.layer_names: + layers_by_spec[group_spec].extend(group.layer_names) + + byte_offset = 0 + for spec, layer_names in layers_by_spec.items(): + if isinstance(spec, (HiSparseHotSpec, HiSparseResidentSpec)): + if not layout.is_block_outermost: + raise ValueError("HiSparse requires a block-outermost KV layout.") + layer_stride = spec.page_size_bytes + block_stride = bytes_per_block + else: + layer_stride, block_stride, _, _, _ = compute_layout_strides( + spec, + num_blocks, + len(layer_names), + layout, + fixed_strides=(None, interleaved_block_stride, None, None, None), + ) + offset = ( + byte_offset + * max(layer_stride, spec.page_size_bytes) + // spec.page_size_bytes + ) + tensors.append( + KVCacheTensor( + size=size, + layers=layer_names, + layer_stride=layer_stride, + block_stride=block_stride, + offset=offset, + host_resident=host_resident, + ) + ) + byte_offset += len(layer_names) * spec.page_size_bytes + return tensors + + +def get_hisparse_kv_cache_config( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], + available_memory: int, + host_budget: int, +) -> KVCacheConfig: + from vllm.v1.core.kv_cache_utils import ( + _get_kv_cache_bytes_per_block, + may_override_num_blocks, + validate_kv_cache_layout, + ) + + hisparse_layout = create_hisparse_layout(vllm_config, kv_cache_groups, host_budget) + device_groups = hisparse_layout.device_groups + layout = vllm_config.cache_config.get_resolved_kv_cache_layout() + validate_kv_cache_layout(layout, device_groups) + bytes_per_block = _get_kv_cache_bytes_per_block(device_groups) + num_blocks = may_override_num_blocks( + vllm_config, available_memory // bytes_per_block + ) + size = bytes_per_block * num_blocks + kv_cache_tensors = _build_hisparse_kv_cache_tensors( + device_groups, num_blocks, size, layout, bytes_per_block + ) + + host_groups = [hisparse_layout.source_group] + host_layout = KVCacheLayout.LBNHC + validate_kv_cache_layout(host_layout, host_groups) + host_bytes_per_block = _get_kv_cache_bytes_per_block(host_groups) + host_size = host_bytes_per_block * hisparse_layout.host_num_blocks + kv_cache_tensors[:0] = _build_hisparse_kv_cache_tensors( + host_groups, + hisparse_layout.host_num_blocks, + host_size, + host_layout, + host_bytes_per_block, + host_resident=True, + ) + logger.info_once( + "HiSparse HMA: %.1f GiB host source (%d blocks), %.1f GiB shared " + "GPU indexer/resident/hot pool (%d blocks).", + host_size / 2**30, + hisparse_layout.host_num_blocks, + size / 2**30, + num_blocks, + ) + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=kv_cache_tensors, + kv_cache_groups=[*host_groups, *device_groups], + hisparse_host_num_blocks=hisparse_layout.host_num_blocks, + prefix_cache_retention_interval=( + vllm_config.cache_config.prefix_cache_retention_interval + ), + ) diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py new file mode 100644 index 0000000000..96235f1570 --- /dev/null +++ b/vllm/v1/hisparse/runtime.py @@ -0,0 +1,1100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVIDIA HiSparse resident, host, and hot-cache data plane.""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, TypeAlias + +import psutil +import torch + +from vllm.config import CUDAGraphMode, VllmConfig +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.utils.math_utils import round_up +from vllm.utils.torch_utils import current_stream +from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor + +logger = init_logger(__name__) + +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup + +# fp8_ds_mla KV row: 512 B quantized NoPE + 16 B scales + 128 B RoPE. +FP8_DS_MLA_ROW_BYTES = 656 +HiSparseTopKResult: TypeAlias = torch.Tensor | tuple[torch.Tensor, torch.Tensor] + + +def is_hisparse_decode_batch(attn_metadata: Any | None) -> bool: + return ( + attn_metadata is not None + and attn_metadata.num_decode_tokens == attn_metadata.num_actual_tokens + and attn_metadata.max_query_len == 1 + and attn_metadata.num_reqs == attn_metadata.num_actual_tokens + ) + + +def _get_max_decode_query_len(vllm_config: VllmConfig) -> int: + speculative_config = getattr(vllm_config, "speculative_config", None) + if ( + speculative_config is not None + and speculative_config.num_speculative_tokens is not None + ): + return 1 + speculative_config.num_speculative_tokens * ( + 2 if speculative_config.parallel_drafting else 1 + ) + return 1 + + +def _get_max_swap_rows(vllm_config: VllmConfig) -> int: + max_query_len = _get_max_decode_query_len(vllm_config) + scheduler_config = vllm_config.scheduler_config + return min( + scheduler_config.max_num_batched_tokens, + scheduler_config.max_num_seqs * max_query_len, + ) + + +@dataclass(frozen=True) +class PagedCacheView: + cache: torch.Tensor + attention_cache: torch.Tensor + block_size: int + attention_block_stride: int + + @classmethod + def bind( + cls, + raw_tensor: torch.Tensor, + *, + dtype: torch.dtype, + row_width: int, + byte_offset: int, + block_stride: int, + num_blocks: int, + block_size: int, + ) -> PagedCacheView: + itemsize = dtype.itemsize + assert byte_offset % itemsize == 0 and block_stride % itemsize == 0 + cache = torch.as_strided( + raw_tensor.view(dtype), + size=(num_blocks, block_size, row_width), + stride=(block_stride // itemsize, row_width, 1), + storage_offset=byte_offset // itemsize, + ) + row_bytes = row_width * itemsize + if block_stride % (block_size * row_bytes): + return cls(cache, cache, block_size, block_size) + attention_cache = ( + raw_tensor[byte_offset:].view(dtype).view(-1, block_size, row_width) + ) + return cls(cache, attention_cache, block_size, block_stride // row_bytes) + + +@dataclass(frozen=True) +class ResolvedHiSparseConfig: + top_k: int + device_buffer_size: int + eager_host_mirror: bool = True + + @classmethod + def from_vllm_config( + cls, + vllm_config: VllmConfig, + model_top_k: int, + block_size: int | None = None, + ) -> ResolvedHiSparseConfig | None: + config = vllm_config.attention_config.hisparse_config + if config is None: + return None + + max_query_len = _get_max_decode_query_len(vllm_config) + min_device_buffer_size = max_query_len * model_top_k + # Retain the worst-case union of every speculative step's top-k and + # leave one additional top-k of LRU slack. + configured_size = config.device_buffer_size + if configured_size is None: + device_buffer_size = (max_query_len + 1) * model_top_k + else: + device_buffer_size = configured_size + + if device_buffer_size < min_device_buffer_size: + raise ValueError( + "HiSparse device_buffer_size must cover every decode query's " + "index_topk rows. Got " + f"device_buffer_size={device_buffer_size}, " + f"max_decode_query_len={max_query_len}, " + f"index_topk={model_top_k}; expected at least " + f"{min_device_buffer_size}." + ) + max_device_buffer_size = torch.iinfo(torch.int16).max + 1 + if device_buffer_size > max_device_buffer_size: + raise ValueError( + "HiSparse device_buffer_size exceeds the int16 slot-index " + f"limit: got {device_buffer_size}, maximum is " + f"{max_device_buffer_size}." + ) + if configured_size is not None and block_size is not None: + padding = -device_buffer_size % block_size + if padding: + logger.warning( + "HiSparse device_buffer_size=%d is not aligned to the " + "%d-token kernel block size and allocates %d unused " + "padding rows per hot group. Use %d to use all allocated " + "rows.", + device_buffer_size, + block_size, + padding, + device_buffer_size + padding, + ) + return cls( + top_k=model_top_k, + device_buffer_size=device_buffer_size, + eager_host_mirror=config.eager_host_mirror, + ) + + +def check_hisparse_host_memory(pool_bytes: int) -> None: + """Reject a physical host-pool allocation that cannot fit in RAM.""" + mem = psutil.virtual_memory() + if pool_bytes > mem.available * 0.95: + raise ValueError( + f"HiSparse pinned host pool needs ~{pool_bytes / 2**30:.0f} GiB " + f"but only {mem.available / 2**30:.0f} GiB of RAM is available. " + "Lower the HiSparseConnector host_pool_gib or leave headroom " + "for co-tenants." + ) + + +def allocate_pinned_host_pool(size: int) -> tuple[torch.Tensor, torch.Tensor]: + """Allocate and deterministically register an exact-size host KV region.""" + page = 4096 + padded_size = round_up(size, page) + backing = torch.empty(padded_size + page, dtype=torch.int8, device="cpu") + aligned_offset = (-backing.data_ptr()) % page + registered = backing[aligned_offset : aligned_offset + padded_size] + pin_tensor(registered) + return registered[:size], registered + + +class HiSparseHostPool: + """Keep the single host backing and its exact CUDA registration range.""" + + def __init__(self) -> None: + self.backing: torch.Tensor | None = None + self.registered: torch.Tensor | None = None + + def allocate(self, size: int) -> torch.Tensor: + assert self.backing is None, "HiSparse host pool is already allocated." + check_hisparse_host_memory(size) + self.backing, self.registered = allocate_pinned_host_pool(size) + return self.backing + + +def release_pinned_state( + runtimes: list[HiSparseRuntime], pinned_host_pools: list[torch.Tensor] +) -> None: + """Synchronize and release registered host KV pools.""" + if pinned_host_pools: + try: + torch.accelerator.synchronize() + except RuntimeError as e: + logger.warning( + "HiSparse: CUDA context unusable at teardown (%s); leaving " + "%d host-pool tensors pinned for kernel exit reclaim.", + e, + len(pinned_host_pools), + ) + return + + cudart = torch.cuda.cudart() + release_start = time.perf_counter() + freed_bytes = 0 + while pinned_host_pools: + tensor = pinned_host_pools[-1] + err = cudart.cudaHostUnregister(tensor.data_ptr()) + if err.value != 0: + logger.warning( + "HiSparse: cudaHostUnregister failed (code=%d); leaving " + "%d host-pool tensors pinned for kernel exit reclaim.", + err.value, + len(pinned_host_pools), + ) + cudart.cudaGetLastError() + break + freed_bytes += tensor.nbytes + pinned_host_pools.pop() + if freed_bytes: + logger.info( + "HiSparse: unpinned %.1f GiB of host pool in %.1fs.", + freed_bytes / 2**30, + time.perf_counter() - release_start, + ) + + for runtime in runtimes: + del runtime._host_cache + del runtime.registered_host_pool + del runtime.hot_backing + + +def hisparse_prefill_staging_remap( + block_table: torch.Tensor, block_size: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Renumber blocks without a data-dependent CUDA output allocation.""" + block_ids = torch.cat( + ( + torch.zeros(1, dtype=torch.int32, device=block_table.device), + block_table.flatten().clamp(min=0).to(torch.int32), + ) + ) + sorted_ids, permutation = torch.sort(block_ids) + is_new = torch.cat( + ( + torch.ones(1, dtype=torch.bool, device=block_table.device), + sorted_ids[1:] != sorted_ids[:-1], + ) + ) + compact_sorted = torch.cumsum(is_new, dim=0, dtype=torch.int32) - 1 + inverse = torch.empty_like(compact_sorted) + inverse.scatter_(0, permutation, compact_sorted) + + unique_ids = torch.full_like(sorted_ids, -1) + unique_ids.scatter_(0, compact_sorted, sorted_ids) + valid = ( + torch.arange(sorted_ids.numel(), device=block_table.device) + <= compact_sorted[-1] + ) + unique_ids = torch.where(valid, unique_ids, -1) + new_bt = inverse[1:].view_as(block_table) + offsets = torch.arange(block_size, dtype=torch.int32, device=block_table.device) + row_ids = ( + torch.where( + valid[:, None], + unique_ids[:, None] * block_size + offsets, + -1, + ) + ).view(1, -1) + return new_bt, row_ids + + +@dataclass +class HiSparsePrefillStagingPlan: + block_table: torch.Tensor + row_ids: torch.Tensor + dst_rows: torch.Tensor + miss_mask: torch.Tensor + block_size: int + # Host rows with a GPU-resident copy (adopted shadow pages): the flat + # resident-cache row to read instead of DMAing from host, -1 for misses. + gpu_row_ids: torch.Tensor | None = None + gpu_source_key: tuple[int, int] | None = None + + def ensure_gpu_sources( + self, + resident_block_table: torch.Tensor, + resident_block_size: int, + ) -> None: + """Resolve which staged rows can be served from the resident cache. + + Computed once per plan (the resident block table is shared by every + layer in the group); non-null resident pages become miss_mask=0 rows + gathered device-to-device by ``gather_prefill_cache``. + """ + source_key = (resident_block_table.data_ptr(), resident_block_size) + if self.gpu_source_key == source_key: + return + block_size = self.block_size + if resident_block_size <= 0 or block_size % resident_block_size != 0: + return + device = self.row_ids.device + num_unique = self.row_ids.shape[1] // block_size + host_ids = self.row_ids[0].view(num_unique, block_size)[:, 0] // block_size + new_bt = self.block_table.to(torch.int64) + num_rows, num_cols = new_bt.shape + if num_rows == 0 or resident_block_table.shape[0] < num_rows: + return + # One representative (row, col) per unique host block: any request + # referencing the block holds an equivalent (refcounted) resident view. + flat_pos = torch.arange(num_rows * num_cols, device=device) + rep = torch.full( + (num_unique,), num_rows * num_cols, device=device, dtype=torch.int64 + ) + rep.scatter_reduce_( + 0, new_bt.reshape(-1), flat_pos, reduce="amin", include_self=True + ) + rep_row = (rep // num_cols).clamp(max=resident_block_table.shape[0] - 1) + rep_col = rep % num_cols + pages_per_block = block_size // resident_block_size + res_pos = rep_col[:, None] * pages_per_block + torch.arange( + pages_per_block, device=device + ) + res_cols = resident_block_table.shape[1] + res_blocks = torch.where( + res_pos < res_cols, + torch.gather( + resident_block_table.to(torch.int64)[rep_row], + 1, + res_pos.clamp(max=max(res_cols - 1, 0)), + ), + torch.zeros_like(res_pos), + ) + offsets = torch.arange(block_size, device=device) + per_off_block = res_blocks[:, offsets // resident_block_size] + gpu_rows = ( + per_off_block * resident_block_size + + (offsets % resident_block_size)[None, :] + ) + # Block id 0 is the null block in both id spaces. + valid_rows = self.row_ids.view(num_unique, block_size) >= 0 + hit = (per_off_block > 0) & (host_ids[:, None] > 0) & valid_rows + self.gpu_row_ids = torch.where(hit, gpu_rows, -1).reshape(1, -1).to(torch.int32) + self.miss_mask = ((self.gpu_row_ids < 0) & valid_rows.view(1, -1)).int() + self.gpu_source_key = source_key + + +def build_hisparse_prefill_staging_plan( + block_table: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + staging_block_capacity: int, +) -> HiSparsePrefillStagingPlan: + """Build an asynchronous layer-independent host-cache staging remap.""" + device = block_table.device + used = (seq_lens.to(torch.int64) + block_size - 1) // block_size + packed_offsets = torch.arange(staging_block_capacity, device=device) + row_ends = torch.cumsum(used, dim=0) + rows = torch.searchsorted(row_ends, packed_offsets, right=True) + valid = rows < block_table.shape[0] + rows = rows.clamp(max=block_table.shape[0] - 1) + row_starts = torch.cat((torch.zeros_like(row_ends[:1]), row_ends[:-1])) + columns = packed_offsets - row_starts[rows] + valid &= columns < block_table.shape[1] + flat_positions = rows * block_table.shape[1] + columns.clamp( + max=block_table.shape[1] - 1 + ) + referenced = torch.where(valid, block_table.flatten()[flat_positions], 0) + referenced_bt, row_ids = hisparse_prefill_staging_remap(referenced, block_size) + scratch_index = block_table.numel() + scatter_positions = torch.where(valid, flat_positions, scratch_index) + remapped = torch.where(valid, referenced_bt.flatten(), 0) + new_bt_storage = torch.zeros(scratch_index + 1, dtype=torch.int32, device=device) + new_bt_storage.scatter_(0, scatter_positions, remapped) + new_bt = new_bt_storage[:-1].view_as(block_table) + valid_rows = row_ids >= 0 + dst_rows = torch.arange(row_ids.shape[1], dtype=torch.int32, device=device).view( + 1, -1 + ) + return HiSparsePrefillStagingPlan( + block_table=new_bt, + row_ids=row_ids, + dst_rows=dst_rows, + miss_mask=valid_rows.to(torch.int32), + block_size=block_size, + ) + + +def _has_hisparse_ops() -> bool: + if not hasattr(torch.ops, "_C_cache_ops"): + return False + return ( + hasattr(torch.ops._C_cache_ops, "hisparse_resolve_residency") + and hasattr(torch.ops._C_cache_ops, "hisparse_invalidate_written_slots") + and hasattr(torch.ops._C_cache_ops, "hisparse_gather_plan") + and hasattr(torch.ops._C_cache_ops, "hisparse_gather_compact") + ) + + +class _SharedTopKState: + """CUDA-graph-safe residency result shared by index-sharing layers.""" + + def __init__(self, device: torch.device, max_rows: int, top_k: int) -> None: + self.device_topk_rows = torch.full( + (max_rows, top_k), -1, dtype=torch.int32, device=device + ) + self.physical_topk_indices = torch.empty_like(self.device_topk_rows) + self.swap_host_physical_rows = torch.empty( + (max_rows, top_k), dtype=torch.int32, device=device + ) + self.swap_device_physical_rows = torch.empty_like(self.swap_host_physical_rows) + self.swap_counts = torch.empty(max_rows, dtype=torch.int32, device=device) + self.valid_topk_counts = torch.empty(max_rows, dtype=torch.int32, device=device) + + +def _create_shared_topk_state( + device: torch.device, max_rows: int, top_k: int +) -> _SharedTopKState: + return _SharedTopKState(device, max_rows, top_k) + + +def _create_copy_stream(device: torch.device) -> torch.Stream: + return torch.Stream(device=device) + + +class HiSparseIndexGroup: + """GPU index-resolution state shared by one leader and its followers.""" + + leader: HiSparseRuntime + + def __init__( + self, + device: torch.device, + max_num_reqs: int, + region_stride: int, + max_swap_rows: int, + top_k: int, + copy_stream: torch.Stream | None = None, + logical_topk_ready: torch.Event | None = None, + ) -> None: + self.device_global_indices = torch.full( + (max_num_reqs, region_stride), + -1, + dtype=torch.int32, + device=device, + ) + self.lru_init = torch.arange(region_stride, dtype=torch.int16, device=device) + self.lru_slots = self.lru_init.repeat(max_num_reqs, 1).contiguous() + self.shared_topk = _create_shared_topk_state(device, max_swap_rows, top_k) + self.copy_stream = ( + copy_stream if copy_stream is not None else _create_copy_stream(device) + ) + self.logical_topk_ready = logical_topk_ready + self.followers: list[HiSparseRuntime] = [] + + +class HiSparseRuntime: + """Per-cache host/hot data plane and GPU replacement state.""" + + hot: PagedCacheView + hot_backing: torch.Tensor + hot_block_table: torch.Tensor + _host_cache: torch.Tensor + registered_host_pool: torch.Tensor + + def __init__( + self, + config: ResolvedHiSparseConfig, + max_num_reqs: int, + row_width: int, + kv_dtype: torch.dtype, + device: torch.device | str, + max_swap_rows: int | None = None, + index_group: HiSparseIndexGroup | None = None, + copy_stream: torch.Stream | None = None, + logical_topk_ready: torch.Event | None = None, + ) -> None: + if not _has_hisparse_ops(): + raise RuntimeError( + "HiSparse requires its compiled _C_cache_ops CUDA kernels " + "(host-resident decode has no Python fallback). Rebuild vLLM " + "from source so " + "csrc/libtorch_stable/hisparse_kernels.cu is included." + ) + self.max_num_reqs = max_num_reqs + self.row_width = row_width + self.kv_dtype = kv_dtype + self.device = torch.device(device) + # Logical slots per request. Physical rows come from its ephemeral HMA + # block table and need not be contiguous in the shared slab. + self.region_stride = config.device_buffer_size + + row_bytes = row_width * kv_dtype.itemsize + if row_bytes % 16 != 0: + raise ValueError( + f"HiSparse requires 16-byte aligned KV rows, got {row_bytes}B." + ) + + self._layer_ready_event: torch.Event | None = None + self._swap_staged = False + self._swap_step = 0 + self.is_group_leader = index_group is None + if index_group is None: + index_group = HiSparseIndexGroup( + self.device, + max_num_reqs, + self.region_stride, + max_swap_rows or max_num_reqs, + config.top_k, + copy_stream, + logical_topk_ready, + ) + index_group.leader = self + else: + index_group.followers.append(self) + self.index_group = index_group + + self.eager_host_mirror = config.eager_host_mirror + self.resident_source_index = -1 + self.request_state_indices: torch.Tensor | None = None + + @property + def host_cache(self) -> torch.Tensor: + return self._host_cache + + def bind_hot_cache( + self, + raw_tensor: torch.Tensor, + *, + byte_offset: int, + block_stride: int, + num_blocks: int, + block_size: int, + block_table: torch.Tensor, + ) -> None: + """Bind this runtime's strided view into the shared GPU HMA slab.""" + self.hot = PagedCacheView.bind( + raw_tensor, + dtype=self.kv_dtype, + row_width=self.row_width, + byte_offset=byte_offset, + block_stride=block_stride, + num_blocks=num_blocks, + block_size=block_size, + ) + self.hot_backing = raw_tensor + self.hot_block_table = block_table + + def bind_source_cache( + self, + kv_cache: torch.Tensor, + *, + registered_host_pool: torch.Tensor | None = None, + ) -> None: + if kv_cache.dtype != self.kv_dtype or kv_cache.shape[-1] != self.row_width: + raise ValueError( + "HiSparse runtime bound to a KV cache with mismatched " + f"layout: expected ({self.row_width}, {self.kv_dtype}), got " + f"({kv_cache.shape[-1]}, {kv_cache.dtype})." + ) + # Host-resident pool: the cache itself is the only full-size store; + # every allocated slot is written before the indexer can select it. + if kv_cache.device.type != "cpu": + raise ValueError( + "HiSparse requires a host-resident KV pool; got a KV cache " + f"on {kv_cache.device}." + ) + # The pool is pinned via cudaHostRegister (exact-size, deterministic + # unpin at shutdown); torch's is_pinned() only recognizes its own + # caching-host-allocator memory, so also accept ranges the model + # allocator explicitly registered. + if not (kv_cache.is_pinned() or registered_host_pool is not None): + raise ValueError("HiSparse host-resident KV pool must be pinned memory.") + + self._host_cache = kv_cache.view(-1, kv_cache.shape[-1]) + self.registered_host_pool = ( + registered_host_pool if registered_host_pool is not None else kv_cache + ) + + def gather_prefill_cache( + self, + kv_cache: torch.Tensor, + plan: HiSparsePrefillStagingPlan, + resident_cache: torch.Tensor | None = None, + ) -> torch.Tensor: + """Gather this runtime's host cache using a shared staging plan.""" + if kv_cache.shape[1] != plan.block_size: + raise ValueError( + f"HiSparse staging block size {plan.block_size} does not match " + f"the KV cache block size {kv_cache.shape[1]}." + ) + row_width = kv_cache.shape[-1] + + staged = torch.empty( + ( + plan.row_ids.shape[1] // plan.block_size, + plan.block_size, + row_width, + ), + dtype=kv_cache.dtype, + device=plan.block_table.device, + ) + if plan.gpu_row_ids is not None and resident_cache is not None: + gpu_rows = plan.gpu_row_ids[0].to(torch.long) + src = gpu_rows.clamp_min(0) + resident_block_size = resident_cache.shape[1] + rows = resident_cache[src // resident_block_size, src % resident_block_size] + if rows.dtype != staged.dtype: + rows = rows.contiguous().view(staged.dtype) + staged.view(-1, row_width).copy_(rows) + torch.ops._C_cache_ops.hisparse_gather_plan( + kv_cache.view(-1, row_width), + staged, + plan.row_ids, + plan.dst_rows, + plan.miss_mask, + None, + None, + 0, + ) + return staged + + def reset_hot_state(self) -> None: + """Drop all hot-buffer bookkeeping (hits become misses).""" + group = self.index_group + compute_stream = current_stream() + compute_stream.wait_stream(group.copy_stream) + group.device_global_indices.fill_(-1) + group.lru_slots.copy_(group.lru_init.expand_as(group.lru_slots)) + group.copy_stream.wait_stream(compute_stream) + + def invalidate_sorted_slots( + self, + sorted_slots: torch.Tensor, + state_indices: torch.Tensor, + ) -> None: + """Drop hot copies using shared, preprocessed invalidation inputs.""" + device_global_indices = self.index_group.device_global_indices + active_indices = device_global_indices.index_select(0, state_indices) + positions = torch.searchsorted(sorted_slots, active_indices) + positions.clamp_(max=sorted_slots.numel() - 1) + active_indices.masked_fill_( + sorted_slots[positions] == active_indices, + -1, + ) + device_global_indices.index_copy_(0, state_indices, active_indices) + + def invalidate_written_slots( + self, + written_slots: torch.Tensor, + req_id_per_token: torch.Tensor, + ) -> None: + num_tokens = min(written_slots.numel(), req_id_per_token.numel()) + if num_tokens == 0: + return + assert self.request_state_indices is not None + torch.ops._C_cache_ops.hisparse_invalidate_written_slots( + self.index_group.device_global_indices, + self.request_state_indices, + req_id_per_token[:num_tokens], + written_slots[:num_tokens], + ) + + def begin_forward(self) -> None: + self._swap_step = 0 + + def _step_rows(self, num_tokens: int) -> slice: + start = self._swap_step * num_tokens + stop = start + num_tokens + if stop > self.index_group.shared_topk.physical_topk_indices.shape[0]: + raise ValueError( + "HiSparse swap rows exceed the configured speculative decode " + f"capacity: stop={stop}, capacity=" + f"{self.index_group.shared_topk.physical_topk_indices.shape[0]}." + ) + self._swap_step += 1 + return slice(start, stop) + + def _resolve_residency( + self, + *, + resident: HiSparseCacheHandle | None = None, + req_id_per_token: torch.Tensor, + block_table: torch.Tensor, + topk_indices: torch.Tensor, + block_size: int, + return_valid_counts: bool = False, + shared_rows: slice, + attention_indices_out: torch.Tensor | None = None, + valid_counts_out: torch.Tensor | None = None, + ) -> None: + """Resolve logical top-k positions and compact rows requiring swaps.""" + num_tokens = topk_indices.shape[0] + hot = self.hot + host_cache = self.host_cache + assert hot.block_size == block_size + request_state_indices = self.request_state_indices + assert request_state_indices is not None + group = self.index_group + shared = group.shared_topk + + device_topk_rows = shared.device_topk_rows[shared_rows] + valid_topk_counts = None + if return_valid_counts: + valid_topk_counts = ( + shared.valid_topk_counts[shared_rows] + if valid_counts_out is None + else valid_counts_out + ) + physical_topk_indices = ( + shared.physical_topk_indices[shared_rows] + if attention_indices_out is None + else attention_indices_out + ) + swap_host_physical_rows = shared.swap_host_physical_rows[shared_rows] + swap_device_physical_rows = shared.swap_device_physical_rows[shared_rows] + swap_counts = shared.swap_counts[shared_rows] + + # Padded rows are skipped by the kernel (request_state_indices) and must + # come out as -1 so the attention kernel masks them. + torch.ops._C_cache_ops.hisparse_resolve_residency( + host_cache, + hot.cache, + self.hot_block_table, + topk_indices, + device_topk_rows, + group.device_global_indices, + group.lru_slots, + request_state_indices, + self.region_stride, + None, + physical_topk_indices, + hot.attention_block_stride, + req_id_per_token[:num_tokens].contiguous(), + block_table, + block_size, + None, + valid_topk_counts, + swap_host_physical_rows, + swap_device_physical_rows, + swap_counts, + resident.block_table if resident is not None else None, + resident.view.block_size + if resident is not None and resident.view is not None + else 0, + 0, + ) + + def _swap_rows(self, shared_rows: slice) -> None: + hot = self.hot + shared = self.index_group.shared_topk + torch.ops._C_cache_ops.hisparse_gather_compact( + self.host_cache, + hot.cache, + shared.swap_host_physical_rows[shared_rows], + shared.swap_device_physical_rows[shared_rows], + shared.swap_counts[shared_rows], + ) + + def _resolve_and_stage_group( + self, + *, + resident: HiSparseCacheHandle, + req_id_per_token: torch.Tensor, + block_table: torch.Tensor, + logical_topk_indices: torch.Tensor, + block_size: int, + return_valid_counts: bool, + shared_rows: slice, + attention_indices_out: torch.Tensor | None, + valid_counts_out: torch.Tensor | None, + ) -> None: + group = self.index_group + compute_stream = current_stream() + if group.logical_topk_ready is not None: + group.copy_stream.wait_event(group.logical_topk_ready) + else: + group.copy_stream.wait_stream(compute_stream) + with group.copy_stream: + self._resolve_residency( + resident=resident, + req_id_per_token=req_id_per_token, + block_table=block_table, + topk_indices=logical_topk_indices, + block_size=block_size, + return_valid_counts=return_valid_counts, + shared_rows=shared_rows, + attention_indices_out=attention_indices_out, + valid_counts_out=valid_counts_out, + ) + runtimes = [self] + if resident.decode_batch: + runtimes.extend(group.followers) + for runtime in runtimes: + runtime._swap_rows(shared_rows) + if runtime._layer_ready_event is None: + runtime._layer_ready_event = torch.Event() + runtime._layer_ready_event.record(group.copy_stream) + runtime._swap_staged = True + + def _stage_this_layer(self, shared_rows: slice) -> None: + group = self.index_group + compute_stream = current_stream() + group.copy_stream.wait_stream(compute_stream) + with group.copy_stream: + self._swap_rows(shared_rows) + if self._layer_ready_event is None: + self._layer_ready_event = torch.Event() + self._layer_ready_event.record(group.copy_stream) + self._swap_staged = True + + def wait_for_staged_swap(self) -> None: + if not self._swap_staged or self._layer_ready_event is None: + raise RuntimeError("HiSparse layer swap was not prefetched.") + current_stream().wait_event(self._layer_ready_event) + self._swap_staged = False + + def swap_in( + self, + *, + resident: HiSparseCacheHandle, + req_id_per_token: torch.Tensor, + block_table: torch.Tensor, + logical_topk_indices: torch.Tensor, + block_size: int, + return_valid_counts: bool = False, + attention_indices_out: torch.Tensor | None = None, + valid_counts_out: torch.Tensor | None = None, + ) -> HiSparseTopKResult: + """Resolve once per group and ensure this layer's rows are available.""" + num_tokens = logical_topk_indices.shape[0] + shared_rows = self._step_rows(num_tokens) + group = self.index_group + hot = self.hot + assert hot.block_size == block_size + assert hot.attention_block_stride == group.leader.hot.attention_block_stride + + if self.is_group_leader: + self._resolve_and_stage_group( + resident=resident, + req_id_per_token=req_id_per_token, + block_table=block_table, + logical_topk_indices=logical_topk_indices, + block_size=block_size, + return_valid_counts=return_valid_counts, + shared_rows=shared_rows, + attention_indices_out=attention_indices_out, + valid_counts_out=valid_counts_out, + ) + + if not self._swap_staged: + self._stage_this_layer(shared_rows) + self.wait_for_staged_swap() + + physical_topk_indices = group.shared_topk.physical_topk_indices[shared_rows] + if return_valid_counts: + return ( + physical_topk_indices, + group.shared_topk.valid_topk_counts[shared_rows], + ) + return physical_topk_indices + + +class HiSparseCacheHandle: + """Attention-facing handle for resident KV and sparse offload state.""" + + def __init__(self, runtime: HiSparseRuntime) -> None: + self.view: PagedCacheView | None = None + self.block_table: torch.Tensor | None = None + self.source_block_table: torch.Tensor | None = None + self.slot_mapping: torch.Tensor | None = None + self.runtime = runtime + self.dummy_batch = False + self.decode_batch = False + self.all_context_pages_resident = True + self.num_actual_tokens = 0 + self.num_decode_tokens = 0 + self.req_id_per_token: torch.Tensor | None = None + self.host_mirror_required = False + self.mirror_from_resident = False + self.mirror_slot_mapping: torch.Tensor | None = None + self.mirror_staging_cache: torch.Tensor | None = None + self.mirror_staging_slots: torch.Tensor | None = None + self.submit_layer_mirror: Callable[[], None] | None = None + self.index_group_caches: list[HiSparseCacheHandle] = [self] + + def prepare_group_for_batch(self, attn_metadata: Any | None) -> None: + assert self.runtime.is_group_leader + for cache in self.index_group_caches: + cache._prepare_for_batch(attn_metadata) + + def _prepare_for_batch(self, attn_metadata: Any | None) -> None: + self.dummy_batch = attn_metadata is None + self.runtime.begin_forward() + self.num_actual_tokens = ( + attn_metadata.num_actual_tokens if attn_metadata is not None else 0 + ) + self.num_decode_tokens = ( + attn_metadata.num_decode_tokens if attn_metadata is not None else 0 + ) + self.req_id_per_token = ( + attn_metadata.req_id_per_token if attn_metadata is not None else None + ) + self.decode_batch = is_hisparse_decode_batch(attn_metadata) + self.host_mirror_required = attn_metadata is not None and ( + not self.decode_batch or self.runtime.eager_host_mirror + ) + + def write_target( + self, num_input_rows: int, num_slot_rows: int + ) -> tuple[torch.Tensor, torch.Tensor, int]: + assert self.view is not None and self.slot_mapping is not None + num_rows = min(num_input_rows, num_slot_rows, self.num_actual_tokens) + if self.decode_batch: + num_rows = min(num_rows, self.runtime.max_num_reqs) + return self.view.cache, self.slot_mapping[:num_rows], num_rows + + def mirror_write_target( + self, num_rows: int + ) -> tuple[torch.Tensor, torch.Tensor] | None: + if ( + not self.host_mirror_required + or self.decode_batch + or self.mirror_from_resident + ): + return None + cache = self.mirror_staging_cache + slots = self.mirror_staging_slots + if cache is None or slots is None: + raise RuntimeError("HiSparse prefill mirror staging is not bound.") + if num_rows > slots.numel(): + raise RuntimeError( + "HiSparse prefill mirror exceeds staging capacity: " + f"{num_rows} > {slots.numel()}." + ) + return cache, slots[:num_rows] + + def finish_kv_update(self) -> None: + if self.dummy_batch: + return + if ( + self.submit_layer_mirror is not None + and not self.decode_batch + and get_forward_context().cudagraph_runtime_mode != CUDAGraphMode.FULL + ): + self.submit_layer_mirror() + + def bind_cache( + self, + raw_tensor: torch.Tensor, + *, + byte_offset: int, + block_stride: int, + num_blocks: int, + block_size: int, + block_table: torch.Tensor, + slot_mapping: torch.Tensor, + ) -> None: + self.view = PagedCacheView.bind( + raw_tensor, + dtype=self.runtime.kv_dtype, + row_width=self.runtime.row_width, + byte_offset=byte_offset, + block_stride=block_stride, + num_blocks=num_blocks, + block_size=block_size, + ) + self.block_table = block_table + self.slot_mapping = slot_mapping + + def swap_in( + self, + req_id_per_token: torch.Tensor, + block_table: torch.Tensor, + logical_topk_indices: torch.Tensor, + *, + block_size: int, + return_valid_counts: bool = False, + attention_indices_out: torch.Tensor | None = None, + valid_counts_out: torch.Tensor | None = None, + ) -> HiSparseTopKResult: + return self.runtime.swap_in( + resident=self, + req_id_per_token=req_id_per_token[: logical_topk_indices.shape[0]], + block_table=block_table, + logical_topk_indices=logical_topk_indices, + block_size=block_size, + return_valid_counts=return_valid_counts, + attention_indices_out=attention_indices_out, + valid_counts_out=valid_counts_out, + ) + + +def initialize_hisparse_runtime_buffers( + cache_handles: list[HiSparseCacheHandle], + *, + max_num_reqs: int, + max_num_batched_tokens: int, +) -> None: + """Allocate shared request state and per-layer prefill staging buffers.""" + assert cache_handles + resident = cache_handles[0].view + assert resident is not None + device = resident.cache.device + request_state_indices = torch.full( + (max_num_reqs,), -1, dtype=torch.int32, device=device + ) + staging_blocks = ( + max_num_batched_tokens + resident.block_size - 1 + ) // resident.block_size + mirror_staging_caches = torch.empty( + ( + len(cache_handles), + staging_blocks, + resident.block_size, + resident.cache.shape[-1], + ), + dtype=resident.cache.dtype, + device=device, + ) + mirror_staging_slots = torch.arange( + max_num_batched_tokens, dtype=torch.int64, device=device + ) + for layer_index, cache_handle in enumerate(cache_handles): + cache_handle.runtime.request_state_indices = request_state_indices + cache_handle.mirror_staging_cache = mirror_staging_caches[layer_index] + cache_handle.mirror_staging_slots = mirror_staging_slots + + +def create_hisparse_cache_handle( + vllm_config: VllmConfig, + model_top_k: int, + *, + is_index_group_leader: bool, + row_width: int, + kv_dtype: torch.dtype, + index_group: HiSparseMLAIndexGroup | None = None, + device: torch.device | str | None = None, +) -> HiSparseCacheHandle | None: + config = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k) + if config is None: + return None + + max_num_reqs = vllm_config.scheduler_config.max_num_seqs + max_swap_rows = _get_max_swap_rows(vllm_config) + if device is None: + device = torch.device( + current_platform.device_type, torch.accelerator.current_device_index() + ) + + hisparse_group = ( + None + if is_index_group_leader or index_group is None + else index_group.hisparse_group + ) + runtime = HiSparseRuntime( + config=config, + max_num_reqs=max_num_reqs, + max_swap_rows=max_swap_rows, + row_width=row_width, + kv_dtype=kv_dtype, + device=device, + index_group=hisparse_group, + copy_stream=index_group.side_stream if index_group is not None else None, + logical_topk_ready=( + index_group.logical_topk_ready + if index_group is not None and index_group.has_indexer + else None + ), + ) + if is_index_group_leader and index_group is not None: + index_group.hisparse_group = runtime.index_group + logger.info_once( + "Enabled experimental HiSparse HMA hot cache: top_k=%d, " + "device_buffer_size=%d (%d LRU rows), max_num_seqs=%d.", + config.top_k, + config.device_buffer_size, + config.device_buffer_size, + max_num_reqs, + ) + handle = HiSparseCacheHandle(runtime) + speculative_config = vllm_config.speculative_config + handle.mirror_from_resident = bool( + vllm_config.scheduler_config.async_scheduling + and speculative_config is not None + and speculative_config.uses_draft_kv_cache() + ) + return handle diff --git a/vllm/v1/hisparse/types.py b/vllm/v1/hisparse/types.py new file mode 100644 index 0000000000..8618f3cb20 --- /dev/null +++ b/vllm/v1/hisparse/types.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class SparseKVPageTransfer: + """Copy one logical KV page between cache-manager and worker-owned tiers.""" + + transfer_id: int + destination_block_id: int + source_block_ids: tuple[int, ...] + after_forward: bool + + +@dataclass(frozen=True) +class SparseKVRowMirror: + """Mirror one contiguous resident-row span into the host cache.""" + + source_starts: tuple[int, ...] + destination_start: int + num_rows: int + + +@dataclass +class SparseKVOffloadCommand: + """Opaque scheduler-to-worker command for one model step.""" + + page_transfers: list[SparseKVPageTransfer] diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 1888404a57..fb839bcd68 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -147,6 +147,11 @@ class KVCacheSpecKind(str, Enum): UNKNOWN = "unknown" +class SparseCacheRole(str, Enum): + SPARSE = "sparse" + INDEXER = "indexer" + + @dataclass(frozen=True) class KVCacheSpec: """ @@ -192,6 +197,10 @@ class KVCacheSpec: return kernel_block_size // self.tokens_per_state return 1 + @property + def block_table_token_alignment(self) -> int | None: + return 128 + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: """ The maximum possible memory usage of this KV cache in bytes. @@ -245,6 +254,25 @@ class KVCacheSpec: isinstance(spec, uniform_type_base_spec) for spec in kv_cache_specs.values() ) + @property + def has_layer_views(self) -> bool: + """Whether generic allocation creates per-layer cache views. + + Specs without a per-layer shape keep the raw backing tensor and lay it + out themselves. They also have no attention module of their own, so + they take no part in attention-backend or kernel-block selection. + """ + return True + + @property + def uses_slot_mapping(self) -> bool: + """Whether the worker computes a per-token slot mapping for this spec. + + Specs that address their pages themselves (raw storage, ring buffers) + take no slot mapping row. + """ + return True + def group_kernel_blocks(cache: torch.Tensor, num_blocks: int) -> torch.Tensor: """View a kernel-block-granular layer cache with manager blocks as dim 0. @@ -383,6 +411,70 @@ def create_kv_cache_views( return views +@dataclass(frozen=True, kw_only=True) +class HiSparseHotSpec(KVCacheSpec): + """Ephemeral per-request HiSparse hot-cache allocation.""" + + page_size: int + blocks_per_request: int + + @property + def prefix_cacheable(self) -> bool: + return False + + @property + def page_size_bytes(self) -> int: + return self.page_size + + @property + def block_table_token_alignment(self) -> int | None: + return None + + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: + return self.blocks_per_request * self.page_size + + def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: + return self.blocks_per_request + + @property + def has_layer_views(self) -> bool: + return False + + @property + def uses_slot_mapping(self) -> bool: + return False + + +@dataclass(frozen=True, kw_only=True) +class HiSparseResidentSpec(KVCacheSpec): + """Reclaimable GPU-resident pages for host-backed HiSparse KV.""" + + page_size: int + + @property + def prefix_cacheable(self) -> bool: + return False + + @property + def page_size_bytes(self) -> int: + return self.page_size + + def max_admission_blocks_per_request( + self, max_in_flight_tokens: int, max_model_len: int + ) -> int: + num_tokens = min(max_in_flight_tokens, max_model_len) + return cdiv(num_tokens, self.block_size) + + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: + return cdiv(vllm_config.model_config.max_model_len, self.block_size) * ( + self.page_size + ) + + @property + def has_layer_views(self) -> bool: + return False + + @dataclass(frozen=True, kw_only=True) class AttentionSpec(KVCacheSpec): num_kv_heads: int @@ -557,6 +649,8 @@ class MLAAttentionSpec(FullAttentionSpec): # DeepseekV4 only fields. Non-DeepseekV4 MLA models leave these at defaults. alignment: int | None = None # Default to None for no padding. model_version: str | None = None + cache_role: SparseCacheRole = SparseCacheRole.SPARSE + is_index_group_leader: bool = False storage_block_size: int | None = None """Token width used to view storage when it differs from the kernel block.""" # Group capability enabled when any member flattens a non-causal query block @@ -577,16 +671,20 @@ class MLAAttentionSpec(FullAttentionSpec): cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs) tokens_per_state_set = set(spec.tokens_per_state for spec in specs) model_version_set = set(spec.model_version for spec in specs) + cache_role_set = {spec.cache_role for spec in specs} + index_group_leader_set = {spec.is_index_group_leader for spec in specs} storage_block_size_set = set(spec.storage_block_size for spec in specs) assert ( len(cache_dtype_str_set) == 1 and len(tokens_per_state_set) == 1 and len(model_version_set) == 1 + and len(cache_role_set) == 1 + and len(index_group_leader_set) == 1 and len(storage_block_size_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " - "quantization method, tokens per state, model version, and storage " - "block size." + "quantization method, tokens per state, model version, cache role, " + "index-sharing role, and storage block size." ) merged_spec = cls( block_size=specs[0].block_size, @@ -600,6 +698,8 @@ class MLAAttentionSpec(FullAttentionSpec): cache_dtype_str=cache_dtype_str_set.pop(), tokens_per_state=tokens_per_state_set.pop(), model_version=model_version_set.pop(), + cache_role=cache_role_set.pop(), + is_index_group_leader=index_group_leader_set.pop(), storage_block_size=storage_block_size_set.pop(), non_causal_multi_token_decode=any( spec.non_causal_multi_token_decode for spec in specs @@ -772,6 +872,10 @@ class CircularBufferSpec(AttentionSpec): reads the open group's committed keys from the ring. """ + @property + def block_table_token_alignment(self) -> int | None: + return None + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: # The ring occupies one block per request for its whole lifetime. del vllm_config @@ -792,6 +896,10 @@ class CircularBufferSpec(AttentionSpec): def prefix_cacheable(self) -> bool: return False + @property + def uses_slot_mapping(self) -> bool: + return False + @dataclass(frozen=True, kw_only=True) class SlidingWindowMLASpec(SlidingWindowSpec): @@ -919,6 +1027,10 @@ class MambaSpec(KVCacheSpec): return self.page_size_padded return page_size + @property + def block_table_token_alignment(self) -> int | None: + return None + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: if vllm_config.cache_config.mamba_cache_mode == "all": max_model_len = vllm_config.model_config.max_model_len @@ -1093,6 +1205,10 @@ class UniformTypeKVCacheSpecs(KVCacheSpec): def page_size_bytes(self) -> int: return sum(spec.page_size_bytes for spec in self.kv_cache_specs.values()) + @property + def block_table_token_alignment(self) -> int | None: + return self.first_spec.block_table_token_alignment + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: max_num_pages = max( cdiv(spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes) @@ -1154,6 +1270,10 @@ class UniformTypeKVCacheSpecs(KVCacheSpec): for spec in self.kv_cache_specs.values() ) + @property + def uses_slot_mapping(self) -> bool: + return self.first_spec.uses_slot_mapping + def iter_layer_specs(kv_cache_spec: KVCacheSpec) -> Collection[KVCacheSpec]: """The per-layer specs a KV cache group spec covers. @@ -1259,6 +1379,13 @@ class KVCacheTensor: layer_stride: int block_stride: int offset: int = 0 # byte offset of layers[0]'s block 0 + host_resident: bool = False + + +class KVCacheGroupRole(str, Enum): + DEFAULT = "default" + HISPARSE_SOURCE = "hisparse_source" + HISPARSE_INDEXER = "hisparse_indexer" @dataclass @@ -1274,8 +1401,13 @@ class KVCacheGroupSpec: kv_cache_spec: KVCacheSpec # Whether this group contains EAGLE/MTP draft attention layers. is_eagle_group: bool = False + # Host groups use the dedicated HiSparse pool; others share the device pool. + host_resident: bool = False # Whether this group is part of the externally transferable KV state. + # Ephemeral accelerator-side replicas are rebuilt from their transferable + # source group after a connector load. enable_kv_transfer: bool = True + role: KVCacheGroupRole = KVCacheGroupRole.DEFAULT @dataclass @@ -1300,6 +1432,8 @@ class KVCacheConfig: """Resolved retention policy for local prefix-cache checkpoints.""" kv_cache_layout: str | None = None """The KV cache layout resolved by the engine core, adopted by all workers.""" + hisparse_host_num_blocks: int | None = None + """Capacity of the dedicated HiSparse host-block manager, when enabled.""" @cached_property def transfer_group_ids(self) -> tuple[int, ...]: @@ -1337,6 +1471,21 @@ class KVCacheConfig: ) return tuple(block_ids[group_id] for group_id in self.transfer_group_ids) + @property + def host_group_ids(self) -> tuple[int, ...]: + return tuple( + group_id + for group_id, group in enumerate(self.kv_cache_groups) + if group.host_resident + ) + + def num_blocks_of(self, tensor: KVCacheTensor) -> int: + """Number of blocks addressable by the pool backing ``tensor``.""" + if not tensor.host_resident: + return self.num_blocks + assert self.hisparse_host_num_blocks is not None + return self.hisparse_host_num_blocks + @property def has_mamba_layers(self) -> bool: return any( @@ -1347,9 +1496,11 @@ class KVCacheConfig: @property def has_mixed_precision_kv_cache(self) -> bool: - """Whether attention groups store their KV cache at more than one precision.""" + """Whether device attention caches use more than one precision.""" kv_cache_precisions: set[tuple[torch.dtype, KVQuantMode]] = set() for group in self.kv_cache_groups: + if group.host_resident: + continue kv_cache_precisions.update( (spec.dtype, spec.kv_quant_mode) for spec in iter_layer_specs(group.kv_cache_spec) diff --git a/vllm/v1/kv_cache_spec_registry.py b/vllm/v1/kv_cache_spec_registry.py index 816a6862da..62cb291892 100644 --- a/vllm/v1/kv_cache_spec_registry.py +++ b/vllm/v1/kv_cache_spec_registry.py @@ -34,6 +34,7 @@ class KVCacheSpecMetadata: _REGISTRY_KVCACHESPEC_LIST: dict[type["KVCacheSpec"], KVCacheSpecMetadata] = {} +_REGISTRY_ROLE_MANAGERS: dict[str, type["SingleTypeKVCacheManager"]] = {} class KVCacheSpecRegistry: @@ -102,20 +103,34 @@ class KVCacheSpecRegistry: uniform_type_base_spec=uniform_type_base_spec, ) + @classmethod + def register_role_manager( + cls, role: str, manager_class: type["SingleTypeKVCacheManager"] + ) -> None: + """Select ``manager_class`` for every group with the given role.""" + registered = _REGISTRY_ROLE_MANAGERS.get(role) + assert registered is None or registered is manager_class, ( + f"Conflicting manager registration for group role {role}" + ) + _REGISTRY_ROLE_MANAGERS[role] = manager_class + @classmethod def get_manager_class( - cls, kvcache_spec: "KVCacheSpec" + cls, kvcache_spec: "KVCacheSpec", role: str | None = None ) -> type["SingleTypeKVCacheManager"] | None: """ Get the single type kvcache manager class for a given kvcache spec instance. Args: kvcache_spec: A KVCacheSpec instance + role: The group's role; a role-registered manager takes precedence. Returns: The SingleTypeKVCacheManager class to use for this kvcache_spec """ cls._ensure_registered() + if role is not None and (manager_class := _REGISTRY_ROLE_MANAGERS.get(role)): + return manager_class kvcache_spec_cls = type(kvcache_spec) # Walk up the MRO to find a registered base class diff --git a/vllm/v1/kv_offload/config.py b/vllm/v1/kv_offload/config.py index 637d3bb664..74b619dc0f 100644 --- a/vllm/v1/kv_offload/config.py +++ b/vllm/v1/kv_offload/config.py @@ -14,6 +14,8 @@ class OffloadingGroupConfig: tokens_per_block: int # Layer names belonging to this group. layer_names: tuple[str, ...] + # Original KVCacheConfig group index. + group_id: int @dataclass(frozen=True) diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index e94380825f..d8924b2738 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -281,6 +281,9 @@ class KVConnectorOutput: # IDs of externally computed KV blocks that failed to load. # Requests referencing these blocks should be rescheduled to recompute them invalid_block_ids: set[int] = field(default_factory=set) + # Receive failures keyed by request identity. This remains unambiguous for + # hybrid/multi-pool cache layouts where numeric block IDs overlap. + failed_recving: set[str] = field(default_factory=set) # Configuration describing how many finished sending/receiving # notifications should be expected for each request. This allows # handshake-based connectors like Nixl to update the KVOutputAggregator. @@ -295,6 +298,7 @@ class KVConnectorOutput: and not self.kv_connector_stats and not self.kv_cache_events and not self.invalid_block_ids + and not self.failed_recving and not self.kv_connector_worker_meta ) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 4a85992425..7e7a212b5f 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -18,6 +18,10 @@ from vllm.v1.attention.backend import ( CommonAttentionMetadata, ) from vllm.v1.attention.backends.utils import create_fast_prefill_custom_backend +from vllm.v1.hisparse.binding import ( + init_hisparse_kv_cache, + resolve_hisparse_block_size, +) from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheConfig, @@ -30,11 +34,12 @@ from vllm.v1.worker.utils import ( AttentionGroup, add_kv_sharing_layers_to_kv_cache_groups, allocate_kv_cache, - bind_kv_cache, + bind_kv_cache_to_layers, prepare_kernel_block_sizes, ) if TYPE_CHECKING: + from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, CudaGraphManager, @@ -132,6 +137,7 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: if isinstance(spec, AttentionSpec): spec = attn_module.get_attn_backend().customize_spec(spec) kv_cache_spec[layer_name] = spec + resolve_hisparse_block_size(vllm_config, kv_cache_spec, attn_layers) return kv_cache_spec @@ -198,6 +204,9 @@ def init_attn_backend( kv_cache_config.kv_cache_groups ): layer_names = kv_cache_group_spec.layer_names + if not kv_cache_group_spec.kv_cache_spec.has_layer_views: + attn_groups.append([]) + continue if active_layer_names is not None: layer_names = list(active_layer_names.intersection(layer_names)) @@ -316,22 +325,34 @@ def get_query_lens_mismatch_unsupported_backend( def init_kv_cache( - runner_kv_caches: list[torch.Tensor | list[torch.Tensor]], forward_context: dict[str, Any], kv_cache_config: KVCacheConfig, device: torch.device, kernel_block_sizes: list[int], vllm_config: VllmConfig, kv_cache_allocation_context: AbstractContextManager | None = None, + *, + block_tables: "BlockTables | None" = None, ) -> dict[str, Any]: allocation_context = kv_cache_allocation_context or nullcontext() with allocation_context: - kv_caches = allocate_kv_cache( - kv_cache_config, - device, - vllm_config.cache_config.get_resolved_kv_cache_layout(), - kernel_block_sizes, - ) + if vllm_config.attention_config.hisparse_config is not None: + assert block_tables is not None + kv_caches = init_hisparse_kv_cache( + kv_cache_config, + device, + kernel_block_sizes, + vllm_config, + forward_context, + block_tables, + ) + else: + kv_caches = allocate_kv_cache( + kv_cache_config, + device, + vllm_config.cache_config.get_resolved_kv_cache_layout(), + kernel_block_sizes, + ) for layer_name, target in get_shared_kv_cache_layers(vllm_config).items(): kv_caches[layer_name] = kv_caches[target] # Dual-attention models (e.g. LongCat-Flash) put two Attention modules per @@ -342,10 +363,12 @@ def init_kv_cache( in ("longcat_flash", "longcat_flash_ngram") else 1 ) - bind_kv_cache( - kv_caches, + bindable_caches = { + name: cache for name, cache in kv_caches.items() if name in forward_context + } + bind_kv_cache_to_layers( + bindable_caches, forward_context, - runner_kv_caches, num_attn_module, kv_cache_groups=kv_cache_config.kv_cache_groups, ) @@ -396,6 +419,8 @@ def build_attn_metadata( attn_metadata: dict[str, Any] = {} num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) for i in range(num_kv_cache_groups): + if not attn_groups[i]: + continue block_table = block_tables[i] slot_mapping = slot_mappings[i] # Per-group causal for hybrid drafters (mixed SWA/full attention). diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 287e64ce30..12354f3c22 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -34,6 +34,7 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import round_up from vllm.utils.torch_utils import current_stream +from vllm.v1.hisparse.binding import release_hisparse_profiling_cache from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer @@ -931,6 +932,7 @@ def _teardown_profiling_state(runner: "GPUModelRunner") -> None: if (model := getattr(runner, "model", None)) is not None: layers = itertools.chain(layers, model.modules()) clear_layer_kv_caches(layers) + release_hisparse_profiling_cache(runner.compilation_config.static_forward_context) runner.cache_config.num_gpu_blocks = None runner.maybe_remove_all_loras(runner.lora_config) gc.collect() diff --git a/vllm/v1/worker/gpu/kv_connector.py b/vllm/v1/worker/gpu/kv_connector.py index eaef0c2916..88b7b0cf9d 100644 --- a/vllm/v1/worker/gpu/kv_connector.py +++ b/vllm/v1/worker/gpu/kv_connector.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch @@ -29,7 +29,10 @@ if TYPE_CHECKING: class KVConnector: """KVConnector interface used by GPUModelRunner.""" - def pre_forward(self, scheduler_output: "SchedulerOutput") -> None: + def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> None: + pass + + def finish_forward(self) -> None: pass def post_forward( @@ -43,6 +46,9 @@ class KVConnector: def set_disabled(self, disabled: bool) -> None: pass + def reset_capture_state(self) -> None: + pass + class ActiveKVConnector(KVConnector): def __init__( @@ -54,10 +60,10 @@ class ActiveKVConnector(KVConnector): self.kv_connector.register_kv_caches(kv_caches_dict) self.kv_connector.set_host_xfer_buffer_ops(copy_kv_blocks) - self._pending_load_start = False + self._pending_load_kwargs: dict[str, Any] | None = None self._disabled = False - def pre_forward(self, scheduler_output: "SchedulerOutput") -> None: + def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> None: if self._disabled: return @@ -65,23 +71,31 @@ class ActiveKVConnector(KVConnector): assert kv_connector_metadata is not None self.kv_connector.handle_preemptions(kv_connector_metadata) self.kv_connector.bind_connector_metadata(kv_connector_metadata) + self._pending_load_kwargs = kwargs if scheduler_output.has_sync_kv_loads: # Sync loads need to run before this step's forward. self._start_load_kv() - else: - # Start any async loads in post-forward instead, keeping - # their host-side submission cost off the critical path. - self._pending_load_start = True + # Otherwise start the async load after forward to keep submission + # off the critical path. def _start_load_kv(self) -> None: - self._pending_load_start = False + load_kwargs = self._pending_load_kwargs + assert load_kwargs is not None + self._pending_load_kwargs = None # TODO: sort out KV Connectors' use of forward_context if is_forward_context_available(): - self.kv_connector.start_load_kv(get_forward_context()) + self.kv_connector.start_load_kv(get_forward_context(), **load_kwargs) else: with set_forward_context(None, self.vllm_config): - self.kv_connector.start_load_kv(get_forward_context()) + self.kv_connector.start_load_kv(get_forward_context(), **load_kwargs) + + def finish_forward(self) -> None: + if not self._disabled: + self.kv_connector.finish_forward() + + def reset_capture_state(self) -> None: + self.kv_connector.reset_capture_state() def post_forward( self, finished_req_ids: set[str], wait_for_save: bool = True @@ -89,15 +103,16 @@ class ActiveKVConnector(KVConnector): if self._disabled: return None - if self._pending_load_start: + if self._pending_load_kwargs is not None: self._start_load_kv() output = KVConnectorOutput() if wait_for_save: self.kv_connector.wait_for_save() - output.finished_sending, output.finished_recving = ( - self.kv_connector.get_finished(finished_req_ids) - ) + transfer_results = self.kv_connector.get_transfer_results(finished_req_ids) + output.finished_sending = transfer_results.finished_sending or None + output.finished_recving = transfer_results.finished_recving or None + output.failed_recving = transfer_results.failed_recving output.invalid_block_ids = self.kv_connector.get_block_ids_with_load_errors() output.kv_connector_stats = self.kv_connector.get_kv_connector_stats() output.kv_cache_events = self.kv_connector.get_kv_connector_kv_cache_events() @@ -112,6 +127,7 @@ class ActiveKVConnector(KVConnector): return EMPTY_MODEL_RUNNER_OUTPUT self.pre_forward(scheduler_output) + self.finish_forward() finished_req_ids = scheduler_output.finished_req_ids kv_connector_output = self.post_forward(finished_req_ids, wait_for_save=False) return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index fcadf9e303..14e3823c04 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -67,7 +67,6 @@ from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import ( - CircularBufferSpec, KVCacheConfig, MambaSpec, UniformTypeKVCacheSpecs, @@ -590,7 +589,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): layer_spec = ( spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec ) - slot_mapping_enabled.append(not isinstance(layer_spec, CircularBufferSpec)) + slot_mapping_enabled.append(layer_spec.uses_slot_mapping) # Let each cache type account for CP. Attention KV is DCP-sharded, # while Mamba/GDN recurrent state is replicated across DCP ranks. max_num_blocks = spec.max_num_blocks_per_req( @@ -598,12 +597,11 @@ class GPUModelRunner(LoRAModelRunnerMixin): ) # Preserve each cache type's alignment requirements after applying # its topology-aware block-table width. - if isinstance(layer_spec, (MambaSpec, CircularBufferSpec)): - max_num_blocks = get_block_table_width( - max_num_blocks, spec.block_size, token_alignment=None - ) - else: - max_num_blocks = get_block_table_width(max_num_blocks, spec.block_size) + max_num_blocks = get_block_table_width( + max_num_blocks, + spec.block_size, + token_alignment=spec.block_table_token_alignment, + ) max_num_blocks_per_group.append(max_num_blocks) target_attn_layer_names = None @@ -722,18 +720,20 @@ class GPUModelRunner(LoRAModelRunnerMixin): # to its own attention support. self.speculator.init_cudagraph_manager(cudagraph_mode) - self.kv_caches: list[torch.Tensor] = [] # Capture warmup providers that depend on allocated KV-cache strides. with self.jit_warmup_registry.activate(): kv_caches_dict = init_kv_cache( - self.kv_caches, self.compilation_config.static_forward_context, self.kv_cache_config, self.device, self.kernel_block_sizes, self.vllm_config, kv_cache_allocation_context=kv_cache_allocation_context, + block_tables=self.block_tables, ) + self.kv_caches = [ + cache for cache in kv_caches_dict.values() if cache.device == self.device + ] if is_profiling: self.kv_connector = NO_OP_KV_CONNECTOR else: @@ -1014,6 +1014,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): ): self._dummy_run(**batch) self.adaptive_verification.set_initial_cost_curves(timings) + self.kv_connector.reset_capture_state() end_free_gpu_memory = torch.accelerator.get_memory_info()[0] @@ -1141,13 +1142,17 @@ class GPUModelRunner(LoRAModelRunnerMixin): def update_requests(self, scheduler_output: SchedulerOutput) -> None: # Add new blocks and update num_computed_tokens for the existing requests. reqs = scheduler_output.scheduled_cached_reqs + table_updates = scheduler_output.block_table_updates or {} + for req_id, block_ids in table_updates.items(): + if (idx := self.req_states.req_id_to_index.get(req_id)) is not None: + self.block_tables.append_block_ids(idx, block_ids, overwrite=True) num_computed_tokens_np = self.req_states.num_computed_tokens_np for req_id, num_computed_tokens, req_new_block_ids in zip( reqs.req_ids, reqs.num_computed_tokens, reqs.new_block_ids ): req_index = self.req_states.req_id_to_index[req_id] num_computed_tokens_np[req_index] = num_computed_tokens - if req_new_block_ids is not None: + if req_new_block_ids is not None and req_id not in table_updates: self.block_tables.append_block_ids( req_index, req_new_block_ids, overwrite=False ) @@ -1841,13 +1846,22 @@ class GPUModelRunner(LoRAModelRunnerMixin): ) self.step_timing.forward_start() + connector_kwargs = dict( + scheduler_output=scheduler_output, + request_state_indices=input_batch.idx_mapping, + request_ids=input_batch.req_ids, + num_tokens=input_batch.num_tokens, + ) + # Run model. if batch_desc.cg_mode == CUDAGraphMode.FULL: # Use explicit cudagraph replay for FULL mode. # NOTE(woosuk): Here, we don't need to pass the input tensors, # because they are already copied to the CUDA graph input buffers. assert self.cudagraph_manager is not None - self.kv_connector.pre_forward(scheduler_output) + self.kv_connector.pre_forward( + **connector_kwargs, attn_metadata=attn_metadata + ) model_output = self.cudagraph_manager.run_fullgraph(batch_desc) else: # For piecewise and eager mode, just call model(). @@ -1871,7 +1885,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): skip_compiled=skip_compiled, is_padding=input_batch.is_padding, ): - self.kv_connector.pre_forward(scheduler_output) + self.kv_connector.pre_forward(**connector_kwargs) if ubatch_state is not None: assert self.ubatch_runner is not None model_output = self.ubatch_runner.run( @@ -1889,6 +1903,8 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Eager (NONE): call the raw model directly. model_output = self.model(**model_inputs) + self.kv_connector.finish_forward() + if self.is_last_pp_rank: if self.use_aux_hidden_state_outputs: assert isinstance(model_output, tuple) diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 7072210cb7..56f7d438fc 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -224,10 +224,20 @@ def warmup_kernels( # fixed draft counts here, then restore the manager for capture and serving. adaptive_verification = model_runner.adaptive_verification model_runner.adaptive_verification = None + rejection_sampler = model_runner.rejection_sampler + adaptive_sampling = ( + rejection_sampler is not None and rejection_sampler.enable_adaptive_verification + ) + if adaptive_sampling: + assert rejection_sampler is not None + rejection_sampler.enable_adaptive_verification = False try: _warmup_kernels(model_runner, worker_execute_model, worker_sample_tokens) finally: model_runner.adaptive_verification = adaptive_verification + if adaptive_sampling: + assert rejection_sampler is not None + rejection_sampler.enable_adaptive_verification = True def _warmup_kernels( diff --git a/vllm/v1/worker/kv_connector_model_runner_mixin.py b/vllm/v1/worker/kv_connector_model_runner_mixin.py index 5d6235bf33..6054bb129a 100644 --- a/vllm/v1/worker/kv_connector_model_runner_mixin.py +++ b/vllm/v1/worker/kv_connector_model_runner_mixin.py @@ -94,9 +94,12 @@ class KVConnectorModelRunnerMixin: if wait_for_save and not defer_finalize: kv_connector.wait_for_save() - output.finished_sending, output.finished_recving = ( - kv_connector.get_finished(scheduler_output.finished_req_ids) + transfer_results = kv_connector.get_transfer_results( + scheduler_output.finished_req_ids ) + output.finished_sending = transfer_results.finished_sending + output.finished_recving = transfer_results.finished_recving + output.failed_recving = transfer_results.failed_recving output.invalid_block_ids = kv_connector.get_block_ids_with_load_errors() output.kv_connector_stats = kv_connector.get_kv_connector_stats() diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 067a2b0d22..df2ebc7d36 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -166,6 +166,8 @@ class KVBlockZeroer: kv = static_forward_context[layer_name].kv_cache if not isinstance(kv, torch.Tensor): continue + if kv.device.type != self.device.type: + continue dp = kv.data_ptr() assert kv.shape[0] % num_blocks == 0, ( @@ -366,32 +368,22 @@ def select_common_block_size( return False return True - # Case 1: if the block_size of kv cache manager is supported by all backends, - # return it directly. if block_size_is_supported(backends, kv_manager_block_size): return kv_manager_block_size - # Case 2: otherwise, the block_size must be an `int`-format supported size of - # at least one backend. Iterate over all `int`-format supported sizes in - # descending order and return the first one that is supported by all backends. - # Simple proof: - # If the supported size b is in MultipleOf(x_i) format for all attention - # backends i, and b a factor of kv_manager_block_size, then - # kv_manager_block_size also satisfies MultipleOf(x_i) for all i. We will - # return kv_manager_block_size in case 1. - all_int_supported_sizes = set( - supported_size + # MultipleOf constraints also accept the manager size if they accept a divisor. + # Any remaining candidate must therefore be an explicit size from a backend. + candidates = { + size for backend in backends - for supported_size in backend.get_supported_kernel_block_sizes() - if isinstance(supported_size, int) - ) + for size in backend.get_supported_kernel_block_sizes() + if isinstance(size, int) and kv_manager_block_size % size == 0 + } - for supported_size in sorted(all_int_supported_sizes, reverse=True): - if kv_manager_block_size % supported_size != 0: - continue - if block_size_is_supported(backends, supported_size): - return supported_size - raise ValueError(f"No common block size for {kv_manager_block_size}. ") + for size in sorted(candidates, reverse=True): + if block_size_is_supported(backends, size): + return size + raise ValueError(f"No common block size for {kv_manager_block_size}.") def allocate_kv_cache( @@ -440,7 +432,11 @@ def allocate_kv_cache( if isinstance(spec, UniformTypeKVCacheSpecs): spec = spec.kv_cache_specs[layer_name] - num_blocks = kv_cache_config.num_blocks + if not spec.has_layer_views: + kv_caches.update((name, buf) for name in tensor.layers) + continue + + num_blocks = kv_cache_config.num_blocks_of(tensor) kernel_block_size = None if kernel_block_sizes is not None and group_id < len(kernel_block_sizes): kernel_block_size = kernel_block_sizes[group_id] @@ -485,7 +481,9 @@ def prepare_kernel_block_sizes( kv_cache_spec = next(iter(kv_cache_spec.kv_cache_specs.values())) if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec): continue - if isinstance(kv_cache_spec, AttentionSpec): + if not kv_cache_spec.has_layer_views: + kernel_block_sizes.append(kv_cache_spec.block_size) + elif isinstance(kv_cache_spec, AttentionSpec): # This is an attention backend that supports virtual block splitting. kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size group_backends = [g.backend for g in attn_groups[kv_cache_gid]] @@ -621,7 +619,6 @@ def bind_kv_cache( for layer_name in kv_caches: index2name[extract_layer_index(layer_name, num_attn_module)].append(layer_name) - ordered_layer_names: list[str] = [] for layer_index in sorted(index2name.keys()): layer_names = index2name[layer_index] if len(layer_names) > 1: @@ -635,8 +632,19 @@ def bind_kv_cache( current_platform.check_runner_kv_caches_multi_layer() for layer_name in layer_names: runner_kv_caches.append(kv_caches[layer_name]) - ordered_layer_names.append(layer_name) + bind_kv_cache_to_layers( + kv_caches, forward_context, num_attn_module, kv_cache_groups + ) + + +def bind_kv_cache_to_layers( + kv_caches: dict[str, torch.Tensor], + forward_context: dict[str, Attention], + num_attn_module: int = 1, + kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None, +) -> None: + """Bind layer caches and share ReplaySSM trackers in model-layer order.""" # Bind kv_caches to forward context. Each layer's bind_kv_cache unpacks # its raw allocation into the per-layer view(s) it needs (e.g. Mamba # splits conv/ssm), so the kv_caches dict can hold a single tensor per @@ -644,6 +652,9 @@ def bind_kv_cache( for layer_name, kv_cache in kv_caches.items(): forward_context[layer_name].bind_kv_cache(kv_cache) + ordered_layer_names = sorted( + kv_caches, key=lambda name: extract_layer_index(name, num_attn_module) + ) share_replayssm_ring_trackers(ordered_layer_names, forward_context, kv_cache_groups)