[3/N] HiSparse: host-resident sparse-MLA decode hot-buffering (#53781)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com> Signed-off-by: Matthew Bonanni <matthew@neuralmagic.com> Signed-off-by: Nicolò Lucchesi <nicolo.lucchesi@mistral.ai> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Lucas Wilkinson <lwilkinson@neuralmagic.com> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@mistral.ai> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Matthew Bonanni committed
Sep 12, 2026 at 00:36 UTC
d43bb2f37f87a63d3a9a299c97050443cf6a2520
124 files changed
+14373
-1017
CMakeLists.txt
+1
@@ -476,6 +476,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
476
477
if(VLLM_GPU_LANG STREQUAL "CUDA")
478
list(APPEND VLLM_STABLE_EXT_SRC
479
+ "csrc/libtorch_stable/hisparse_kernels.cu"
480
"csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu"
481
"csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu"
482
"csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu")
csrc/libtorch_stable/cache_kernels.cu
+70
-5
@@ -1458,7 +1458,11 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache_page(
1458
const int32_t block_size, const int32_t total_tokens,
1459
const int64_t block_table_stride, const int64_t cache_block_stride,
1460
const int64_t cache_entry_stride, const int64_t dst_entry_stride,
1461
- const int32_t* __restrict__ seq_starts) {
1461
+ const int32_t* __restrict__ seq_starts,
1462
+ const uint8_t* __restrict__ host_cache,
1463
+ const int32_t* __restrict__ host_row_ids,
1464
+ const int32_t* __restrict__ device_row_ids,
1465
+ const int64_t host_entry_stride) {
1466
constexpr int32_t warps_per_cta = 16;
1467
__shared__ GatherPageTask page;
1468
__shared__ int32_t physical_block;
@@ -1484,8 +1488,22 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache_page(
1488
page_token < page.page_token_end; page_token += warps_per_cta) {
1489
const int32_t output_token =
1490
page.output_token_begin + page_token - page.page_token_begin;
1487
- const uint8_t* token_ptr = src_cache + physical_block * cache_block_stride +
1488
- page_token * cache_entry_stride;
1491
+ const int32_t source_row = physical_block * block_size + page_token;
1492
+ const uint8_t* token_ptr;
1493
+ if (device_row_ids == nullptr) {
1494
+ token_ptr = src_cache + physical_block * cache_block_stride +
1495
+ page_token * cache_entry_stride;
1496
+ } else {
1497
+ const int32_t device_row = device_row_ids[source_row];
1498
+ if (device_row >= 0) {
1499
+ token_ptr = src_cache + (device_row / block_size) * cache_block_stride +
1500
+ (device_row % block_size) * cache_entry_stride;
1501
+ } else {
1502
+ token_ptr =
1503
+ host_cache +
1504
+ static_cast<int64_t>(host_row_ids[source_row]) * host_entry_stride;
1505
+ }
1506
+ }
1507
__nv_bfloat16* dst_ptr = dst + output_token * dst_entry_stride;
1508
gather_and_upconvert_fp8_token(token_ptr, dst_ptr, lane_id);
1509
}
@@ -1694,7 +1712,10 @@ void cp_gather_and_upconvert_fp8_kv_cache(
1712
torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
1713
torch::stable::Tensor const& workspace_starts, // [BATCH]
1714
int64_t batch_size,
1697
- std::optional<torch::stable::Tensor> seq_starts = std::nullopt) {
1715
+ std::optional<torch::stable::Tensor> seq_starts = std::nullopt,
1716
+ std::optional<torch::stable::Tensor> host_cache = std::nullopt,
1717
+ std::optional<torch::stable::Tensor> host_row_ids = std::nullopt,
1718
+ std::optional<torch::stable::Tensor> device_row_ids = std::nullopt) {
1719
torch::stable::accelerator::DeviceGuard device_guard(
1720
src_cache.get_device_index());
1721
const cudaStream_t stream = get_current_cuda_stream();
@@ -1714,6 +1735,12 @@ void cp_gather_and_upconvert_fp8_kv_cache(
1735
"seq_starts must be int32");
1736
}
1737
1738
+ const bool has_host_rows = host_cache.has_value();
1739
+ STD_TORCH_CHECK(
1740
+ has_host_rows == host_row_ids.has_value() &&
1741
+ has_host_rows == device_row_ids.has_value(),
1742
+ "host_cache, host_row_ids, and device_row_ids must be provided together");
1743
+
1744
STD_TORCH_CHECK(src_cache.device() == dst.device(),
1745
"src_cache and dst must be on the same device");
1746
STD_TORCH_CHECK(src_cache.device() == block_table.device(),
@@ -1724,6 +1751,43 @@ void cp_gather_and_upconvert_fp8_kv_cache(
1751
STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
1752
"src_cache and seq_starts must be on the same device");
1753
}
1754
+ const uint8_t* host_cache_ptr = nullptr;
1755
+ const int32_t* host_row_ids_ptr = nullptr;
1756
+ const int32_t* device_row_ids_ptr = nullptr;
1757
+ int64_t host_entry_stride = 0;
1758
+ if (has_host_rows) {
1759
+ auto const& host = host_cache.value();
1760
+ auto const& host_rows = host_row_ids.value();
1761
+ auto const& device_rows = device_row_ids.value();
1762
+ STD_TORCH_CHECK(
1763
+ host.device().is_cpu() &&
1764
+ host.scalar_type() == torch::headeronly::ScalarType::Byte &&
1765
+ host.dim() == 2 && host.is_contiguous() &&
1766
+ host.size(1) == src_cache.size(2),
1767
+ "host_cache must be a contiguous uint8 CPU row matrix matching "
1768
+ "src_cache");
1769
+ cudaPointerAttributes attributes{};
1770
+ const cudaError_t pointer_status =
1771
+ cudaPointerGetAttributes(&attributes, host.const_data_ptr());
1772
+ if (pointer_status != cudaSuccess) {
1773
+ cudaGetLastError();
1774
+ }
1775
+ STD_TORCH_CHECK(
1776
+ pointer_status == cudaSuccess && attributes.type == cudaMemoryTypeHost,
1777
+ "host_cache must be CUDA-accessible pinned memory");
1778
+ STD_TORCH_CHECK(
1779
+ host_rows.is_cuda() && device_rows.is_cuda() &&
1780
+ host_rows.scalar_type() == torch::headeronly::ScalarType::Int &&
1781
+ device_rows.scalar_type() == torch::headeronly::ScalarType::Int &&
1782
+ host_rows.is_contiguous() && device_rows.is_contiguous() &&
1783
+ host_rows.numel() == device_rows.numel() && host_rows.dim() <= 2 &&
1784
+ device_rows.dim() <= 2,
1785
+ "host/device row maps must be matching contiguous CUDA int32 tensors");
1786
+ host_cache_ptr = reinterpret_cast<const uint8_t*>(host.const_data_ptr());
1787
+ host_row_ids_ptr = host_rows.const_data_ptr<int32_t>();
1788
+ device_row_ids_ptr = device_rows.const_data_ptr<int32_t>();
1789
+ host_entry_stride = host.stride(0);
1790
+ }
1791
auto dtype = src_cache.scalar_type();
1792
STD_TORCH_CHECK(
1793
dtype == torch::headeronly::ScalarType::Byte || // uint8
@@ -1766,7 +1830,8 @@ void cp_gather_and_upconvert_fp8_kv_cache(
1830
workspace_starts.const_data_ptr<int32_t>(),
1831
static_cast<int32_t>(batch_size), block_size, total_tokens,
1832
block_table_stride, cache_block_stride, cache_entry_stride,
1769
- dst_entry_stride, seq_starts_ptr);
1833
+ dst_entry_stride, seq_starts_ptr, host_cache_ptr, host_row_ids_ptr,
1834
+ device_row_ids_ptr, host_entry_stride);
1835
}
1836
1837
void cp_gather_and_upconvert_nvfp4_kv_cache(
csrc/libtorch_stable/hisparse_kernels.cu
new
+1226
@@ -0,0 +1,1226 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+//
4
+// HiSparse decode hot-buffer kernels for sparse MLA (GLM-5 / DeepSeek-V3.2).
5
+//
6
+// The swap-in algorithm is a port of SGLang's hisparse
7
+// load_cache_to_device_buffer kernel (sgl jit_kernel/csrc/hisparse.cuh),
8
+// adapted to vLLM addressing:
9
+// - tokens are keyed by their global KV slot id (block_table-converted
10
+// indexer output) instead of in-request positions, so no per-request
11
+// host-location table is needed: host pool row i is global slot i.
12
+// - each batch row owns a fixed region of `region_stride` hot rows;
13
+// slots [0, hot_size) are LRU-managed, slot `hot_size` holds the row's
14
+// newest token (written directly by the KV-cache update).
15
+//
16
+// The kernels only move bytes; they are dtype agnostic.
17
+
18
+#include "torch_utils.h"
19
+#include "ops.h"
20
+#include "../cuda_utils.h"
21
+
22
+#include <algorithm>
23
+#include <cstdint>
24
+#include <optional>
25
+
26
+#include <torch/csrc/stable/c/shim.h>
27
+#include <torch/csrc/stable/ops.h>
28
+#include <torch/headeronly/version.h>
29
+
30
+namespace {
31
+
32
+constexpr int kWarpSize = 32;
33
+// Sentinel in the shared top-k scratch: entry already resolved (hit /
34
+// newest / invalid), no miss handling needed.
35
+constexpr int32_t kTokenDone = -1;
36
+constexpr int32_t kHashEmpty = -1;
37
+
38
+bool is_pinned_cpu_tensor(const torch::stable::Tensor& tensor) {
39
+ cudaPointerAttributes attributes{};
40
+ const auto status =
41
+ cudaPointerGetAttributes(&attributes, tensor.const_data_ptr());
42
+ if (status != cudaSuccess) {
43
+ cudaGetLastError();
44
+ return false;
45
+ }
46
+ return attributes.type == cudaMemoryTypeHost;
47
+}
48
+
49
+__device__ __forceinline__ int32_t hash_slot(int32_t key, int32_t hash_size) {
50
+ // Knuth multiplicative hash for the open-addressing table.
51
+ return static_cast<int32_t>((static_cast<uint32_t>(key) * 2654435761u) %
52
+ static_cast<uint32_t>(hash_size));
53
+}
54
+
55
+// Copy one row of `row_bytes` bytes with a single warp. MLA rows take the
56
+// 16-byte vectorized path; packed indexer rows use the scalar tail-safe path.
57
+__device__ __forceinline__ void copy_row_warp(int lane_id, const char* src,
58
+ char* dst, int64_t row_bytes) {
59
+ const auto alignment = reinterpret_cast<uintptr_t>(src) |
60
+ reinterpret_cast<uintptr_t>(dst) |
61
+ static_cast<uintptr_t>(row_bytes);
62
+ if ((alignment & 15) == 0) {
63
+ const int64_t num_vec = row_bytes / 16;
64
+ const uint4* src4 = reinterpret_cast<const uint4*>(src);
65
+ uint4* dst4 = reinterpret_cast<uint4*>(dst);
66
+ for (int64_t j = lane_id; j < num_vec; j += kWarpSize) {
67
+ __stcg(dst4 + j, __ldcg(src4 + j));
68
+ }
69
+ return;
70
+ }
71
+
72
+ if ((alignment & 3) == 0) {
73
+ const int64_t num_words = row_bytes / 4;
74
+ const unsigned int* src_words = reinterpret_cast<const unsigned int*>(src);
75
+ unsigned int* dst_words = reinterpret_cast<unsigned int*>(dst);
76
+ for (int64_t j = lane_id; j < num_words; j += kWarpSize) {
77
+ __stcg(dst_words + j, __ldcg(src_words + j));
78
+ }
79
+ return;
80
+ }
81
+
82
+ for (int64_t j = lane_id; j < row_bytes; j += kWarpSize) {
83
+ __stcg(dst + j, __ldcg(src + j));
84
+ }
85
+}
86
+
87
+// Zero one row with a single warp (16B vectorized, L2-only stores), same
88
+// addressing contract as copy_row_warp.
89
+__device__ __forceinline__ void zero_row_warp(int lane_id, char* dst,
90
+ int64_t row_bytes) {
91
+ const auto alignment =
92
+ reinterpret_cast<uintptr_t>(dst) | static_cast<uintptr_t>(row_bytes);
93
+ if ((alignment & 15) == 0) {
94
+ const int64_t num_vec = row_bytes / 16;
95
+ uint64_t* dst8 = reinterpret_cast<uint64_t*>(dst);
96
+ for (int64_t j = lane_id; j < num_vec; j += kWarpSize) {
97
+ uint64_t* d = dst8 + j * 2;
98
+ asm volatile("st.global.cg.v2.b64 [%0],{%1,%2};" ::"l"(d), "l"(0ULL),
99
+ "l"(0ULL)
100
+ : "memory");
101
+ }
102
+ return;
103
+ }
104
+
105
+ if ((alignment & 3) == 0) {
106
+ const int64_t num_words = row_bytes / 4;
107
+ unsigned int* dst_words = reinterpret_cast<unsigned int*>(dst);
108
+ for (int64_t j = lane_id; j < num_words; j += kWarpSize) {
109
+ __stcg(dst_words + j, 0u);
110
+ }
111
+ return;
112
+ }
113
+
114
+ for (int64_t j = lane_id; j < row_bytes; j += kWarpSize) {
115
+ __stcg(dst + j, static_cast<char>(0));
116
+ }
117
+}
118
+
119
+__device__ __forceinline__ char* cache_row_ptr(char* cache, int64_t row,
120
+ int32_t block_size,
121
+ int64_t block_stride,
122
+ int64_t row_bytes) {
123
+ return cache + (row / block_size) * block_stride +
124
+ (row % block_size) * row_bytes;
125
+}
126
+
127
+__device__ __forceinline__ void zero_cache_row_warp(int lane_id, char* cache,
128
+ int64_t row,
129
+ int32_t block_size,
130
+ int64_t block_stride,
131
+ int64_t row_bytes) {
132
+ zero_row_warp(lane_id,
133
+ cache_row_ptr(cache, row, block_size, block_stride, row_bytes),
134
+ row_bytes);
135
+}
136
+
137
+// In-place inclusive scan over s_data[offset, count) performed by warp 0,
138
+// carrying `accumulator` across calls. Returns the running total.
139
+__device__ __forceinline__ int warp_inclusive_scan(int32_t* s_data, int lane_id,
140
+ int offset, int count,
141
+ int accumulator) {
142
+ int idx = lane_id + offset;
143
+ int val = (idx < count) ? s_data[idx] : 0;
144
+#pragma unroll
145
+ for (int i = 1; i < 32; i *= 2) {
146
+ int n = __shfl_up_sync(0xffffffff, val, i);
147
+ if (lane_id >= i) val += n;
148
+ }
149
+ val += accumulator;
150
+ if (idx < count) {
151
+ s_data[idx] = val;
152
+ }
153
+ return __shfl_sync(0xffffffff, val, 31);
154
+}
155
+
156
+__device__ __forceinline__ int64_t
157
+get_physical_hot_row(const int32_t* hot_block_table, int32_t row,
158
+ int64_t table_stride, int32_t block_size, int32_t slot) {
159
+ const int32_t block =
160
+ hot_block_table[static_cast<int64_t>(row) * table_stride +
161
+ slot / block_size];
162
+ return static_cast<int64_t>(block) * block_size + slot % block_size;
163
+}
164
+
165
+__device__ __forceinline__ void store_hot_index(
166
+ int32_t* hot_indices, int32_t* attention_indices, int32_t index,
167
+ int32_t physical_row, int32_t hot_block_size,
168
+ int64_t attention_block_stride) {
169
+ hot_indices[index] = physical_row;
170
+ if (attention_indices != nullptr) {
171
+ attention_indices[index] =
172
+ physical_row < 0
173
+ ? -1
174
+ : (physical_row / hot_block_size) * attention_block_stride +
175
+ physical_row % hot_block_size;
176
+ }
177
+}
178
+
179
+// One block per batch row.
180
+//
181
+// Shared memory layout (int32 region followed by int16 region):
182
+// s_topk[top_k] top-k global ids; reused as miss scratch
183
+// s_chunk_off[nbc + 1] prefix sums for hit (then miss) compaction
184
+// s_evict_off[nbc + 1] prefix sums for evictable compaction
185
+// s_hash_keys[hash_size] open addressing: global id -> top-k index
186
+// s_counters[3] hits, phase-1 resolved, valid count
187
+// s_lru_out[hot_size] int16, compacted slots: [hits fwd | evict bwd]
188
+// s_hash_vals[hash_size] int16 hash values (top-k index)
189
+// Valid global ids must be unique within each row.
190
+__global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
191
+ const int32_t* __restrict__ hot_block_table, // [max_rows, hot_blocks]
192
+ const int32_t* __restrict__ global_indices, // global or request-relative
193
+ const int32_t* __restrict__ request_ids, // [num_rows] or nullptr
194
+ const int32_t* __restrict__ source_block_table, // [num_reqs, max_blocks]
195
+ const int32_t* __restrict__ resident_block_table,
196
+ int32_t* __restrict__ resolved_global_indices, // [num_rows, top_k]
197
+ int32_t* __restrict__ valid_counts, // [num_rows] or nullptr
198
+ int32_t* __restrict__ swap_host_physical_rows, // [num_rows, top_k]
199
+ int32_t* __restrict__ swap_device_physical_rows, // [num_rows, top_k]
200
+ int32_t* __restrict__ swap_counts, // [num_rows]
201
+ int32_t* __restrict__ hot_indices, // [num_rows, top_k]
202
+ int32_t* __restrict__ attention_indices, // [num_rows, top_k]
203
+ int32_t* __restrict__ miss_mask, // [num_rows, top_k] or nullptr
204
+ int32_t* __restrict__ device_global_indices, // [max_rows, region_stride]
205
+ int16_t* __restrict__ lru_slots, // [max_rows, hot_size]
206
+ const int32_t* __restrict__ request_state_indices, // [num_requests] or
207
+ // nullptr
208
+ const int32_t request_state_count, const int64_t host_rows,
209
+ const int64_t hot_table_stride, const int32_t hot_block_size,
210
+ const int32_t top_k, const int32_t hot_size, const int32_t hash_size,
211
+ const int64_t region_stride, const int64_t attention_block_stride,
212
+ const int64_t source_bt_stride, const int32_t source_num_reqs,
213
+ const int32_t source_num_blocks, const int32_t source_block_size,
214
+ const int64_t resident_bt_stride, const int32_t resident_num_reqs,
215
+ const int32_t resident_num_blocks, const int32_t resident_block_size,
216
+ const int32_t resident_null_block, const int64_t input_row_stride,
217
+ const int64_t attention_row_stride, const int64_t valid_count_stride) {
218
+ const int NUM_WARPS = blockDim.x / kWarpSize;
219
+ const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize;
220
+ const int num_token_chunks = (top_k + kWarpSize - 1) / kWarpSize;
221
+
222
+ const int batch_row = blockIdx.x;
223
+ const int request_row =
224
+ request_ids != nullptr ? request_ids[batch_row] : batch_row;
225
+ const int state_row =
226
+ request_state_indices != nullptr && request_row >= 0 &&
227
+ request_row < request_state_count
228
+ ? request_state_indices[request_row]
229
+ : (request_state_indices == nullptr ? request_row : -1);
230
+ // V2 publishes -1 for CUDA-graph padding rows.
231
+ if (state_row < 0) {
232
+ for (int i = threadIdx.x; i < top_k; i += blockDim.x) {
233
+ const int64_t index = static_cast<int64_t>(batch_row) * top_k + i;
234
+ hot_indices[index] = -1;
235
+ if (attention_indices != nullptr) {
236
+ attention_indices[static_cast<int64_t>(batch_row) *
237
+ attention_row_stride +
238
+ i] = -1;
239
+ }
240
+ if (resolved_global_indices != nullptr) {
241
+ resolved_global_indices[index] = -1;
242
+ }
243
+ }
244
+ if (valid_counts != nullptr && threadIdx.x == 0) {
245
+ valid_counts[static_cast<int64_t>(batch_row) * valid_count_stride] = 0;
246
+ }
247
+ if (swap_counts != nullptr && threadIdx.x == 0) {
248
+ swap_counts[batch_row] = 0;
249
+ }
250
+ return;
251
+ }
252
+ const int tid = threadIdx.x;
253
+ const int warp_id = tid / kWarpSize;
254
+ const int lane_id = tid % kWarpSize;
255
+ const unsigned int lanes_before = ((unsigned int)1 << lane_id) - 1;
256
+
257
+ const int32_t* row_topk =
258
+ global_indices + static_cast<int64_t>(batch_row) * input_row_stride;
259
+ int32_t* row_out = hot_indices + static_cast<int64_t>(batch_row) * top_k;
260
+ int32_t* row_attention =
261
+ attention_indices != nullptr
262
+ ? attention_indices +
263
+ static_cast<int64_t>(batch_row) * attention_row_stride
264
+ : nullptr;
265
+ int32_t* row_miss = (miss_mask != nullptr)
266
+ ? miss_mask + static_cast<int64_t>(batch_row) * top_k
267
+ : nullptr;
268
+ int32_t* row_dgi =
269
+ device_global_indices + static_cast<int64_t>(state_row) * region_stride;
270
+ int16_t* row_lru = lru_slots + static_cast<int64_t>(state_row) * hot_size;
271
+
272
+ extern __shared__ char smem_raw[];
273
+ int32_t* s_topk = reinterpret_cast<int32_t*>(smem_raw);
274
+ int32_t* s_chunk_off = s_topk + top_k;
275
+ int32_t* s_evict_off = s_chunk_off + (num_buffer_chunks + 1);
276
+ int32_t* s_hash_keys = s_evict_off + (num_buffer_chunks + 1);
277
+ int32_t* s_counters = s_hash_keys + hash_size;
278
+ int16_t* s_lru_out = reinterpret_cast<int16_t*>(s_counters + 3);
279
+ int16_t* s_hash_vals = s_lru_out + hot_size;
280
+
281
+ if (tid < 3) {
282
+ s_counters[tid] = 0;
283
+ }
284
+ __syncthreads();
285
+
286
+ // Phase 1: translate request-relative positions and resolve resident rows.
287
+ for (int i = tid; i < top_k; i += blockDim.x) {
288
+ const int32_t token_index = row_topk[i];
289
+ int32_t g = token_index;
290
+ int32_t resident_row = -1;
291
+ if (source_block_table != nullptr) {
292
+ const int32_t request_id = request_ids[batch_row];
293
+ const int32_t source_block =
294
+ token_index >= 0 ? token_index / source_block_size : -1;
295
+ if (request_id >= 0 && request_id < source_num_reqs &&
296
+ source_block >= 0 && source_block < source_num_blocks) {
297
+ const int32_t physical_block =
298
+ source_block_table[static_cast<int64_t>(request_id) *
299
+ source_bt_stride +
300
+ source_block];
301
+ g = physical_block > 0 ? physical_block * source_block_size +
302
+ token_index % source_block_size
303
+ : -1;
304
+ } else {
305
+ g = -1;
306
+ }
307
+ if (resident_block_table != nullptr) {
308
+ const int32_t resident_block =
309
+ token_index >= 0 ? token_index / resident_block_size : -1;
310
+ if (request_id >= 0 && request_id < resident_num_reqs &&
311
+ resident_block >= 0 && resident_block < resident_num_blocks) {
312
+ const int32_t physical_block =
313
+ resident_block_table[static_cast<int64_t>(request_id) *
314
+ resident_bt_stride +
315
+ resident_block];
316
+ if (physical_block != resident_null_block && physical_block >= 0) {
317
+ resident_row = physical_block * resident_block_size +
318
+ token_index % resident_block_size;
319
+ }
320
+ }
321
+ }
322
+ }
323
+ if (g >= host_rows) {
324
+ g = -1;
325
+ }
326
+ if (resolved_global_indices != nullptr) {
327
+ resolved_global_indices[static_cast<int64_t>(batch_row) * top_k + i] = g;
328
+ }
329
+ if (resident_row >= 0 || g >= 0) atomicAdd(&s_counters[2], 1);
330
+ if (row_miss != nullptr) row_miss[i] = 0;
331
+ if (resident_row >= 0) {
332
+ store_hot_index(row_out, row_attention, i, resident_row, hot_block_size,
333
+ attention_block_stride);
334
+ s_topk[i] = kTokenDone;
335
+ atomicAdd(&s_counters[1], 1);
336
+ } else if (g < 0) {
337
+ store_hot_index(row_out, row_attention, i, -1, hot_block_size,
338
+ attention_block_stride);
339
+ s_topk[i] = kTokenDone;
340
+ atomicAdd(&s_counters[1], 1);
341
+ } else {
342
+ s_topk[i] = g;
343
+ }
344
+ }
345
+ __syncthreads();
346
+ if (valid_counts != nullptr && tid == 0) {
347
+ valid_counts[static_cast<int64_t>(batch_row) * valid_count_stride] =
348
+ s_counters[2];
349
+ }
350
+ // Fully resident rows need only request-relative page translation. Avoid
351
+ // scanning or rewriting the hot LRU when no selected row can consult it.
352
+ if (s_counters[1] == top_k) {
353
+ if (tid == 0 && swap_counts != nullptr) {
354
+ swap_counts[batch_row] = 0;
355
+ }
356
+ return;
357
+ }
358
+
359
+ for (int i = tid; i < hash_size; i += blockDim.x) {
360
+ s_hash_keys[i] = kHashEmpty;
361
+ }
362
+ for (int i = tid; i < num_buffer_chunks + 1; i += blockDim.x) {
363
+ s_chunk_off[i] = 0;
364
+ s_evict_off[i] = 0;
365
+ }
366
+ __syncthreads();
367
+
368
+ for (int i = tid; i < top_k; i += blockDim.x) {
369
+ const int32_t g = s_topk[i];
370
+ if (g == kTokenDone) {
371
+ continue;
372
+ }
373
+ int slot = hash_slot(g, hash_size);
374
+ while (true) {
375
+ const int32_t old = atomicCAS(&s_hash_keys[slot], kHashEmpty, g);
376
+ if (old == kHashEmpty || old == g) {
377
+ s_hash_vals[slot] = static_cast<int16_t>(i);
378
+ break;
379
+ }
380
+ slot = (slot + 1) % hash_size;
381
+ }
382
+ }
383
+ __syncthreads();
384
+
385
+ // Phase 2: walk hot slots in LRU order, classify hit / evictable, and
386
+ // compact them (hits forward, evictables backward) into s_lru_out.
387
+ const int iters_buffer = (num_buffer_chunks + NUM_WARPS - 1) / NUM_WARPS;
388
+ int total_hit_count = 0;
389
+ int total_evict_count = 0;
390
+ for (int iter = 0; iter < iters_buffer; iter++) {
391
+ const int chunk_idx = warp_id + iter * NUM_WARPS;
392
+ const bool has_valid_chunk = chunk_idx < num_buffer_chunks;
393
+
394
+ const int pos = chunk_idx * kWarpSize + lane_id;
395
+ const bool has_valid_pos = has_valid_chunk && (pos < hot_size);
396
+ const int16_t slot = has_valid_pos ? row_lru[pos] : int16_t(-1);
397
+ // Corruption tripwire: lru/dgi are long-lived device state; if some
398
+ // external writer (e.g. stray RDMA into reused VRAM) corrupts a slot
399
+ // out of [0, region_stride), treat it as no-hit so it degrades to a re-miss
400
+ // instead of an unbounded dgi read (phases 3/5 bound the write side).
401
+ const int32_t cached_g =
402
+ (slot >= 0 && slot < region_stride) ? row_dgi[slot] : -1;
403
+
404
+ int found_topk_idx = -1;
405
+ if (cached_g >= 0) {
406
+ int h = hash_slot(cached_g, hash_size);
407
+ while (true) {
408
+ const int32_t k = s_hash_keys[h];
409
+ if (k == cached_g) {
410
+ found_topk_idx = static_cast<int32_t>(s_hash_vals[h]);
411
+ break;
412
+ }
413
+ if (k == kHashEmpty) break;
414
+ h = (h + 1) % hash_size;
415
+ }
416
+ }
417
+ const bool is_hit = found_topk_idx >= 0;
418
+ const bool is_evictable = has_valid_pos && !is_hit;
419
+
420
+ if (is_hit) {
421
+ s_topk[found_topk_idx] = kTokenDone;
422
+ store_hot_index(row_out, row_attention, found_topk_idx,
423
+ static_cast<int32_t>(get_physical_hot_row(
424
+ hot_block_table, request_row, hot_table_stride,
425
+ hot_block_size, slot)),
426
+ hot_block_size, attention_block_stride);
427
+ }
428
+
429
+ int local_hit_off = 0;
430
+ int local_evict_off = 0;
431
+ if (has_valid_chunk) {
432
+ const unsigned int hit_mask = __ballot_sync(0xFFFFFFFF, is_hit);
433
+ const unsigned int evict_mask = __ballot_sync(0xFFFFFFFF, is_evictable);
434
+ local_hit_off = __popc(hit_mask & lanes_before);
435
+ local_evict_off = __popc(evict_mask & lanes_before);
436
+ if (lane_id == 0) {
437
+ s_chunk_off[chunk_idx + 1] = __popc(hit_mask);
438
+ s_evict_off[chunk_idx + 1] = __popc(evict_mask);
439
+ }
440
+ }
441
+ __syncthreads();
442
+
443
+ if (warp_id == 0) {
444
+ total_hit_count =
445
+ warp_inclusive_scan(s_chunk_off, lane_id, chunk_idx + 1,
446
+ num_buffer_chunks + 1, total_hit_count);
447
+ total_evict_count =
448
+ warp_inclusive_scan(s_evict_off, lane_id, chunk_idx + 1,
449
+ num_buffer_chunks + 1, total_evict_count);
450
+ if (tid == 0) {
451
+ s_counters[0] = total_hit_count;
452
+ }
453
+ }
454
+ __syncthreads();
455
+
456
+ if (is_hit) {
457
+ const int off = s_chunk_off[chunk_idx] + local_hit_off;
458
+ s_lru_out[off] = slot;
459
+ }
460
+ if (is_evictable) {
461
+ const int off = s_evict_off[chunk_idx] + local_evict_off;
462
+ s_lru_out[hot_size - 1 - off] = slot;
463
+ }
464
+ }
465
+ __syncthreads();
466
+
467
+ // Reset prefix sums for the miss compaction (token chunks <= buffer
468
+ // chunks because hot_size >= top_k).
469
+ for (int i = tid; i < num_token_chunks + 1; i += blockDim.x) {
470
+ s_chunk_off[i] = 0;
471
+ }
472
+ __syncthreads();
473
+
474
+ // Phase 3: compact misses, assign them eviction slots (oldest first) and
475
+ // record the new ownership in device_global_indices.
476
+ const int iters_token = (num_token_chunks + NUM_WARPS - 1) / NUM_WARPS;
477
+ int miss_running_total = 0;
478
+ for (int iter = 0; iter < iters_token; iter++) {
479
+ const int chunk_idx = warp_id + iter * NUM_WARPS;
480
+ const bool has_valid_chunk = chunk_idx < num_token_chunks;
481
+
482
+ const int i = chunk_idx * kWarpSize + lane_id;
483
+ const bool has_valid_token = has_valid_chunk && (i < top_k);
484
+
485
+ int32_t g = 0;
486
+ bool is_miss = false;
487
+ if (has_valid_token) {
488
+ is_miss = s_topk[i] != kTokenDone;
489
+ if (is_miss) {
490
+ g = s_topk[i];
491
+ }
492
+ }
493
+
494
+ int local_miss_off = 0;
495
+ if (has_valid_chunk) {
496
+ const unsigned int miss_mask = __ballot_sync(0xFFFFFFFF, is_miss);
497
+ local_miss_off = __popc(miss_mask & lanes_before);
498
+ if (lane_id == 0) {
499
+ s_chunk_off[chunk_idx + 1] = __popc(miss_mask);
500
+ }
501
+ }
502
+ __syncthreads();
503
+
504
+ if (warp_id == 0) {
505
+ miss_running_total =
506
+ warp_inclusive_scan(s_chunk_off, lane_id, chunk_idx + 1,
507
+ num_token_chunks + 1, miss_running_total);
508
+ }
509
+ __syncthreads();
510
+
511
+ if (is_miss) {
512
+ const int m = s_chunk_off[chunk_idx] + local_miss_off;
513
+ const int16_t evict_slot = s_lru_out[hot_size - 1 - m];
514
+ if (evict_slot < 0 || evict_slot >= region_stride) {
515
+ // Corruption tripwire: an out-of-range slot (corrupted lru state,
516
+ // see phase 2) must not become a hot-cache/dgi write. Resolve the
517
+ // entry as invalid (-1, masked by attention, not in the miss set);
518
+ // it re-misses on a later step. Phase 5 re-checks the same
519
+ // s_lru_out value, so its copy is skipped consistently.
520
+ store_hot_index(row_out, row_attention, i, -1, hot_block_size,
521
+ attention_block_stride);
522
+ if (swap_host_physical_rows != nullptr) {
523
+ const int64_t compact_index =
524
+ static_cast<int64_t>(batch_row) * top_k + m;
525
+ swap_host_physical_rows[compact_index] = g;
526
+ swap_device_physical_rows[compact_index] = -1;
527
+ }
528
+ } else {
529
+ // Reuse s_topk as compacted miss scratch: m < i always holds (done
530
+ // entries are skipped), so writes never overrun pending reads.
531
+ s_topk[m] = g;
532
+ const int32_t physical_row = static_cast<int32_t>(
533
+ get_physical_hot_row(hot_block_table, request_row, hot_table_stride,
534
+ hot_block_size, evict_slot));
535
+ store_hot_index(row_out, row_attention, i, physical_row, hot_block_size,
536
+ attention_block_stride);
537
+ if (swap_host_physical_rows != nullptr) {
538
+ const int64_t compact_index =
539
+ static_cast<int64_t>(batch_row) * top_k + m;
540
+ swap_host_physical_rows[compact_index] = g;
541
+ swap_device_physical_rows[compact_index] = physical_row;
542
+ }
543
+ if (row_miss != nullptr) row_miss[i] = 1;
544
+ row_dgi[evict_slot] = g;
545
+ }
546
+ }
547
+ }
548
+ __syncthreads();
549
+
550
+ const int total_hits = s_counters[0];
551
+ const int total_misses = top_k - total_hits - s_counters[1];
552
+ if (swap_counts != nullptr && tid == 0) {
553
+ swap_counts[batch_row] = total_misses;
554
+ }
555
+ // Phase 4: write back the LRU order: stale evictables at the front,
556
+ // freshly loaded misses next, then hits at MRU.
557
+ const int total_evictable = hot_size - total_hits;
558
+ const int remaining_evictable = total_evictable - total_misses;
559
+ for (int i = tid; i < hot_size; i += blockDim.x) {
560
+ if (i < remaining_evictable) {
561
+ row_lru[i] = s_lru_out[hot_size - 1 - total_misses - i];
562
+ } else if (i < remaining_evictable + total_misses) {
563
+ row_lru[i] = s_lru_out[hot_size - 1 - (i - remaining_evictable)];
564
+ } else {
565
+ row_lru[i] = s_lru_out[i - remaining_evictable - total_misses];
566
+ }
567
+ }
568
+}
569
+
570
+// Shared-layer plan replay. Given a plan (hot_indices + miss_mask) already
571
+// computed by the group's index-producing layer, gather THIS
572
+// layer's own missed KV rows into the planned hot slots. No LRU resolution:
573
+// index-sharing shared layers see identical global_indices/hot_indices, so the
574
+// slot assignment is identical -- only the per-layer bytes differ. Fixed shape
575
+// (num_rows x top_k), so it is CUDA-graph-capture safe.
576
+__global__ void hisparse_gather_plan_kernel(
577
+ const char* __restrict__ host_cache, // [host_rows, row_bytes]
578
+ char* __restrict__ hot_cache, // [hot_rows, row_bytes]
579
+ const int32_t* __restrict__ global_indices, // [num_rows, top_k]
580
+ const int32_t* __restrict__ hot_indices, // [num_rows, top_k] abs hot rows
581
+ const int32_t* __restrict__ miss_mask, // [num_rows, top_k]
582
+ int32_t* __restrict__ attention_indices, // [num_rows, top_k]
583
+ const int32_t* __restrict__ request_state_indices, // [num_rows] or nullptr
584
+ const int64_t host_rows, const int64_t hot_rows, const int64_t row_bytes,
585
+ const int64_t hot_block_stride, const int32_t hot_block_size,
586
+ const int32_t top_k, const int64_t attention_block_stride) {
587
+ const int NUM_WARPS = blockDim.x / kWarpSize;
588
+ // Columns are interleaved across gridDim.y blocks per row so few-row
589
+ // launches (local-prefill staging's single-row layout, small decode
590
+ // batches) still fill the device with outstanding host reads instead of
591
+ // starving on one block per row.
592
+ const int row = blockIdx.x;
593
+ const bool is_padding =
594
+ request_state_indices != nullptr && request_state_indices[row] < 0;
595
+ const int warp_id = threadIdx.x / kWarpSize;
596
+ const int lane_id = threadIdx.x % kWarpSize;
597
+ const int col_start = blockIdx.y * NUM_WARPS + warp_id;
598
+ const int col_stride = gridDim.y * NUM_WARPS;
599
+ const int64_t base = static_cast<int64_t>(row) * top_k;
600
+ for (int col = col_start; col < top_k; col += col_stride) {
601
+ const int32_t dst = hot_indices[base + col];
602
+ if (lane_id == 0 && attention_indices != nullptr) {
603
+ attention_indices[base + col] =
604
+ is_padding
605
+ ? -1
606
+ : (dst < 0 ? -1
607
+ : (dst / hot_block_size) * attention_block_stride +
608
+ dst % hot_block_size);
609
+ }
610
+ if (is_padding) {
611
+ continue;
612
+ }
613
+ if (miss_mask[base + col] == 0) {
614
+ continue;
615
+ }
616
+ const int32_t g = global_indices[base + col];
617
+ if (g < 0 || dst < 0 || dst >= hot_rows) {
618
+ continue;
619
+ }
620
+ if (g < host_rows) {
621
+ copy_row_warp(lane_id, host_cache + static_cast<int64_t>(g) * row_bytes,
622
+ cache_row_ptr(hot_cache, dst, hot_block_size,
623
+ hot_block_stride, row_bytes),
624
+ row_bytes);
625
+ } else {
626
+ // No source row for g: zero the planned slot rather than serving
627
+ // whatever bytes it held (see the swap-in kernel's phase 5).
628
+ zero_cache_row_warp(lane_id, hot_cache, dst, hot_block_size,
629
+ hot_block_stride, row_bytes);
630
+ }
631
+ }
632
+}
633
+
634
+// Copy the compact swap rows produced by the residency resolver.
635
+__global__ void hisparse_gather_compact_kernel(
636
+ const char* __restrict__ host_cache, char* __restrict__ hot_cache,
637
+ const int32_t* __restrict__ miss_global_indices,
638
+ const int32_t* __restrict__ miss_hot_indices,
639
+ const int32_t* __restrict__ miss_counts, const int64_t host_rows,
640
+ const int64_t hot_rows, const int64_t row_bytes,
641
+ const int64_t hot_block_stride, const int32_t hot_block_size,
642
+ const int32_t top_k) {
643
+ const int NUM_WARPS = blockDim.x / kWarpSize;
644
+ const int row = blockIdx.x;
645
+ const int warp_id = threadIdx.x / kWarpSize;
646
+ const int lane_id = threadIdx.x % kWarpSize;
647
+ const int miss_count = min(max(miss_counts[row], 0), top_k);
648
+ const int col_start = blockIdx.y * NUM_WARPS + warp_id;
649
+ const int col_stride = gridDim.y * NUM_WARPS;
650
+ const int64_t base = static_cast<int64_t>(row) * top_k;
651
+ for (int col = col_start; col < miss_count; col += col_stride) {
652
+ const int32_t g = miss_global_indices[base + col];
653
+ const int32_t dst = miss_hot_indices[base + col];
654
+ if (g < 0 || dst < 0 || dst >= hot_rows) {
655
+ continue;
656
+ }
657
+ if (g < host_rows) {
658
+ copy_row_warp(lane_id, host_cache + static_cast<int64_t>(g) * row_bytes,
659
+ cache_row_ptr(hot_cache, dst, hot_block_size,
660
+ hot_block_stride, row_bytes),
661
+ row_bytes);
662
+ } else {
663
+ zero_cache_row_warp(lane_id, hot_cache, dst, hot_block_size,
664
+ hot_block_stride, row_bytes);
665
+ }
666
+ }
667
+}
668
+
669
+__global__ void hisparse_invalidate_written_slots_kernel(
670
+ int32_t* __restrict__ device_global_indices,
671
+ const int32_t* __restrict__ request_state_indices,
672
+ const int32_t* __restrict__ req_id_per_token,
673
+ const int64_t* __restrict__ written_slots, const int64_t num_tokens,
674
+ const int64_t num_request_ids, const int64_t num_state_rows,
675
+ const int64_t region_stride) {
676
+ const int64_t token_idx = blockIdx.x;
677
+ if (token_idx >= num_tokens) {
678
+ return;
679
+ }
680
+ const int32_t req_idx = req_id_per_token[token_idx];
681
+ if (req_idx < 0 || req_idx >= num_request_ids) {
682
+ return;
683
+ }
684
+ const int32_t state_idx = request_state_indices[req_idx];
685
+ const int64_t written_slot = written_slots[token_idx];
686
+ if (state_idx < 0 || state_idx >= num_state_rows || written_slot < 0) {
687
+ return;
688
+ }
689
+ int32_t* row = device_global_indices + state_idx * region_stride;
690
+ for (int64_t offset = threadIdx.x; offset < region_stride;
691
+ offset += blockDim.x) {
692
+ if (row[offset] == written_slot) {
693
+ row[offset] = -1;
694
+ }
695
+ }
696
+}
697
+
698
+int64_t check_2d_rows(const torch::stable::Tensor& t, const char* name,
699
+ int64_t row_bytes) {
700
+ STD_TORCH_CHECK(t.dim() == 2, name, " must be 2D");
701
+ STD_TORCH_CHECK(t.is_contiguous(), name, " must be contiguous");
702
+ STD_TORCH_CHECK(t.size(1) * t.element_size() == row_bytes, name,
703
+ " row width mismatch");
704
+ return t.size(0);
705
+}
706
+
707
+} // namespace
708
+
709
+void hisparse_invalidate_written_slots(
710
+ torch::stable::Tensor& device_global_indices,
711
+ torch::stable::Tensor const& request_state_indices,
712
+ torch::stable::Tensor const& req_id_per_token,
713
+ torch::stable::Tensor const& written_slots) {
714
+ STD_TORCH_CHECK(device_global_indices.is_cuda() &&
715
+ request_state_indices.is_cuda() &&
716
+ req_id_per_token.is_cuda() && written_slots.is_cuda(),
717
+ "HiSparse invalidation tensors must be on CUDA");
718
+ STD_TORCH_CHECK(
719
+ device_global_indices.scalar_type() ==
720
+ torch::headeronly::ScalarType::Int &&
721
+ device_global_indices.dim() == 2 &&
722
+ device_global_indices.is_contiguous(),
723
+ "device_global_indices must be a contiguous 2D int32 CUDA tensor");
724
+ STD_TORCH_CHECK(
725
+ request_state_indices.scalar_type() ==
726
+ torch::headeronly::ScalarType::Int &&
727
+ request_state_indices.dim() == 1 &&
728
+ request_state_indices.is_contiguous(),
729
+ "request_state_indices must be a contiguous 1D int32 CUDA tensor");
730
+ STD_TORCH_CHECK(
731
+ req_id_per_token.scalar_type() == torch::headeronly::ScalarType::Int &&
732
+ req_id_per_token.dim() == 1 && req_id_per_token.is_contiguous(),
733
+ "req_id_per_token must be a contiguous 1D int32 CUDA tensor");
734
+ STD_TORCH_CHECK(
735
+ written_slots.scalar_type() == torch::headeronly::ScalarType::Long &&
736
+ written_slots.dim() == 1 && written_slots.is_contiguous(),
737
+ "written_slots must be a contiguous 1D int64 CUDA tensor");
738
+ STD_TORCH_CHECK(req_id_per_token.numel() == written_slots.numel(),
739
+ "req_id_per_token and written_slots must have equal length");
740
+ const int device_index = device_global_indices.get_device_index();
741
+ STD_TORCH_CHECK(
742
+ request_state_indices.get_device_index() == device_index &&
743
+ req_id_per_token.get_device_index() == device_index &&
744
+ written_slots.get_device_index() == device_index,
745
+ "HiSparse invalidation tensors must be on the same CUDA device");
746
+
747
+ const int64_t num_tokens = written_slots.numel();
748
+ const int64_t num_state_rows = device_global_indices.size(0);
749
+ const int64_t region_stride = device_global_indices.size(1);
750
+ if (num_tokens == 0 || num_state_rows == 0 || region_stride == 0) {
751
+ return;
752
+ }
753
+ STD_TORCH_CHECK(num_tokens <= INT32_MAX,
754
+ "HiSparse invalidation grid exceeds CUDA limits");
755
+
756
+ constexpr int kBlockSize = 256;
757
+ const torch::stable::accelerator::DeviceGuard device_guard(device_index);
758
+ const cudaStream_t stream = get_current_cuda_stream();
759
+ hisparse_invalidate_written_slots_kernel<<<static_cast<int>(num_tokens),
760
+ kBlockSize, 0, stream>>>(
761
+ device_global_indices.mutable_data_ptr<int32_t>(),
762
+ request_state_indices.const_data_ptr<int32_t>(),
763
+ req_id_per_token.const_data_ptr<int32_t>(),
764
+ written_slots.const_data_ptr<int64_t>(), num_tokens,
765
+ request_state_indices.numel(), num_state_rows, region_stride);
766
+ const cudaError_t launch_error = cudaGetLastError();
767
+ STD_TORCH_CHECK(launch_error == cudaSuccess,
768
+ "HiSparse invalidation kernel launch failed: ",
769
+ cudaGetErrorString(launch_error));
770
+}
771
+
772
+void hisparse_resolve_residency(
773
+ torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache,
774
+ torch::stable::Tensor const& hot_block_table,
775
+ torch::stable::Tensor const& global_indices,
776
+ torch::stable::Tensor& hot_indices,
777
+ torch::stable::Tensor& device_global_indices,
778
+ torch::stable::Tensor& lru_slots,
779
+ std::optional<torch::stable::Tensor> const& request_state_indices,
780
+ int64_t region_stride,
781
+ std::optional<torch::stable::Tensor> const& miss_mask,
782
+ std::optional<torch::stable::Tensor> const& attention_indices,
783
+ int64_t attention_block_stride,
784
+ std::optional<torch::stable::Tensor> const& request_ids,
785
+ std::optional<torch::stable::Tensor> const& source_block_table,
786
+ int64_t source_block_size,
787
+ std::optional<torch::stable::Tensor> const& resolved_global_indices,
788
+ std::optional<torch::stable::Tensor> const& valid_counts,
789
+ std::optional<torch::stable::Tensor> const& swap_host_physical_rows,
790
+ std::optional<torch::stable::Tensor> const& swap_device_physical_rows,
791
+ std::optional<torch::stable::Tensor> const& swap_counts,
792
+ std::optional<torch::stable::Tensor> const& resident_block_table,
793
+ int64_t resident_block_size, int64_t resident_null_block) {
794
+ STD_TORCH_CHECK(
795
+ host_cache.device().is_cpu() && is_pinned_cpu_tensor(host_cache),
796
+ "host_cache must be pinned CPU memory");
797
+ STD_TORCH_CHECK(hot_cache.is_cuda(), "hot_cache must be on CUDA");
798
+ STD_TORCH_CHECK(
799
+ hot_block_table.is_cuda() &&
800
+ hot_block_table.scalar_type() == torch::headeronly::ScalarType::Int &&
801
+ hot_block_table.dim() == 2,
802
+ "hot_block_table must be a 2D int32 CUDA tensor");
803
+ STD_TORCH_CHECK(global_indices.is_cuda() && hot_indices.is_cuda() &&
804
+ device_global_indices.is_cuda() && lru_slots.is_cuda(),
805
+ "index tensors must be on CUDA");
806
+ STD_TORCH_CHECK(
807
+ global_indices.scalar_type() == torch::headeronly::ScalarType::Int &&
808
+ hot_indices.scalar_type() == torch::headeronly::ScalarType::Int,
809
+ "global_indices/hot_indices must be int32");
810
+ STD_TORCH_CHECK(
811
+ device_global_indices.scalar_type() == torch::headeronly::ScalarType::Int,
812
+ "device_global_indices must be int32");
813
+ STD_TORCH_CHECK(
814
+ lru_slots.scalar_type() == torch::headeronly::ScalarType::Short,
815
+ "lru_slots must be int16");
816
+ STD_TORCH_CHECK(global_indices.dim() == 2 && global_indices.stride(1) == 1,
817
+ "global_indices must be row-major 2D");
818
+ const int64_t num_rows = global_indices.size(0);
819
+ STD_TORCH_CHECK(num_rows >= 0 && num_rows <= INT32_MAX,
820
+ "num_rows must fit the CUDA grid");
821
+ STD_TORCH_CHECK(hot_indices.size(0) == num_rows &&
822
+ hot_indices.size(1) == global_indices.size(1) &&
823
+ hot_indices.is_contiguous(),
824
+ "hot_indices must have one contiguous row per launch row");
825
+ STD_TORCH_CHECK(
826
+ device_global_indices.dim() == 2 && device_global_indices.is_contiguous(),
827
+ "device_global_indices must be contiguous 2D");
828
+ STD_TORCH_CHECK(
829
+ lru_slots.dim() == 2 &&
830
+ lru_slots.size(0) == device_global_indices.size(0) &&
831
+ lru_slots.is_contiguous(),
832
+ "lru_slots must be contiguous with one row per request state");
833
+ STD_TORCH_CHECK(hot_cache.dim() == 3,
834
+ "hot_cache must be [num_blocks, block_size, row_width]");
835
+ const int64_t row_bytes = hot_cache.size(-1) * hot_cache.element_size();
836
+ STD_TORCH_CHECK(row_bytes % 16 == 0, "KV rows must be 16-byte aligned");
837
+ STD_TORCH_CHECK(hot_cache.stride(1) * hot_cache.element_size() == row_bytes,
838
+ "hot-cache rows must be contiguous");
839
+ const int64_t hot_block_size = hot_cache.size(1);
840
+ const int64_t hot_rows = hot_cache.size(0) * hot_block_size;
841
+ const int64_t hot_block_stride =
842
+ hot_cache.stride(0) * hot_cache.element_size();
843
+ const int64_t host_rows = check_2d_rows(host_cache, "host_cache", row_bytes);
844
+ const auto launch_rows = static_cast<int32_t>(num_rows);
845
+ const auto top_k = static_cast<int32_t>(global_indices.size(1));
846
+ const auto hot_size = static_cast<int32_t>(lru_slots.size(1));
847
+ STD_TORCH_CHECK(hot_size >= top_k, "hot buffer size must be >= top_k");
848
+ STD_TORCH_CHECK(hot_size <= 32768, "hot buffer size must fit int16 slots");
849
+ STD_TORCH_CHECK(region_stride == hot_size,
850
+ "region_stride must match the LRU size");
851
+ STD_TORCH_CHECK(device_global_indices.size(1) == region_stride,
852
+ "device_global_indices must cover the full hot region");
853
+ STD_TORCH_CHECK(hot_block_table.size(1) >=
854
+ (region_stride + hot_block_size - 1) / hot_block_size,
855
+ "hot_block_table has too few columns");
856
+ STD_TORCH_CHECK(hot_rows < INT32_MAX, "hot indices must fit int32");
857
+
858
+ const int32_t* request_ids_ptr = nullptr;
859
+ const int32_t* source_block_table_ptr = nullptr;
860
+ int64_t source_bt_stride = 0;
861
+ int32_t source_num_reqs = 0;
862
+ int32_t source_num_blocks = 0;
863
+ STD_TORCH_CHECK(
864
+ request_ids.has_value() == source_block_table.has_value(),
865
+ "request_ids and source_block_table must be provided together");
866
+ if (source_block_table.has_value()) {
867
+ auto const& req = request_ids.value();
868
+ auto const& table = source_block_table.value();
869
+ STD_TORCH_CHECK(
870
+ req.is_cuda() &&
871
+ req.scalar_type() == torch::headeronly::ScalarType::Int &&
872
+ req.numel() >= launch_rows && req.is_contiguous(),
873
+ "request_ids must be contiguous int32 on CUDA with one entry per row");
874
+ STD_TORCH_CHECK(
875
+ table.is_cuda() &&
876
+ table.scalar_type() == torch::headeronly::ScalarType::Int &&
877
+ table.dim() == 2 && table.stride(1) == 1,
878
+ "source_block_table must be a row-major 2D int32 CUDA tensor");
879
+ STD_TORCH_CHECK(source_block_size > 0,
880
+ "source_block_size must be positive");
881
+ request_ids_ptr = req.const_data_ptr<int32_t>();
882
+ source_block_table_ptr = table.const_data_ptr<int32_t>();
883
+ source_bt_stride = table.stride(0);
884
+ source_num_reqs = static_cast<int32_t>(table.size(0));
885
+ source_num_blocks = static_cast<int32_t>(table.size(1));
886
+ }
887
+ const int32_t required_hot_rows =
888
+ request_ids_ptr != nullptr ? source_num_reqs : launch_rows;
889
+ STD_TORCH_CHECK(hot_block_table.size(0) >= required_hot_rows,
890
+ "hot_block_table has too few rows");
891
+
892
+ const int32_t* resident_block_table_ptr = nullptr;
893
+ int64_t resident_bt_stride = 0;
894
+ int32_t resident_num_reqs = 0;
895
+ int32_t resident_num_blocks = 0;
896
+ if (resident_block_table.has_value()) {
897
+ auto const& table = resident_block_table.value();
898
+ STD_TORCH_CHECK(
899
+ source_block_table.has_value() && request_ids.has_value(),
900
+ "resident lookup requires request_ids and source_block_table");
901
+ STD_TORCH_CHECK(
902
+ table.is_cuda() &&
903
+ table.scalar_type() == torch::headeronly::ScalarType::Int &&
904
+ table.dim() == 2 && table.stride(1) == 1,
905
+ "resident_block_table must be a row-major 2D int32 CUDA tensor");
906
+ STD_TORCH_CHECK(resident_block_size == hot_block_size,
907
+ "resident and hot block sizes must match");
908
+ STD_TORCH_CHECK(resident_null_block >= 0,
909
+ "resident null block must be non-negative");
910
+ resident_block_table_ptr = table.const_data_ptr<int32_t>();
911
+ resident_bt_stride = table.stride(0);
912
+ resident_num_reqs = static_cast<int32_t>(table.size(0));
913
+ resident_num_blocks = static_cast<int32_t>(table.size(1));
914
+ }
915
+
916
+ int32_t* resolved_global_indices_ptr = nullptr;
917
+ if (resolved_global_indices.has_value()) {
918
+ auto const& resolved = resolved_global_indices.value();
919
+ STD_TORCH_CHECK(
920
+ resolved.is_cuda() &&
921
+ resolved.scalar_type() == torch::headeronly::ScalarType::Int &&
922
+ resolved.dim() == 2 && resolved.size(0) == launch_rows &&
923
+ resolved.size(1) == global_indices.size(1) &&
924
+ resolved.is_contiguous(),
925
+ "resolved_global_indices must be contiguous int32 matching indices");
926
+ resolved_global_indices_ptr = resolved.mutable_data_ptr<int32_t>();
927
+ }
928
+
929
+ int32_t* valid_counts_ptr = nullptr;
930
+ if (valid_counts.has_value()) {
931
+ auto const& counts = valid_counts.value();
932
+ STD_TORCH_CHECK(
933
+ counts.is_cuda() &&
934
+ counts.scalar_type() == torch::headeronly::ScalarType::Int &&
935
+ counts.dim() == 1 && counts.size(0) == launch_rows,
936
+ "valid_counts must be int32 on CUDA with one entry per row");
937
+ valid_counts_ptr = counts.mutable_data_ptr<int32_t>();
938
+ }
939
+
940
+ int32_t* swap_host_physical_rows_ptr = nullptr;
941
+ int32_t* swap_device_physical_rows_ptr = nullptr;
942
+ int32_t* swap_counts_ptr = nullptr;
943
+ const bool has_compact_swaps = swap_host_physical_rows.has_value();
944
+ STD_TORCH_CHECK(has_compact_swaps == swap_device_physical_rows.has_value() &&
945
+ has_compact_swaps == swap_counts.has_value(),
946
+ "compact swap tensors must be provided together");
947
+ if (has_compact_swaps) {
948
+ auto const& globals = swap_host_physical_rows.value();
949
+ auto const& hots = swap_device_physical_rows.value();
950
+ auto const& counts = swap_counts.value();
951
+ STD_TORCH_CHECK(
952
+ globals.is_cuda() && hots.is_cuda() && counts.is_cuda() &&
953
+ globals.scalar_type() == torch::headeronly::ScalarType::Int &&
954
+ hots.scalar_type() == torch::headeronly::ScalarType::Int &&
955
+ counts.scalar_type() == torch::headeronly::ScalarType::Int &&
956
+ globals.dim() == 2 && hots.dim() == 2 &&
957
+ globals.size(0) == launch_rows &&
958
+ globals.size(1) == global_indices.size(1) &&
959
+ hots.size(0) == launch_rows &&
960
+ hots.size(1) == global_indices.size(1) &&
961
+ counts.numel() >= launch_rows && globals.is_contiguous() &&
962
+ hots.is_contiguous() && counts.is_contiguous(),
963
+ "compact swaps must be contiguous int32 matching indices");
964
+ swap_host_physical_rows_ptr = globals.mutable_data_ptr<int32_t>();
965
+ swap_device_physical_rows_ptr = hots.mutable_data_ptr<int32_t>();
966
+ swap_counts_ptr = counts.mutable_data_ptr<int32_t>();
967
+ }
968
+
969
+ const int32_t* request_state_ptr = nullptr;
970
+ int32_t request_state_count = 0;
971
+ if (request_state_indices.has_value()) {
972
+ auto const& state_indices = request_state_indices.value();
973
+ STD_TORCH_CHECK(
974
+ state_indices.is_cuda() &&
975
+ state_indices.scalar_type() == torch::headeronly::ScalarType::Int &&
976
+ state_indices.dim() == 1 && state_indices.is_contiguous() &&
977
+ state_indices.numel() <= INT32_MAX,
978
+ "request_state_indices must be contiguous 1D int32 on CUDA");
979
+ request_state_ptr = state_indices.const_data_ptr<int32_t>();
980
+ request_state_count = static_cast<int32_t>(state_indices.numel());
981
+ } else {
982
+ STD_TORCH_CHECK(device_global_indices.size(0) >= launch_rows,
983
+ "device_global_indices has too few rows");
984
+ }
985
+
986
+ // Optional output: 1 at columns requiring a host-to-device swap, 0 elsewhere.
987
+ int32_t* miss_mask_ptr = nullptr;
988
+ if (miss_mask.has_value()) {
989
+ auto const& mm = miss_mask.value();
990
+ STD_TORCH_CHECK(
991
+ mm.is_cuda() &&
992
+ mm.scalar_type() == torch::headeronly::ScalarType::Int &&
993
+ mm.dim() == 2 && mm.is_contiguous() && mm.size(0) == launch_rows &&
994
+ mm.size(1) == global_indices.size(1),
995
+ "miss_mask must be a contiguous int32 CUDA tensor matching "
996
+ "global_indices");
997
+ miss_mask_ptr = mm.mutable_data_ptr<int32_t>();
998
+ }
999
+
1000
+ int32_t* attention_indices_ptr = nullptr;
1001
+ if (attention_indices.has_value()) {
1002
+ auto const& indices = attention_indices.value();
1003
+ STD_TORCH_CHECK(
1004
+ indices.is_cuda() &&
1005
+ indices.scalar_type() == torch::headeronly::ScalarType::Int &&
1006
+ indices.dim() == 2 && indices.size(0) == launch_rows &&
1007
+ indices.stride(1) == 1 && indices.size(1) == global_indices.size(1),
1008
+ "attention_indices must be row-major int32 matching global_indices");
1009
+ STD_TORCH_CHECK(attention_block_stride >= hot_block_size,
1010
+ "attention block stride must cover one hot block");
1011
+ attention_indices_ptr = indices.mutable_data_ptr<int32_t>();
1012
+ }
1013
+
1014
+ if (launch_rows == 0 || top_k == 0) {
1015
+ return;
1016
+ }
1017
+
1018
+ constexpr int kBlockSize = 1024;
1019
+ const int hash_size = 2 * top_k;
1020
+ const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize;
1021
+ const size_t smem_bytes =
1022
+ sizeof(int32_t) * (top_k + 2 * (num_buffer_chunks + 1) + hash_size + 3) +
1023
+ sizeof(int16_t) * (hot_size + hash_size);
1024
+
1025
+ const torch::stable::accelerator::DeviceGuard device_guard(
1026
+ hot_cache.get_device_index());
1027
+ const cudaStream_t stream = get_current_cuda_stream();
1028
+ const int64_t attention_row_stride =
1029
+ attention_indices.has_value() ? attention_indices.value().stride(0) : 0;
1030
+ const int64_t valid_count_stride =
1031
+ valid_counts.has_value() ? valid_counts.value().stride(0) : 0;
1032
+ auto kernel = hisparse_resolve_residency_kernel;
1033
+ if (smem_bytes > 48 * 1024) {
1034
+ const cudaError_t attribute_error = cudaFuncSetAttribute(
1035
+ kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
1036
+ STD_TORCH_CHECK(attribute_error == cudaSuccess,
1037
+ "failed to configure HiSparse swap-in shared memory: ",
1038
+ cudaGetErrorString(attribute_error));
1039
+ }
1040
+ kernel<<<launch_rows, kBlockSize, smem_bytes, stream>>>(
1041
+ hot_block_table.const_data_ptr<int32_t>(),
1042
+ global_indices.const_data_ptr<int32_t>(), request_ids_ptr,
1043
+ source_block_table_ptr, resident_block_table_ptr,
1044
+ resolved_global_indices_ptr, valid_counts_ptr,
1045
+ swap_host_physical_rows_ptr, swap_device_physical_rows_ptr,
1046
+ swap_counts_ptr, hot_indices.mutable_data_ptr<int32_t>(),
1047
+ attention_indices_ptr, miss_mask_ptr,
1048
+ device_global_indices.mutable_data_ptr<int32_t>(),
1049
+ lru_slots.mutable_data_ptr<int16_t>(), request_state_ptr,
1050
+ request_state_count, host_rows, hot_block_table.stride(0), hot_block_size,
1051
+ top_k, hot_size, hash_size, region_stride, attention_block_stride,
1052
+ source_bt_stride, source_num_reqs, source_num_blocks,
1053
+ static_cast<int32_t>(source_block_size), resident_bt_stride,
1054
+ resident_num_reqs, resident_num_blocks,
1055
+ static_cast<int32_t>(resident_block_size),
1056
+ static_cast<int32_t>(resident_null_block), global_indices.stride(0),
1057
+ attention_row_stride, valid_count_stride);
1058
+ const cudaError_t launch_error = cudaGetLastError();
1059
+ STD_TORCH_CHECK(launch_error == cudaSuccess,
1060
+ "HiSparse residency kernel launch failed: ",
1061
+ cudaGetErrorString(launch_error));
1062
+}
1063
+
1064
+void hisparse_gather_plan(
1065
+ torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache,
1066
+ torch::stable::Tensor const& global_indices,
1067
+ torch::stable::Tensor const& hot_indices,
1068
+ torch::stable::Tensor const& miss_mask,
1069
+ std::optional<torch::stable::Tensor> const& request_state_indices,
1070
+ std::optional<torch::stable::Tensor> const& attention_indices,
1071
+ int64_t attention_block_stride) {
1072
+ STD_TORCH_CHECK(
1073
+ host_cache.device().is_cpu() && is_pinned_cpu_tensor(host_cache),
1074
+ "host_cache must be pinned CPU memory");
1075
+ STD_TORCH_CHECK(hot_cache.is_cuda(), "hot_cache must be on CUDA");
1076
+ STD_TORCH_CHECK(
1077
+ global_indices.is_cuda() && hot_indices.is_cuda() && miss_mask.is_cuda(),
1078
+ "plan tensors must be on CUDA");
1079
+ STD_TORCH_CHECK(
1080
+ global_indices.scalar_type() == torch::headeronly::ScalarType::Int &&
1081
+ hot_indices.scalar_type() == torch::headeronly::ScalarType::Int &&
1082
+ miss_mask.scalar_type() == torch::headeronly::ScalarType::Int,
1083
+ "plan tensors must be int32");
1084
+ STD_TORCH_CHECK(global_indices.dim() == 2 && global_indices.is_contiguous(),
1085
+ "global_indices must be contiguous 2D");
1086
+ STD_TORCH_CHECK(
1087
+ hot_indices.size(0) == global_indices.size(0) &&
1088
+ hot_indices.size(1) == global_indices.size(1) &&
1089
+ miss_mask.size(0) == global_indices.size(0) &&
1090
+ miss_mask.size(1) == global_indices.size(1) &&
1091
+ hot_indices.is_contiguous() && miss_mask.is_contiguous(),
1092
+ "hot_indices/miss_mask must match contiguous 2D global_indices");
1093
+
1094
+ STD_TORCH_CHECK(hot_cache.dim() == 2 || hot_cache.dim() == 3,
1095
+ "hot_cache must be a 2D staging buffer or paged 3D cache");
1096
+ const int64_t row_bytes = hot_cache.size(-1) * hot_cache.element_size();
1097
+ STD_TORCH_CHECK(row_bytes % 16 == 0, "KV rows must be 16-byte aligned");
1098
+ STD_TORCH_CHECK(
1099
+ hot_cache.stride(hot_cache.dim() - 2) * hot_cache.element_size() ==
1100
+ row_bytes,
1101
+ "hot-cache rows must be contiguous");
1102
+ const int64_t host_rows = check_2d_rows(host_cache, "host_cache", row_bytes);
1103
+
1104
+ const auto num_rows = static_cast<int32_t>(global_indices.size(0));
1105
+ const auto top_k = static_cast<int32_t>(global_indices.size(1));
1106
+
1107
+ const int32_t* request_state_ptr = nullptr;
1108
+ if (request_state_indices.has_value()) {
1109
+ auto const& state_indices = request_state_indices.value();
1110
+ STD_TORCH_CHECK(
1111
+ state_indices.is_cuda() &&
1112
+ state_indices.scalar_type() == torch::headeronly::ScalarType::Int &&
1113
+ state_indices.numel() >= num_rows,
1114
+ "request_state_indices must be int32 on CUDA with one entry per row");
1115
+ request_state_ptr = state_indices.const_data_ptr<int32_t>();
1116
+ }
1117
+
1118
+ int32_t* attention_indices_ptr = nullptr;
1119
+ if (attention_indices.has_value()) {
1120
+ auto const& indices = attention_indices.value();
1121
+ STD_TORCH_CHECK(
1122
+ indices.is_cuda() &&
1123
+ indices.scalar_type() == torch::headeronly::ScalarType::Int &&
1124
+ indices.dim() == 2 && indices.is_contiguous() &&
1125
+ indices.size(0) == global_indices.size(0) &&
1126
+ indices.size(1) == global_indices.size(1),
1127
+ "attention_indices must be a contiguous int32 CUDA tensor matching "
1128
+ "global_indices");
1129
+ attention_indices_ptr = indices.mutable_data_ptr<int32_t>();
1130
+ }
1131
+
1132
+ if (num_rows == 0 || top_k == 0) {
1133
+ return;
1134
+ }
1135
+
1136
+ // Match the swap-in kernel's block size: the gather serves 3 of every 4
1137
+ // layers' misses (index-sharing replay), so per-row copy parallelism is
1138
+ // the throughput limiter on cold rows.
1139
+ constexpr int kBlockSize = 1024;
1140
+ constexpr int kNumWarps = kBlockSize / kWarpSize;
1141
+ // Interleave columns over enough blocks per row to cover the device even
1142
+ // for few-row launches (local-prefill staging's single-row layout, small
1143
+ // decode batches), keeping >= 1 column per warp.
1144
+ constexpr int kTargetBlocks = 256;
1145
+ const int max_chunks = std::max(1, (top_k + kNumWarps - 1) / kNumWarps);
1146
+ const int num_chunks =
1147
+ std::min(max_chunks, std::max(1, kTargetBlocks / num_rows));
1148
+ const dim3 grid(num_rows, num_chunks);
1149
+ const int32_t hot_block_size =
1150
+ hot_cache.dim() == 3 ? static_cast<int32_t>(hot_cache.size(1)) : 1;
1151
+ if (attention_indices.has_value()) {
1152
+ STD_TORCH_CHECK(attention_block_stride >= hot_block_size,
1153
+ "attention block stride must cover one hot block");
1154
+ }
1155
+ const int64_t hot_rows = hot_cache.size(0) * hot_block_size;
1156
+ const int64_t hot_block_stride =
1157
+ hot_cache.stride(0) * hot_cache.element_size();
1158
+ const torch::stable::accelerator::DeviceGuard device_guard(
1159
+ hot_cache.get_device_index());
1160
+ const cudaStream_t stream = get_current_cuda_stream();
1161
+ hisparse_gather_plan_kernel<<<grid, kBlockSize, 0, stream>>>(
1162
+ static_cast<const char*>(host_cache.const_data_ptr()),
1163
+ static_cast<char*>(hot_cache.mutable_data_ptr()),
1164
+ global_indices.const_data_ptr<int32_t>(),
1165
+ hot_indices.const_data_ptr<int32_t>(),
1166
+ miss_mask.const_data_ptr<int32_t>(), attention_indices_ptr,
1167
+ request_state_ptr, host_rows, hot_rows, row_bytes, hot_block_stride,
1168
+ hot_block_size, top_k, attention_block_stride);
1169
+}
1170
+
1171
+void hisparse_gather_compact(torch::stable::Tensor const& host_cache,
1172
+ torch::stable::Tensor& hot_cache,
1173
+ torch::stable::Tensor const& miss_global_indices,
1174
+ torch::stable::Tensor const& miss_hot_indices,
1175
+ torch::stable::Tensor const& miss_counts) {
1176
+ STD_TORCH_CHECK(
1177
+ host_cache.device().is_cpu() && is_pinned_cpu_tensor(host_cache),
1178
+ "host_cache must be pinned CPU memory");
1179
+ STD_TORCH_CHECK(hot_cache.is_cuda(), "hot_cache must be on CUDA");
1180
+ STD_TORCH_CHECK(
1181
+ miss_global_indices.is_cuda() && miss_hot_indices.is_cuda() &&
1182
+ miss_counts.is_cuda() &&
1183
+ miss_global_indices.scalar_type() ==
1184
+ torch::headeronly::ScalarType::Int &&
1185
+ miss_hot_indices.scalar_type() ==
1186
+ torch::headeronly::ScalarType::Int &&
1187
+ miss_counts.scalar_type() == torch::headeronly::ScalarType::Int &&
1188
+ miss_global_indices.dim() == 2 && miss_hot_indices.dim() == 2 &&
1189
+ miss_global_indices.size(0) == miss_hot_indices.size(0) &&
1190
+ miss_global_indices.size(1) == miss_hot_indices.size(1) &&
1191
+ miss_counts.numel() >= miss_global_indices.size(0) &&
1192
+ miss_global_indices.is_contiguous() &&
1193
+ miss_hot_indices.is_contiguous() && miss_counts.is_contiguous(),
1194
+ "compact miss plan must be matching contiguous int32 CUDA tensors");
1195
+ STD_TORCH_CHECK(hot_cache.dim() == 3,
1196
+ "hot_cache must be [num_blocks, block_size, row_width]");
1197
+ const int64_t row_bytes = hot_cache.size(-1) * hot_cache.element_size();
1198
+ STD_TORCH_CHECK(row_bytes % 16 == 0, "KV rows must be 16-byte aligned");
1199
+ STD_TORCH_CHECK(hot_cache.stride(1) * hot_cache.element_size() == row_bytes,
1200
+ "hot-cache rows must be contiguous");
1201
+ const int64_t host_rows = check_2d_rows(host_cache, "host_cache", row_bytes);
1202
+ const auto num_rows = static_cast<int32_t>(miss_global_indices.size(0));
1203
+ const auto top_k = static_cast<int32_t>(miss_global_indices.size(1));
1204
+ if (num_rows == 0 || top_k == 0) {
1205
+ return;
1206
+ }
1207
+
1208
+ constexpr int kBlockSize = 512;
1209
+ constexpr int kTargetBlocks = 64;
1210
+ const int num_chunks = std::min(8, std::max(1, kTargetBlocks / num_rows));
1211
+ const dim3 grid(num_rows, num_chunks);
1212
+ const int32_t hot_block_size = static_cast<int32_t>(hot_cache.size(1));
1213
+ const int64_t hot_rows = hot_cache.size(0) * hot_block_size;
1214
+ const int64_t hot_block_stride =
1215
+ hot_cache.stride(0) * hot_cache.element_size();
1216
+ const torch::stable::accelerator::DeviceGuard device_guard(
1217
+ hot_cache.get_device_index());
1218
+ const cudaStream_t stream = get_current_cuda_stream();
1219
+ hisparse_gather_compact_kernel<<<grid, kBlockSize, 0, stream>>>(
1220
+ static_cast<const char*>(host_cache.const_data_ptr()),
1221
+ static_cast<char*>(hot_cache.mutable_data_ptr()),
1222
+ miss_global_indices.const_data_ptr<int32_t>(),
1223
+ miss_hot_indices.const_data_ptr<int32_t>(),
1224
+ miss_counts.const_data_ptr<int32_t>(), host_rows, hot_rows, row_bytes,
1225
+ hot_block_stride, hot_block_size, top_k);
1226
+}
csrc/libtorch_stable/ops.h
+52
-1
@@ -661,6 +661,54 @@ void concat_and_cache_mla_grouped(
661
std::optional<torch::stable::Tensor> kv_scales,
662
const std::string& kv_cache_dtype);
663
664
+#ifndef USE_ROCM
665
+// HiSparse kernels use raw PTX in the row copy; CUDA-only.
666
+void hisparse_resolve_residency(
667
+ torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache,
668
+ torch::stable::Tensor const& hot_block_table,
669
+ torch::stable::Tensor const& global_indices,
670
+ torch::stable::Tensor& hot_indices,
671
+ torch::stable::Tensor& device_global_indices,
672
+ torch::stable::Tensor& lru_slots,
673
+ std::optional<torch::stable::Tensor> const& request_state_indices,
674
+ int64_t region_stride,
675
+ std::optional<torch::stable::Tensor> const& miss_mask,
676
+ std::optional<torch::stable::Tensor> const& attention_indices,
677
+ int64_t attention_block_stride,
678
+ std::optional<torch::stable::Tensor> const& request_ids,
679
+ std::optional<torch::stable::Tensor> const& source_block_table,
680
+ int64_t source_block_size,
681
+ std::optional<torch::stable::Tensor> const& resolved_global_indices,
682
+ std::optional<torch::stable::Tensor> const& valid_counts,
683
+ std::optional<torch::stable::Tensor> const& swap_host_physical_rows,
684
+ std::optional<torch::stable::Tensor> const& swap_device_physical_rows,
685
+ std::optional<torch::stable::Tensor> const& swap_counts,
686
+ std::optional<torch::stable::Tensor> const& resident_block_table,
687
+ int64_t resident_block_size, int64_t resident_null_block);
688
+
689
+void hisparse_invalidate_written_slots(
690
+ torch::stable::Tensor& device_global_indices,
691
+ torch::stable::Tensor const& request_state_indices,
692
+ torch::stable::Tensor const& req_id_per_token,
693
+ torch::stable::Tensor const& written_slots);
694
+
695
+void hisparse_gather_plan(
696
+ torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache,
697
+ torch::stable::Tensor const& global_indices,
698
+ torch::stable::Tensor const& hot_indices,
699
+ torch::stable::Tensor const& miss_mask,
700
+ std::optional<torch::stable::Tensor> const& request_state_indices,
701
+ std::optional<torch::stable::Tensor> const& attention_indices,
702
+ int64_t attention_block_stride);
703
+
704
+void hisparse_gather_compact(torch::stable::Tensor const& host_cache,
705
+ torch::stable::Tensor& hot_cache,
706
+ torch::stable::Tensor const& miss_global_indices,
707
+ torch::stable::Tensor const& miss_hot_indices,
708
+ torch::stable::Tensor const& miss_counts);
709
+
710
+#endif // !USE_ROCM
711
+
712
// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
713
void concat_and_cache_mla_rope_fused(
714
torch::stable::Tensor& positions, torch::stable::Tensor& q_pe,
@@ -704,7 +752,10 @@ void cp_gather_and_upconvert_fp8_kv_cache(
752
torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
753
torch::stable::Tensor const& workspace_starts, // [BATCH]
754
int64_t batch_size,
707
- std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
755
+ std::optional<torch::stable::Tensor> seq_starts = std::nullopt,
756
+ std::optional<torch::stable::Tensor> host_cache = std::nullopt,
757
+ std::optional<torch::stable::Tensor> host_row_ids = std::nullopt,
758
+ std::optional<torch::stable::Tensor> device_row_ids = std::nullopt);
759
760
// Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace
761
void cp_gather_and_upconvert_nvfp4_kv_cache(
csrc/libtorch_stable/torch_bindings.cpp
+62
-1
@@ -987,6 +987,57 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
987
" Tensor? kv_scales=None,"
988
" str kv_cache_dtype='auto') -> ()");
989
990
+#ifndef USE_ROCM
991
+ ops.def(
992
+ "hisparse_resolve_residency(Tensor host_cache,"
993
+ " Tensor! hot_cache,"
994
+ " Tensor hot_block_table,"
995
+ " Tensor global_indices,"
996
+ " Tensor! hot_indices,"
997
+ " Tensor! device_global_indices,"
998
+ " Tensor! lru_slots,"
999
+ " Tensor? request_state_indices,"
1000
+ " int region_stride,"
1001
+ " Tensor(a!)? miss_mask=None,"
1002
+ " Tensor(b!)? attention_indices=None,"
1003
+ " int attention_block_stride=0,"
1004
+ " Tensor? request_ids=None,"
1005
+ " Tensor? source_block_table=None,"
1006
+ " int source_block_size=0,"
1007
+ " Tensor(c!)? resolved_global_indices=None,"
1008
+ " Tensor(d!)? valid_counts=None,"
1009
+ " Tensor(e!)? swap_host_physical_rows=None,"
1010
+ " Tensor(f!)? swap_device_physical_rows=None,"
1011
+ " Tensor(g!)? swap_counts=None,"
1012
+ " Tensor? resident_block_table=None,"
1013
+ " int resident_block_size=0,"
1014
+ " int resident_null_block=0) -> ()");
1015
+
1016
+ ops.def(
1017
+ "hisparse_invalidate_written_slots(Tensor! device_global_indices,"
1018
+ " Tensor request_state_indices,"
1019
+ " Tensor req_id_per_token,"
1020
+ " Tensor written_slots) -> ()");
1021
+
1022
+ ops.def(
1023
+ "hisparse_gather_plan(Tensor host_cache,"
1024
+ " Tensor! hot_cache,"
1025
+ " Tensor global_indices,"
1026
+ " Tensor hot_indices,"
1027
+ " Tensor miss_mask,"
1028
+ " Tensor? request_state_indices,"
1029
+ " Tensor(a!)? attention_indices=None,"
1030
+ " int attention_block_stride=0) -> ()");
1031
+
1032
+ ops.def(
1033
+ "hisparse_gather_compact(Tensor host_cache,"
1034
+ " Tensor! hot_cache,"
1035
+ " Tensor miss_global_indices,"
1036
+ " Tensor miss_hot_indices,"
1037
+ " Tensor miss_counts) -> ()");
1038
+
1039
+#endif // !USE_ROCM
1040
+
1041
// Rotate Q and K, then write to kv cache for MLA
1042
ops.def(
1043
"concat_and_cache_mla_rope_fused("
@@ -1023,7 +1074,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
1074
ops.def(
1075
"cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
1076
"Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? "
1026
- "seq_starts) -> ()");
1077
+ "seq_starts, Tensor? host_cache=None, Tensor? host_row_ids=None, Tensor? "
1078
+ "device_row_ids=None) -> ()");
1079
1080
ops.def(
1081
"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) {
1143
ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla));
1144
ops.impl("concat_and_cache_mla_grouped",
1145
TORCH_BOX(&concat_and_cache_mla_grouped));
1146
+
1147
+#ifndef USE_ROCM
1148
+ ops.impl("hisparse_resolve_residency",
1149
+ TORCH_BOX(&hisparse_resolve_residency));
1150
+ ops.impl("hisparse_invalidate_written_slots",
1151
+ TORCH_BOX(&hisparse_invalidate_written_slots));
1152
+ ops.impl("hisparse_gather_plan", TORCH_BOX(&hisparse_gather_plan));
1153
+ ops.impl("hisparse_gather_compact", TORCH_BOX(&hisparse_gather_compact));
1154
+#endif // !USE_ROCM
1155
ops.impl("concat_and_cache_mla_rope_fused",
1156
TORCH_BOX(&concat_and_cache_mla_rope_fused));
1157
ops.impl("convert_fp8", TORCH_BOX(&convert_fp8));
docs/design/hisparse.md
new
+244
@@ -0,0 +1,244 @@
1
+# HiSparse local KV offload architecture
2
+
3
+Status: experimental
4
+
5
+## The short version
6
+
7
+There are three different jobs:
8
+
9
+1. The normal KV cache system manages GPU block pools and tables.
10
+2. `HiSparseCoordinator` manages logical host blocks, source-prefix identity, and
11
+ host/GPU residency transitions.
12
+3. `HiSparseConnector` carries residency work between scheduler and worker;
13
+ `HiSparseWorker` coordinates transfers, while per-cache
14
+ `HiSparseRuntime` objects own host/hot views and GPU replacement state.
15
+
16
+Neither worker-side object allocates or frees logical blocks. HMA provides the
17
+GPU allocation shared by the resident and hot groups; it does not manage CPU
18
+memory or KV identity.
19
+
20
+```text
21
+request ────► HiSparseCoordinator ── source blocks + residency policy
22
+ │
23
+ ├──► KV cache manager ── resident/hot GPU leases (HMA)
24
+ │
25
+ └──► HiSparseConnector ── host bytes + copies + GPU LRU
26
+```
27
+
28
+This is a local KV connector. When P/D or another offload connector is also
29
+configured, `MultiConnector` composes it with `HiSparseConnector`.
30
+
31
+`host_pool_gib` is configured on `HiSparseConnector` and is the usable
32
+host-cache capacity per data-parallel replica, not a node-wide memory budget.
33
+Tensor-parallel ranks hold replicated views of that logical cache. Those views
34
+may use private per-rank backing or one shared
35
+physical allocation without changing the configured capacity. Physical host
36
+memory consumption is therefore topology- and implementation-dependent. The
37
+realized capacity may be slightly smaller because the budget is rounded down
38
+to complete host blocks.
39
+
40
+## Ownership
41
+
42
+| Thing | Owner | What “owner” means |
43
+| --- | --- | --- |
44
+| HiSparse source and prefix identity | `HiSparseCoordinator` | maps tokens to logical host blocks |
45
+| Resident GPU block leases | normal KV cache manager | allocates and frees HMA blocks |
46
+| Resident block tables | normal KV cache manager | tells attention where resident pages are |
47
+| Residency transitions | `HiSparseCoordinator` | plans spill-before-free transactions |
48
+| Logical host block allocation | `HiSparseCoordinator` | owns the separate CPU block pool and its lifecycle |
49
+| Pinned host-pool lifecycle | `HiSparseWorker` | worker-wide backing and teardown |
50
+| Per-cache host view and hot contents | `HiSparseRuntime` | binds host/hot storage and fills cache-manager-provided hot leases |
51
+| Hot row map and LRU | `HiSparseRuntime` | resolves hits and chooses victims on GPU |
52
+| Resident-cache route | `HiSparseCacheHandle` | exposes resident or host/hot resolution to attention |
53
+| Sparse attention | attention backend | consumes a device cache and physical row IDs |
54
+| HMA | allocator | provides GPU capacity; owns no KV meaning |
55
+
56
+The key distinction is logical allocation versus contents. `HiSparseCoordinator`
57
+owns host block IDs and their request/prefix associations. `HiSparseWorker` and
58
+its per-cache runtimes own the corresponding bytes. The normal cache manager
59
+sees only device pools.
60
+The source group has `block_pool_id=None`; device-pool consumers must narrow it
61
+before indexing, so host ownership cannot masquerade as a numeric GPU pool.
62
+
63
+## Code boundary
64
+
65
+```text
66
+scheduler process worker process
67
+
68
+HiSparseConnector HiSparseConnector
69
+ └─ HiSparseCoordinator └─ HiSparseWorker
70
+ │ │
71
+ │ connector metadata ├─ host bytes
72
+ │ - page transfers ├─ copy scheduling
73
+ │ - block-table replacements └─ per-layer hot state
74
+ └───────────────────────────────────────────────►│
75
+ ◄──────── connector worker metadata ─────────────┘
76
+ enqueued and completed transfer IDs
77
+```
78
+
79
+The command travels in `kv_connector_metadata`; transfer updates return in
80
+`KVConnectorOutput.kv_connector_worker_meta`. The model runner does not
81
+interpret page transfers. Enqueue acknowledgements let the scheduler release
82
+source leases in stream order; completion acknowledgements publish the copied
83
+host pages.
84
+
85
+## Resident device pages
86
+
87
+Resident pages are intentionally outside `HiSparseRuntime`.
88
+
89
+KV-cache initialization binds cache-manager allocations to the attention-facing
90
+`HiSparseCacheHandle` before constructing `HiSparseWorker`. That same
91
+handle's runtime retains the resident source index needed by a transfer plan.
92
+There is no second resident object or registration wrapper.
93
+
94
+```text
95
+KV cache setup
96
+ │
97
+ ├─ bind resident allocation ──► HiSparseCacheHandle
98
+ │ cache + block table + slot mapping
99
+ │
100
+ ├─ bind host/hot allocation ──► HiSparseRuntime
101
+ │ host + hot + GPU LRU
102
+ │
103
+ └─ register cache handles ────► HiSparseWorker
104
+ step-level transfers
105
+
106
+HiSparseWorker registers the same HiSparseCacheHandle objects directly
107
+```
108
+
109
+Attention construction links each layer to the most recent layer that actually
110
+owns an indexer. This releases a follower's duplicate LRU tensors before GPU
111
+memory profiling. Cache binding only attaches storage; it does not infer
112
+semantic groups from the physical packed-tensor order. The construction cursor
113
+is discarded with the worker's pinned state.
114
+
115
+Every HiSparse decode batch uses the same fused resolver. It checks resident
116
+pages first, then hot rows, then pinned host memory. A resident hit exits inside
117
+the kernel before hot-LRU lookup or host copying; there is no framework-level
118
+residency route or separate CUDA graph. No CPU decision is added to the decode
119
+path. The resolver consumes the existing graph-stable request mapping from
120
+attention metadata; neither the worker nor individual cache handles keep a
121
+duplicate mapping.
122
+
123
+Speculative decoding resolves and consumes each verification step in order.
124
+Each step receives distinct replayable plan rows while sharing the request's
125
+hot-cache state, so a later step cannot reuse a hot row before an earlier step
126
+has consumed it.
127
+
128
+## P/D import target
129
+
130
+The decoder chooses the landing target once per request from the normal cache
131
+admission calculation. If the complete imported prefix fits the device pools,
132
+NIXL transfers it directly into resident GPU pages. Otherwise, if the fixed
133
+host-backed GPU footprint and host source blocks fit, the request imports into
134
+the host tier. There is no context-length threshold or other heuristic, and a
135
+request waiting for capacity retains its choice across admission retries.
136
+
137
+A host import reads through a bounded decoder-GPU staging pool before copying
138
+into registered host memory. Pages needed immediately are mirrored into their
139
+resident destinations during that copy. Both landing targets then use the same
140
+fused decode resolver described above.
141
+
142
+## Indexer KV offloading
143
+
144
+HiSparse does not keep a private CPU copy of indexer KV. The indexer remains a
145
+normal prefix-cacheable GPU cache group. If `OffloadingConnector` is configured
146
+with HiSparse, it stores and restores that group through the generic KV
147
+offloading path; HiSparse continues to own only the sparse MLA host tier.
148
+
149
+The two prefix sources can have different hit lengths. When the HiSparse host
150
+prefix extends beyond the GPU-resident indexer prefix, the scheduler asks
151
+`OffloadingConnector` to restore only the missing indexer suffix, capped at the
152
+host prefix boundary. If that suffix is unavailable, all groups fall back to
153
+the shorter prefix they share. NIXL P/D transfers continue to place indexer KV
154
+directly in its GPU group.
155
+
156
+## Spill transaction
157
+
158
+A resident block cannot be reused until its contents have been handed to the
159
+worker.
160
+
161
+```text
162
+HiSparseCoordinator HiSparseWorker
163
+ │ │
164
+ │ pin source and destination leases │
165
+ │── SparseKVPageTransfer ─────────────►│
166
+ │ │ enqueue GPU-to-host copy
167
+ │◄── enqueued transfer ID ────────────│
168
+ │ replace resident table entry │
169
+ │ release resident lease to HMA │
170
+ │ │ copy reaches its event
171
+ │◄── completed transfer ID ───────────│
172
+ │ mark host page valid │
173
+ │ release destination host lease │
174
+```
175
+
176
+“Enqueued” means the copy has entered the worker stream. Stream ordering makes
177
+it safe to reuse the resident GPU block for later work, but the host page is
178
+not yet published. “Completed” means the worker has observed the copy's event;
179
+only then does the coordinator publish the host page for prefix reuse and
180
+release its destination lease. A host-write event separately protects direct
181
+CPU readers from writes already queued on the accelerator.
182
+
183
+The worker transfer contains only its transfer ID and physical copy
184
+coordinates. Request identity and logical page state remain in the scheduler.
185
+
186
+## Hot lookup and LRU
187
+
188
+The NVIDIA CUDA path keeps replacement entirely on the accelerator:
189
+
190
+```text
191
+top-K logical positions
192
+ │
193
+ ▼
194
+resident page? ── yes ──► resident physical row
195
+ │ no
196
+ ▼
197
+hot row? ──────── yes ──► existing hot physical row + update GPU LRU
198
+ │ no
199
+ ▼
200
+choose GPU LRU victim ──► copy pinned host row ──► hot physical row
201
+```
202
+
203
+ROCm is not currently supported because the fused HiSparse cache operations are
204
+implemented only by CUDA kernels. A future platform-specific worker may provide
205
+the same command, output, and cache-resolution boundaries.
206
+
207
+## Main classes
208
+
209
+| Class | Inherits / implements | Responsibility |
210
+| --- | --- | --- |
211
+| `HiSparseCoordinator` | plain scheduler component | host allocation, source prefixes, resident leases, and spill state machine |
212
+| `HiSparseConnector` | `KVConnectorBase_V1`, `SupportsHMA` | scheduler/worker metadata and lifecycle boundary |
213
+| `HiSparseResidentManager` | `SingleTypeKVCacheManager` | normal block-pool bookkeeping with host-backed holes |
214
+| `PagedCacheView` | immutable data object | shared resident/hot HMA tensor binding |
215
+| `HiSparseWorker` | connector-owned worker component | worker-wide transfer scheduling and host-pool lifecycle |
216
+| `HiSparseRuntime` | plain worker-owned component | per-cache host/hot tensors, GPU LRU, and fused resolution |
217
+| `HiSparseCacheHandle` | plain attention component | resident view and fused cache resolution |
218
+| `SparseKVOffloadCommand` | dataclass | opaque scheduler-to-worker work |
219
+
220
+## Performance invariants
221
+
222
+- Resident hits bypass hot-LRU lookup and host copies inside the fused resolver.
223
+- Hot lookup, victim selection, and LRU updates stay on the GPU.
224
+- A hot miss still copies directly from registered pinned host memory.
225
+- Top-K resolution stays inside the attention invocation and remains graph
226
+ capturable.
227
+- Compatible layers still share one miss plan.
228
+- Index-sharing followers release their private LRU state before memory sizing.
229
+- Indexer KV is untouched by HiSparse unless a generic KV offloader is configured.
230
+- Resident and hot leases can still share one packed HMA allocation.
231
+- No device scalar readback or CPU/device metadata round trip is added.
232
+- The abstraction wraps the fused kernel; it does not add another kernel
233
+ launch.
234
+- When HiSparse is disabled, the scheduler does not construct an offload
235
+ command or empty update table.
236
+
237
+## What remains platform-specific
238
+
239
+The command/result and attention-layer boundaries can be shared. The host
240
+allocator, copy implementation, hot layout, and replacement policy should stay
241
+platform-specific. NVIDIA uses the current accelerator LRU and fused host/hot
242
+kernel. ROCm is not currently supported; AMD or other accelerator backends can
243
+implement their own worker without forcing NVIDIA's policy into the shared
244
+boundary.
tests/evals/gsm8k/gsm8k_eval.py
+7
@@ -364,6 +364,12 @@ def main() -> None:
364
type=int,
365
help="Maximum number of concurrent requests",
366
)
367
+ parser.add_argument(
368
+ "--request-timeout-seconds",
369
+ type=float,
370
+ default=600,
371
+ help="Timeout for each request, including time waiting for a connection",
372
+ )
373
parser.add_argument("--save-results", type=str, help="Save results to JSON file")
374
375
args = parser.parse_args()
@@ -377,6 +383,7 @@ def main() -> None:
383
temperature=args.temperature,
384
seed=args.seed,
385
max_concurrency=args.max_concurrency,
386
+ request_timeout_seconds=args.request_timeout_seconds,
387
)
388
389
# Print results to terminal
tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py
+8
-4
@@ -217,7 +217,8 @@ def test_flashinfer_mla_dense_decode_unified_slot_view():
217
assert (out_ref - out_view).abs().max().item() == 0.0
218
219
220
-def test_flashmla_fp8_sparse_decode_unified_slot_view():
220
+@pytest.mark.parametrize("next_n", [1, 2, 4])
221
+def test_flashmla_fp8_sparse_decode_unified_slot_view(next_n: int):
222
"""FlashMLA fp8 sparse decode (DeepSeek V3.2/V4 DSA path) must read a
223
unified-slot block-major view bit-identically to a contiguous fp8_ds_mla
224
cache, with finite nonzero output."""
@@ -240,7 +241,10 @@ def test_flashmla_fp8_sparse_decode_unified_slot_view():
241
n_layers = 3
242
layer = 1
243
243
- q = torch.randn(batch, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1
244
+ q = (
245
+ torch.randn(batch, next_n, h_q, head_dim, device=dev, dtype=torch.bfloat16)
246
+ * 0.1
247
+ )
248
249
# Structurally valid fp8 ds_mla payload: 512B fp8 + 16B f32 scales + 128B
250
# bf16 rope (random bytes corrupt the scale region and yield NaNs).
@@ -273,13 +277,13 @@ def test_flashmla_fp8_sparse_decode_unified_slot_view():
277
278
# Sparse indices: each batch uses its own disjoint blocks.
279
blocks_per_batch = num_blocks // batch
276
- idx = torch.full((batch, 1, topk), -1, device=dev, dtype=torch.int32)
280
+ idx = torch.full((batch, next_n, topk), -1, device=dev, dtype=torch.int32)
281
for b in range(batch):
282
slots: list[int] = []
283
for blk in range(b * blocks_per_batch, (b + 1) * blocks_per_batch):
284
slots.extend(blk * page + off for off in range(page))
285
slots_t = torch.tensor(slots[:topk], device=dev, dtype=torch.int32)
282
- idx[b, 0, : slots_t.numel()] = slots_t
286
+ idx[b, :, : slots_t.numel()] = slots_t
287
288
def run(kc):
289
meta, num_splits = fm.get_mla_metadata()
tests/kernels/test_cp_gather_fp8.py
+42
@@ -316,6 +316,48 @@ def test_cp_gather_fp8_shuffled_blocks():
316
assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:])
317
318
319
+def test_cp_gather_fp8_mixes_resident_and_pinned_host_rows():
320
+ seq_lens = [7, 5]
321
+ block_size = 4
322
+ (
323
+ cache,
324
+ block_table,
325
+ workspace_starts,
326
+ num_reqs,
327
+ total_tokens,
328
+ expected,
329
+ ) = _build_test_case(seq_lens, block_size)
330
+
331
+ device_cache = cache.clone()
332
+ host_cache = cache.cpu().pin_memory().view(-1, ENTRY_BYTES)
333
+ device_rows = torch.arange(host_cache.shape[0], dtype=torch.int32, device="cuda")
334
+ use_device = device_rows % 2 == 0
335
+ host_rows = device_rows.clone()
336
+ device_rows.masked_fill_(~use_device, -1)
337
+
338
+ host_cache[use_device.cpu()] = 0
339
+ device_cache.view(-1, ENTRY_BYTES)[~use_device] = 0
340
+ dst = torch.empty(
341
+ total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda"
342
+ )
343
+
344
+ ops.cp_gather_and_upconvert_fp8_kv_cache(
345
+ device_cache,
346
+ dst,
347
+ block_table,
348
+ workspace_starts,
349
+ num_reqs,
350
+ host_cache=host_cache,
351
+ host_row_ids=host_rows,
352
+ device_row_ids=device_rows,
353
+ )
354
+
355
+ torch.testing.assert_close(
356
+ dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2
357
+ )
358
+ assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:])
359
+
360
+
361
@pytest.mark.parametrize(
362
"gather_seq_lens,seq_starts",
363
[
tests/kernels/test_fused_deepseek_v32_norm_rope.py
+70
@@ -234,6 +234,76 @@ def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str
234
assert (topk == -1).all(), "topk buffer not cleared on indexer layer"
235
236
237
+def test_fused_norm_rope_packed_indexer_block_stride():
238
+ """Indexer writes use their HMA block stride and independent slot map."""
239
+ torch.manual_seed(7)
240
+ dev = "cuda"
241
+ num_tokens = block_size = 2
242
+ pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
243
+ q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
244
+ kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
245
+ k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
246
+ qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
247
+ kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
248
+ ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)
249
+ ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
250
+ ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
251
+ cos_sin = make_cos_sin(32, ROPE_DIM, dev)
252
+
253
+ idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4
254
+ packed_block_stride = block_size * idx_row + 64
255
+ backing = torch.zeros(2 * packed_block_stride, device=dev, dtype=torch.uint8)
256
+ idx_cache = torch.as_strided(
257
+ backing,
258
+ (2, block_size, idx_row),
259
+ (packed_block_stride, idx_row, 1),
260
+ )
261
+ mla_cache = torch.zeros(
262
+ 1,
263
+ block_size,
264
+ KV_LORA + ROPE_DIM,
265
+ device=dev,
266
+ dtype=torch.bfloat16,
267
+ )
268
+ topk = torch.zeros(num_tokens, 8, device=dev, dtype=torch.int32)
269
+
270
+ K.fused_norm_rope(
271
+ pos,
272
+ q_c,
273
+ qw,
274
+ EPS,
275
+ kv_c,
276
+ kvw,
277
+ EPS,
278
+ k_pe,
279
+ cos_sin,
280
+ ik,
281
+ ikw,
282
+ ikb,
283
+ EPS,
284
+ cos_sin,
285
+ topk,
286
+ slot_mapping=torch.arange(num_tokens, device=dev, dtype=torch.int64),
287
+ indexer_slot_mapping=torch.arange(
288
+ block_size, block_size + num_tokens, device=dev, dtype=torch.int64
289
+ ),
290
+ indexer_k_cache=idx_cache,
291
+ mla_kv_cache=mla_cache,
292
+ has_indexer=True,
293
+ )
294
+
295
+ ik_ref = rope(layer_norm(ik, ikw, ikb), pos, cos_sin, interleave=False)
296
+ q_ref, s_ref = ue8m0_quant(ik_ref)
297
+ packed = idx_cache[1].reshape(-1)
298
+ values = (
299
+ packed[: block_size * INDEX_HEAD_DIM].view(FP8).view(block_size, INDEX_HEAD_DIM)
300
+ )
301
+ scales = packed[block_size * INDEX_HEAD_DIM :].view(torch.float32)
302
+ assert_fp8(values, q_ref, "packed indexer-K fp8")
303
+ torch.testing.assert_close(scales, s_ref, rtol=0, atol=0)
304
+ assert (backing[block_size * idx_row : packed_block_stride] == 0).all()
305
+
306
+
307
def test_fused_norm_rope_materializes_qk_without_local_cache_slots():
308
"""DCP non-owner ranks still need valid query shards for query AllGather."""
309
torch.manual_seed(7)
tests/model_executor/layers/test_mla_short_prefill_indexer.py
+50
@@ -7,6 +7,7 @@ import pytest
7
import torch
8
9
import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer
10
+import vllm.models.deepseek_v32.attention as deepseek_attention
11
from vllm.config import CUDAGraphMode
12
from vllm.models.deepseek_v32 import attention as deepseek_v32_attention
13
from vllm.models.deepseek_v32.attention import DeepseekV32Attention
@@ -16,6 +17,51 @@ INDEXER_LAYER = "model.layers.0.self_attn.indexer.k_cache"
17
MLA_LAYER = "model.layers.0.self_attn.attn"
18
19
20
+def test_sparse_attention_refreshes_batch_state_inside_eager_segment(
21
+ monkeypatch: pytest.MonkeyPatch,
22
+) -> None:
23
+ metadata = SimpleNamespace(num_actual_tokens=1)
24
+
25
+ class BatchStateRefreshed(Exception):
26
+ pass
27
+
28
+ def refresh_batch_state(attn_metadata):
29
+ assert attn_metadata is metadata
30
+ raise BatchStateRefreshed
31
+
32
+ layer = SimpleNamespace(
33
+ indexer=None,
34
+ skip_topk=True,
35
+ layer_name=MLA_LAYER,
36
+ impl=SimpleNamespace(
37
+ prepare_for_batch=refresh_batch_state,
38
+ record_logical_topk_ready=lambda: None,
39
+ ),
40
+ )
41
+ monkeypatch.setattr(
42
+ deepseek_attention,
43
+ "get_attention_context",
44
+ lambda layer_name: (metadata, None, torch.empty(1), None),
45
+ )
46
+
47
+ with pytest.raises(BatchStateRefreshed):
48
+ DeepseekV32Attention._sparse_indexer_and_attn(
49
+ layer,
50
+ torch.empty(1, dtype=torch.long),
51
+ torch.empty(1, 1),
52
+ torch.empty(1, 1, 1),
53
+ torch.empty(1, 1, 1),
54
+ None,
55
+ None,
56
+ None,
57
+ None,
58
+ None,
59
+ torch.empty(1, 1, 1),
60
+ torch.empty(1, 1, 1),
61
+ torch.empty(1, 1),
62
+ )
63
+
64
+
65
def make_indexer_metadata(
66
*,
67
num_decodes: int = 0,
@@ -246,6 +292,10 @@ def test_deepseek_v32_dispatches_selected_mha(
292
_fp8_query=fp8_query,
293
_use_sparse_mha=lambda _: True,
294
rotary_emb=lambda _positions, q: (q + 1, None),
295
+ impl=SimpleNamespace(
296
+ record_logical_topk_ready=lambda: None,
297
+ prepare_for_batch=lambda _: None,
298
+ ),
299
forward_impl=record_forward_impl,
300
)
301
q_nope = torch.randn(2, 1, 2)
tests/model_executor/test_flashinfer_autotune_warmup.py
+13
-2
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from inspect import signature
5
from types import SimpleNamespace
6
from typing import Any
7
from unittest.mock import Mock, call, patch
@@ -76,8 +77,16 @@ def test_flashinfer_autotune_token_counts_are_bounded_and_deduplicated():
77
assert token_counts == (32,)
78
79
79
-def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run():
80
+@pytest.mark.parametrize("skip_attn", [False, True])
81
+def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(skip_attn):
82
+ from vllm.v1.worker.gpu.model_runner import GPUModelRunner as V2Runner
83
+ from vllm.v1.worker.gpu_model_runner import GPUModelRunner as V1Runner
84
+
85
runner = _make_runner([])
86
+ dummy_run_signature = signature((V2Runner if skip_attn else V1Runner)._dummy_run)
87
+ runner._dummy_run.side_effect = lambda **kwargs: dummy_run_signature.bind(
88
+ None, **kwargs
89
+ )
90
max_buckets = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 8192)
91
deferred_buckets = (1, 2, 4, 8, 16, 32, 64, 128)
92
@@ -93,7 +102,7 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run():
102
) as get_buckets,
103
patch("vllm.utils.flashinfer.autotune") as autotune,
104
):
96
- _run_flashinfer_autotune_dummy_runs(runner)
105
+ _run_flashinfer_autotune_dummy_runs(runner, skip_attn=skip_attn)
106
107
assert get_buckets.call_args_list == [call(8192), call(128)]
108
assert autotune.call_args_list == [
@@ -106,11 +115,13 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run():
115
skip_eplb=True,
116
is_profile=True,
117
randomize_inputs=True,
118
+ **({"skip_attn": True} if skip_attn else {}),
119
),
120
call(
121
num_tokens=128,
122
skip_eplb=True,
123
is_profile=True,
124
randomize_inputs=True,
125
+ **({"skip_attn": True} if skip_attn else {}),
126
),
127
]
tests/test_config.py
+66
-2
@@ -19,10 +19,12 @@ import vllm.config.vllm as vllm_config_module
19
import vllm.envs as envs
20
from vllm.compilation.backends import VllmBackend
21
from vllm.config import (
22
+ AttentionConfig,
23
CacheConfig,
24
CompilationConfig,
25
DeviceConfig,
26
EngramConfig,
27
+ HiSparseConfig,
28
KernelConfig,
29
KVTransferConfig,
30
ModelConfig,
@@ -279,6 +281,62 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected):
281
assert envs.VLLM_USE_V2_MODEL_RUNNER is expected
282
283
284
+def test_hisparse_requires_v2_model_runner():
285
+ config = object.__new__(VllmConfig)
286
+ config.attention_config = AttentionConfig(hisparse_config=HiSparseConfig())
287
+
288
+ with patch.object(envs, "VLLM_USE_V2_MODEL_RUNNER", None):
289
+ assert config.use_v2_model_runner
290
+ with (
291
+ patch.object(envs, "VLLM_USE_V2_MODEL_RUNNER", False),
292
+ pytest.raises(ValueError, match="requires Model Runner V2"),
293
+ ):
294
+ _ = config.use_v2_model_runner
295
+
296
+
297
+def test_hisparse_rejects_decode_context_parallelism(monkeypatch):
298
+ monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
299
+ monkeypatch.setattr(current_platform, "device_count", lambda: 2)
300
+ with pytest.raises(ValueError, match="decode context parallelism"):
301
+ VllmConfig(
302
+ attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
303
+ parallel_config=ParallelConfig(
304
+ tensor_parallel_size=2,
305
+ decode_context_parallel_size=2,
306
+ ),
307
+ )
308
+
309
+
310
+def test_hisparse_rejects_pipeline_parallelism(monkeypatch):
311
+ monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
312
+ monkeypatch.setattr(current_platform, "device_count", lambda: 2)
313
+ with pytest.raises(ValueError, match="pipeline parallelism"):
314
+ VllmConfig(
315
+ attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
316
+ parallel_config=ParallelConfig(pipeline_parallel_size=2),
317
+ )
318
+
319
+
320
+def test_hisparse_rejects_disabled_hybrid_kv_cache_manager(monkeypatch):
321
+ monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
322
+ monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True)
323
+ with pytest.raises(ValueError, match="requires the hybrid KV cache manager"):
324
+ VllmConfig(
325
+ attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
326
+ scheduler_config=SchedulerConfig(
327
+ max_model_len=2048,
328
+ is_encoder_decoder=False,
329
+ disable_hybrid_kv_cache_manager=True,
330
+ ),
331
+ )
332
+
333
+
334
+def test_hisparse_rejects_non_cuda(monkeypatch):
335
+ monkeypatch.setattr(current_platform, "is_cuda", lambda: False)
336
+ with pytest.raises(ValueError, match="requires NVIDIA CUDA"):
337
+ VllmConfig(attention_config=AttentionConfig(hisparse_config=HiSparseConfig()))
338
+
339
+
340
def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch):
341
"""ROCm keeps the DSA models (DeepSeek V3.2/V4, GLM-5.2) on their compiled
342
MRV1 paths and off breakable cudagraphs by default."""
@@ -305,7 +363,8 @@ def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch):
363
# (warning_once args must be hashable for its lru_cache).
364
monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False)
365
config = SimpleNamespace(
308
- model_config=SimpleNamespace(architectures=["DeepseekV32ForCausalLM"])
366
+ model_config=SimpleNamespace(architectures=["DeepseekV32ForCausalLM"]),
367
+ attention_config=AttentionConfig(),
368
)
369
assert VllmConfig.use_v2_model_runner.fget(config) is False
370
finally:
@@ -385,6 +444,7 @@ def test_dsa_models_default_to_mrv2_and_breakable_cudagraph(
444
)
445
config = SimpleNamespace(
446
model_config=model_config,
447
+ attention_config=AttentionConfig(),
448
speculative_config=SimpleNamespace(method="mtp") if with_mtp else None,
449
parallel_config=SimpleNamespace(prefill_context_parallel_size=1),
450
compilation_config=CompilationConfig(
@@ -897,7 +957,10 @@ def test_models_default_to_v2_model_runner(model_config, expected, monkeypatch):
957
monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False)
958
monkeypatch.setattr(vllm_config_module, "HAS_TRITON", True)
959
monkeypatch.setattr(current_platform, "is_rocm", lambda: False)
900
- config = SimpleNamespace(model_config=model_config)
960
+ config = SimpleNamespace(
961
+ model_config=model_config,
962
+ attention_config=AttentionConfig(),
963
+ )
964
config._get_v2_model_runner_unsupported_features = lambda: []
965
966
assert VllmConfig.use_v2_model_runner.fget(config) is expected
@@ -2787,6 +2850,7 @@ def test_draft_sample_method_gumbel_is_rejected():
2850
def _watermarked_vllm_config() -> VllmConfig:
2851
config = object.__new__(VllmConfig)
2852
config.watermark_config = WatermarkConfig(key=42)
2853
+ config.attention_config = AttentionConfig()
2854
config.speculative_config = None
2855
return config
2856
tests/v1/attention/test_backend_per_kind.py
+54
-2
@@ -1,16 +1,19 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Tests for per-KV-group attention backend selection (backend_per_kind)."""
3
+"""Tests for attention configuration and per-KV-group backend selection."""
4
+
5
+from types import SimpleNamespace
6
7
import pytest
8
7
-from vllm.config.attention import AttentionConfig
9
+from vllm.config.attention import AttentionConfig, HiSparseConfig
10
from vllm.model_executor.layers.attention.attention import (
11
_largest_kernel_block_within,
12
)
13
from vllm.v1.attention.backend import AttentionType, MultipleOf
14
from vllm.v1.attention.backends.registry import AttentionBackendEnum
15
from vllm.v1.attention.selector import get_attn_spec_kind
16
+from vllm.v1.hisparse.runtime import ResolvedHiSparseConfig
17
from vllm.v1.kv_cache_interface import KVCacheSpecKind
18
19
@@ -86,3 +89,52 @@ def test_backend_per_kind_rejects_unknown_kind():
89
90
def test_backend_per_kind_defaults_empty():
91
assert AttentionConfig().backend_per_kind == {}
92
+
93
+
94
+def test_hisparse_device_buffer_size_boundaries():
95
+ vllm_config = SimpleNamespace(
96
+ attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
97
+ speculative_config=None,
98
+ )
99
+ resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
100
+ assert resolved is not None
101
+ assert resolved.device_buffer_size == 256
102
+
103
+ vllm_config.attention_config.hisparse_config = HiSparseConfig(
104
+ device_buffer_size=127
105
+ )
106
+ with pytest.raises(ValueError, match="expected at least 128"):
107
+ ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
108
+
109
+ vllm_config.attention_config.hisparse_config = HiSparseConfig(
110
+ device_buffer_size=32768
111
+ )
112
+ resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
113
+ assert resolved is not None
114
+ assert resolved.device_buffer_size == 32768
115
+
116
+ vllm_config.attention_config.hisparse_config = HiSparseConfig(
117
+ device_buffer_size=32769
118
+ )
119
+ with pytest.raises(ValueError, match="int16 slot-index limit"):
120
+ ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
121
+
122
+
123
+def test_hisparse_device_buffer_covers_speculative_window():
124
+ vllm_config = SimpleNamespace(
125
+ attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
126
+ speculative_config=SimpleNamespace(
127
+ num_speculative_tokens=3,
128
+ parallel_drafting=False,
129
+ ),
130
+ )
131
+
132
+ resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
133
+ assert resolved is not None
134
+ assert resolved.device_buffer_size == 5 * 128
135
+
136
+ vllm_config.attention_config.hisparse_config = HiSparseConfig(
137
+ device_buffer_size=4 * 128 - 1
138
+ )
139
+ with pytest.raises(ValueError, match="expected at least 512"):
140
+ ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
tests/v1/attention/test_mla_backends.py
+1
@@ -200,6 +200,7 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype(
200
layer = SimpleNamespace(
201
kv_cache_dtype=cache_dtype,
202
head_size=576,
203
+ indexer=None,
204
non_causal_multi_token_decode=False,
205
sliding_window=None,
206
)
tests/v1/attention/test_mla_context_chunks.py
+23
@@ -55,6 +55,7 @@ def build_chunked_context(
55
def test_sparse_context_lengths_do_not_force_gpu_sync(monkeypatch):
56
"""Sparse MLA must pin computed context lengths before the H2D copy."""
57
builder = SimpleNamespace(
58
+ _context_lens_cpu=torch.empty(2, dtype=torch.int32, pin_memory=True),
59
chunked_prefill_workspace=torch.empty((2048, 1)),
60
chunked_prefill_workspace_size=1024,
61
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
114
query_lens = [4] * len(context_lens)
115
metadata = build_chunked_context(context_lens, query_lens, workspace_size)
116
assert metadata is not None
117
+ assert metadata.context_lens.tolist() == context_lens
118
119
gathered: dict[int, list[tuple[int, int]]] = {}
120
previous_request_start = -1
@@ -253,6 +255,7 @@ def test_dcp_chunks_fit_the_per_rank_row_budget():
255
dcp_local_block_size=interleave,
256
)
257
assert metadata is not None
258
+ assert metadata.context_lens.tolist() == context_lens
259
virtual_block_size = interleave * dcp_world_size
260
261
local_cursor: dict[int, int] = {}
@@ -330,3 +333,23 @@ def test_dcp_reorg_uses_each_chunks_local_starts():
333
)
334
335
torch.testing.assert_close(reorganized, torch.cat(expected))
336
+
337
+
338
+def test_chunked_context_metadata_owns_context_lengths():
339
+ """Reusing CPU scratch must not overwrite an earlier batch's context lengths."""
340
+ context_lens = torch.tensor([16, 32], dtype=torch.int32)
341
+ metadata = build_mla_chunked_context_metadata(
342
+ context_lens_cpu=context_lens,
343
+ prefill_query_start_loc_cpu=torch.tensor([0, 4, 8], dtype=torch.int32),
344
+ chunked_prefill_workspace=torch.empty((64, 1)),
345
+ chunked_prefill_workspace_size=64,
346
+ block_size=BLOCK_SIZE,
347
+ align_chunk_to_block=True,
348
+ device=torch.device("cpu"),
349
+ dcp_world_size=1,
350
+ dcp_local_block_size=1,
351
+ dcp_virtual_block_size=1,
352
+ )
353
+ context_lens.zero_()
354
+ assert metadata is not None
355
+ assert metadata.context_lens.tolist() == [16, 32]
tests/v1/attention/test_sparse_mla_backends.py
+1946
-161
@@ -3,7 +3,9 @@
3
"""Unit tests for the sparse MLA backends and utilities."""
4
5
import math
6
+from collections import deque
7
from types import MethodType, SimpleNamespace
8
+from unittest.mock import MagicMock
9
10
import pytest
11
import torch
@@ -20,13 +22,19 @@ from tests.v1.attention.utils import (
22
create_vllm_config,
23
)
24
from vllm import _custom_ops as ops
23
-from vllm.config import set_current_vllm_config
25
+from vllm.config import HiSparseConfig, SpeculativeConfig, set_current_vllm_config
26
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
27
+ HiSparseConnectorWorker,
28
+)
29
+from vllm.model_executor.layers.attention import mla_attention
30
from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha
31
from vllm.model_executor.layers.attention.sparse_mla_attention import (
32
GLOBAL_TOPK_MASK_MAX_BYTES,
33
SparseMLACommonImpl,
34
+ SparseMLAPrefillMetadata,
35
_masked_mha_workspace_fits,
36
_topk_mask_shape,
37
+ _use_dense_mha_prefill,
38
)
39
from vllm.model_executor.layers.linear import ColumnParallelLinear
40
from vllm.platforms import current_platform
@@ -41,7 +49,6 @@ if not current_platform.is_cuda():
49
allow_module_level=True,
50
)
51
44
-import vllm.v1.attention.backends.mla.flashinfer_mla_sparse as flashinfer_sparse_mod
52
from vllm.model_executor.layers.attention.mla_attention import (
53
_canonicalize_sparse_mla_kv_cache_dtype,
54
)
@@ -49,10 +56,19 @@ from vllm.model_executor.layers.attention.sparse_mla_attention import (
56
SharedTopkIndicesBuffer,
57
)
58
from vllm.utils.math_utils import cdiv
59
+from vllm.utils.torch_utils import current_stream
60
+from vllm.v1.attention.backends.mla import index_group as index_group_module
61
+from vllm.v1.attention.backends.mla.flashattn_mla_sparse import (
62
+ FlashAttnMLASparseImpl,
63
+)
64
from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
65
FlashInferMLASparseImpl,
66
+ FlashInferMLASparseMetadataBuilder,
67
FlashInferMLASparseTRTLLMBackend,
68
)
69
+from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import (
70
+ FlashInferMLASparseSM120Impl,
71
+)
72
from vllm.v1.attention.backends.mla.flashmla_sparse import (
73
FlashMLASparseBackend,
74
FlashMLASparseImpl,
@@ -60,12 +76,29 @@ from vllm.v1.attention.backends.mla.flashmla_sparse import (
76
FlashMLASparseMetadataBuilder,
77
triton_convert_req_index_to_global_index,
78
)
63
-from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadataBuilder
79
+from vllm.v1.attention.backends.mla.index_group import (
80
+ HiSparseMLAIndexGroup,
81
+ SparseMLAIndexGroupBuilder,
82
+)
83
+from vllm.v1.attention.backends.mla.indexer import (
84
+ DeepseekV32IndexerMetadataBuilder,
85
+)
86
+from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
87
from vllm.v1.attention.backends.utils import (
88
split_decodes_and_prefills,
89
split_prefill_chunks,
90
)
91
from vllm.v1.attention.ops import flashmla
92
+from vllm.v1.hisparse import runtime as hisparse_runtime
93
+from vllm.v1.hisparse.runtime import (
94
+ HiSparseCacheHandle,
95
+ HiSparseRuntime,
96
+ ResolvedHiSparseConfig,
97
+ _has_hisparse_ops,
98
+ build_hisparse_prefill_staging_plan,
99
+ hisparse_prefill_staging_remap,
100
+)
101
+from vllm.v1.hisparse.types import SparseKVRowMirror
102
103
SPARSE_BACKEND_BATCH_SPECS = {
104
name: BATCH_SPECS[name]
@@ -113,10 +146,11 @@ def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch):
146
impl.bmm2_scale = None
147
impl.is_nope_mla = True
148
impl.need_to_return_lse_for_decode = False
116
- monkeypatch.setattr(
117
- flashinfer_sparse_mod,
118
- "triton_convert_req_index_to_global_index",
119
- lambda *args, **kwargs: (topk, torch.ones(1, dtype=torch.int32)),
149
+ impl.index_group = MagicMock()
150
+ impl.index_group_index = 0
151
+ impl.index_group.convert_logical_to_physical_topk.return_value = (
152
+ topk,
153
+ torch.ones(1, dtype=torch.int32),
154
)
155
156
import flashinfer.decode
@@ -142,6 +176,41 @@ def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch):
176
assert recorded_scale != kv_lora_rank**-0.5
177
178
179
+def test_hisparse_routes_prefill_to_sparse_mqa():
180
+ layer = SimpleNamespace(hisparse_cache=object())
181
+
182
+ assert not mla_attention.MLAAttention._use_sparse_mha(layer, SimpleNamespace())
183
+
184
+
185
+def test_hisparse_metadata_keeps_short_prefill_indexer_enabled():
186
+ config = SimpleNamespace(
187
+ attention_config=SimpleNamespace(
188
+ hisparse_config=object(), sparse_mla_force_mqa=False
189
+ )
190
+ )
191
+
192
+ assert not _use_dense_mha_prefill(config, prefill_max_seq_len=512, topk_tokens=2048)
193
+
194
+
195
+def test_unified_mla_prepares_batch_without_slot_mapping(monkeypatch):
196
+ layer = SimpleNamespace(
197
+ prepare_kv_cache_update=MagicMock(),
198
+ update_kv_cache=MagicMock(),
199
+ )
200
+ metadata = object()
201
+ monkeypatch.setattr(
202
+ mla_attention,
203
+ "get_attention_context",
204
+ lambda _: (metadata, layer, torch.empty(0), None),
205
+ )
206
+ tensor = torch.empty(0, device=DEVICE_TYPE)
207
+
208
+ mla_attention.unified_mla_kv_cache_update(tensor, tensor, "layer", "auto", tensor)
209
+
210
+ layer.prepare_kv_cache_update.assert_called_once_with(metadata)
211
+ layer.update_kv_cache.assert_not_called()
212
+
213
+
214
def _float_to_e8m0_truncate(f: float) -> float:
215
"""Simulate SM100's float -> e8m0 -> bf16 scale conversion.
216
e8m0 format only stores the exponent (power of 2).
@@ -618,8 +687,6 @@ def test_sparse_backend_decode_correctness(
687
688
# The sparse builder clones the layer's dense-MHA prefill backend from
689
# static_forward_context; register a mock layer carrying one.
621
- from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
622
-
690
prefill_backend = get_mla_prefill_backend(vllm_config)(
691
num_heads=num_heads,
692
scale=scale,
@@ -972,6 +1039,8 @@ def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens():
1039
),
1040
block_table=torch.arange(40, dtype=torch.int32, device=device).view(4, 10),
1041
block_size=block_size,
1042
+ num_decode_tokens=num_mqa_tokens,
1043
+ num_actual_tokens=num_mqa_tokens,
1044
)
1045
assert attn_metadata.req_id_per_token.shape[0] == num_batch_tokens
1046
@@ -990,11 +1059,30 @@ def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens():
1059
def _stub_kernel(q, kv, indices, lengths, actual_num_heads):
1060
captured["indices"] = indices
1061
captured["actual_num_heads"] = actual_num_heads
993
- return torch.zeros(q.shape[0], q.shape[1], 512, dtype=q.dtype, device=q.device)
1062
+ return (
1063
+ torch.zeros(q.shape[0], q.shape[1], 512, dtype=q.dtype, device=q.device),
1064
+ None,
1065
+ )
1066
+
1067
+ def _convert_topk(indices, metadata, **kwargs):
1068
+ return triton_convert_req_index_to_global_index(
1069
+ metadata.req_id_per_token[: indices.shape[0]],
1070
+ metadata.block_table,
1071
+ indices,
1072
+ BLOCK_SIZE=metadata.block_size,
1073
+ BLOCK_STRIDE_ROWS=kwargs["block_stride_rows"],
1074
+ NUM_TOPK_TOKENS=indices.shape[1],
1075
+ return_valid_counts=kwargs["return_valid_counts"],
1076
+ )
1077
995
- stub_impl = SimpleNamespace(_bf16_flash_mla_kernel=_stub_kernel)
1078
+ stub_impl = SimpleNamespace(
1079
+ _bf16_flash_mla_kernel=_stub_kernel,
1080
+ _convert_logical_to_physical_topk=_convert_topk,
1081
+ index_group=None,
1082
+ index_group_index=0,
1083
+ )
1084
997
- out = FlashMLASparseImpl._forward_bf16_kv(
1085
+ out, _ = FlashMLASparseImpl._forward_bf16_kv(
1086
stub_impl, q, kv_cache, topk_indices, attn_metadata, q.shape[1]
1087
)
1088
@@ -1283,8 +1371,6 @@ def test_sparse_backend_prefill_correctness(
1371
1372
# The sparse builder clones the layer's dense-MHA prefill backend from
1373
# static_forward_context; register a mock layer carrying one.
1286
- from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
1287
-
1374
prefill_backend = get_mla_prefill_backend(vllm_config)(
1375
num_heads=num_heads,
1376
scale=scale,
@@ -1524,189 +1610,1775 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int):
1610
assert torch.all(result[row, num_valid:] == -1)
1611
1612
1527
-def test_flashmla_cache_dtype_aliases_use_ds_layout():
1528
- # kv-cache dtype aliases are canonicalized to fp8_ds_mla before the layer
1529
- # stores kv_cache_dtype, so they cannot bypass the gate.
1530
- for alias in ("fp8", "fp8_e4m3"):
1531
- assert (
1532
- _canonicalize_sparse_mla_kv_cache_dtype(FlashMLASparseBackend, alias)
1533
- == "fp8_ds_mla"
1534
- )
1613
+def test_sparse_mla_index_group_converts_decode_indices_once(monkeypatch):
1614
+ device = torch.device(DEVICE_TYPE)
1615
+ logical = torch.full((2, 128), -1, dtype=torch.int32, device=device)
1616
+ logical[:, :4] = torch.tensor(
1617
+ [[0, 2, 5, -1], [1, 3, 7, -1]], dtype=torch.int32, device=device
1618
+ )
1619
+ builder = SparseMLAIndexGroupBuilder(logical)
1620
+ group, leader_index = builder.register_layer(True)
1621
+ follower_group, follower_index = builder.register_layer(False)
1622
+ assert follower_group is group
1623
+
1624
+ block_table = torch.tensor([[2, 3], [4, 5]], dtype=torch.int32, device=device)
1625
+ metadata = SimpleNamespace(
1626
+ block_table=block_table,
1627
+ block_size=4,
1628
+ req_id_per_token=torch.tensor([0, 1], dtype=torch.int32, device=device),
1629
+ num_decode_tokens=2,
1630
+ num_actual_tokens=2,
1631
+ )
1632
+ calls = []
1633
+ convert = index_group_module.triton_convert_req_index_to_global_index
1634
1635
+ def convert_spy(*args, **kwargs):
1636
+ calls.append((args, kwargs))
1637
+ return convert(*args, **kwargs)
1638
1537
-def test_flashmla_fp8_metadata_reuses_common_batch_split():
1538
- builder = SimpleNamespace(
1539
- device=torch.device(DEVICE_TYPE),
1540
- vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=8)),
1639
+ monkeypatch.setattr(
1640
+ index_group_module,
1641
+ "triton_convert_req_index_to_global_index",
1642
+ convert_spy,
1643
)
1542
- common_metadata = SimpleNamespace(
1543
- num_actual_tokens=1,
1544
- seq_lens_cpu_upper_bound=torch.tensor([1]),
1545
- query_start_loc_cpu=torch.tensor([0, 1]),
1546
- block_table_tensor=torch.zeros(1, 1, dtype=torch.int32, device=DEVICE_TYPE),
1644
+ group.set_logical_topk_ready(leader_index)
1645
+ leader_result = group.convert_logical_to_physical_topk(
1646
+ leader_index,
1647
+ logical,
1648
+ metadata,
1649
+ block_stride_rows=None,
1650
+ return_valid_counts=True,
1651
)
1548
- metadata = FlashMLASparseMetadata(
1549
- num_reqs=1,
1550
- max_query_len=1,
1551
- max_seq_len=1,
1552
- num_actual_tokens=1,
1553
- query_start_loc=torch.tensor([0, 1], device=DEVICE_TYPE),
1554
- slot_mapping=torch.tensor([0], device=DEVICE_TYPE),
1555
- block_table=torch.zeros(1, 1, dtype=torch.int32, device=DEVICE_TYPE),
1556
- req_id_per_token=torch.zeros(1, dtype=torch.int32, device=DEVICE_TYPE),
1557
- num_decodes=0,
1558
- num_prefills=1,
1559
- num_decode_tokens=0,
1652
+ follower_result = group.convert_logical_to_physical_topk(
1653
+ follower_index,
1654
+ logical,
1655
+ metadata,
1656
+ block_stride_rows=None,
1657
+ return_valid_counts=True,
1658
)
1659
+ torch.accelerator.synchronize()
1660
1562
- fp8_metadata = FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode(
1563
- builder, common_metadata, metadata
1661
+ assert len(calls) == 1
1662
+ leader_indices, leader_counts = leader_result
1663
+ follower_indices, follower_counts = follower_result
1664
+ torch.testing.assert_close(follower_indices, leader_indices)
1665
+ torch.testing.assert_close(follower_counts, leader_counts)
1666
+
1667
+
1668
+def test_sparse_mla_index_group_falls_back_for_prefill_sized_batch(monkeypatch):
1669
+ device = torch.device(DEVICE_TYPE)
1670
+ logical = torch.full((4, 128), -1, dtype=torch.int32, device=device)
1671
+ logical[:, 0] = torch.tensor([0, 1, 2, 3], dtype=torch.int32, device=device)
1672
+ builder = SparseMLAIndexGroupBuilder(logical, max_decode_rows=2)
1673
+ group, leader_index = builder.register_layer(True)
1674
+ follower_group, follower_index = builder.register_layer(False)
1675
+ assert follower_group is group
1676
+
1677
+ metadata = SimpleNamespace(
1678
+ block_table=torch.tensor(
1679
+ [[2], [3], [4], [5]], dtype=torch.int32, device=device
1680
+ ),
1681
+ block_size=4,
1682
+ req_id_per_token=torch.arange(4, dtype=torch.int32, device=device),
1683
+ )
1684
+ calls = 0
1685
+ convert = index_group_module.triton_convert_req_index_to_global_index
1686
+
1687
+ def convert_spy(*args, **kwargs):
1688
+ nonlocal calls
1689
+ calls += 1
1690
+ return convert(*args, **kwargs)
1691
+
1692
+ monkeypatch.setattr(
1693
+ index_group_module,
1694
+ "triton_convert_req_index_to_global_index",
1695
+ convert_spy,
1696
+ )
1697
+ leader_result = group.convert_logical_to_physical_topk(
1698
+ leader_index,
1699
+ logical,
1700
+ metadata,
1701
+ block_stride_rows=None,
1702
+ return_valid_counts=False,
1703
+ )
1704
+ follower_result = group.convert_logical_to_physical_topk(
1705
+ follower_index,
1706
+ logical,
1707
+ metadata,
1708
+ block_stride_rows=None,
1709
+ return_valid_counts=False,
1710
)
1711
+ torch.accelerator.synchronize()
1712
1566
- assert fp8_metadata.num_decodes == 0
1567
- assert fp8_metadata.num_prefills == 1
1568
- assert fp8_metadata.num_decode_tokens == 0
1569
- assert fp8_metadata.num_prefill_tokens == 1
1713
+ assert calls == 2
1714
+ torch.testing.assert_close(follower_result, leader_result)
1715
1716
1572
-def test_flashmla_common_metadata_requires_uniform_decodes():
1573
- common_metadata = SimpleNamespace(
1574
- max_query_len=3,
1575
- num_reqs=3,
1576
- num_actual_tokens=6,
1577
- query_start_loc_cpu=torch.tensor([0, 1, 3, 6]),
1578
- is_prefilling=None,
1717
+def test_sparse_mla_index_groups_own_distinct_physical_buffers():
1718
+ logical = torch.empty((8, 128), dtype=torch.int32, device=DEVICE_TYPE)
1719
+ builder = SparseMLAIndexGroupBuilder(logical, max_decode_rows=2)
1720
+
1721
+ first, _ = builder.register_layer(True)
1722
+ follower, _ = builder.register_layer(False)
1723
+ second, _ = builder.register_layer(True)
1724
+
1725
+ assert follower is first
1726
+ assert second is not first
1727
+ assert first.physical_topk_indices[:2].shape == (2, 128)
1728
+ assert second.physical_topk_indices[:2].shape == (2, 128)
1729
+ assert first.physical_topk_indices.data_ptr() != (
1730
+ second.physical_topk_indices.data_ptr()
1731
)
1732
1581
- split = split_decodes_and_prefills(
1582
- common_metadata,
1583
- decode_threshold=128,
1584
- require_uniform=FlashMLASparseMetadataBuilder.require_uniform_decodes,
1733
+
1734
+# HiSparse is host-resident-only and kernel-only: runtime construction
1735
+# raises without the compiled CUDA ops.
1736
+requires_hisparse_ops = pytest.mark.skipif(
1737
+ not _has_hisparse_ops(),
1738
+ reason="HiSparse CUDA ops not compiled",
1739
+)
1740
+
1741
+
1742
+def _make_hisparse_runtime(
1743
+ *,
1744
+ top_k: int = 4,
1745
+ device_buffer_size: int = 5,
1746
+ max_num_reqs: int = 2,
1747
+ row_width: int = 8,
1748
+ block_size: int = 64,
1749
+ max_swap_rows: int | None = None,
1750
+ index_group: hisparse_runtime.HiSparseIndexGroup | None = None,
1751
+) -> HiSparseRuntime:
1752
+ runtime = HiSparseRuntime(
1753
+ config=ResolvedHiSparseConfig(
1754
+ top_k=top_k,
1755
+ device_buffer_size=device_buffer_size,
1756
+ ),
1757
+ max_num_reqs=max_num_reqs,
1758
+ row_width=row_width,
1759
+ kv_dtype=torch.float32,
1760
+ device=DEVICE_TYPE,
1761
+ max_swap_rows=max_swap_rows,
1762
+ index_group=index_group,
1763
+ )
1764
+ blocks_per_request = cdiv(runtime.region_stride, block_size)
1765
+ # Leave one extra physical block for resident-tier tests. HMA group layouts
1766
+ # alias this backing, but distinct group allocations use distinct block IDs.
1767
+ num_blocks = max_num_reqs * blocks_per_request + 1
1768
+ raw = torch.zeros(
1769
+ num_blocks * block_size * row_width,
1770
+ dtype=torch.float32,
1771
+ device=DEVICE_TYPE,
1772
+ ).view(torch.int8)
1773
+ block_table = torch.arange(
1774
+ max_num_reqs * blocks_per_request,
1775
+ dtype=torch.int32,
1776
+ device=DEVICE_TYPE,
1777
+ ).view(max_num_reqs, blocks_per_request)
1778
+ runtime.bind_hot_cache(
1779
+ raw,
1780
+ byte_offset=0,
1781
+ block_stride=block_size * row_width * torch.float32.itemsize,
1782
+ num_blocks=num_blocks,
1783
+ block_size=block_size,
1784
+ block_table=block_table,
1785
+ )
1786
+ runtime.request_state_indices = torch.arange(
1787
+ max_num_reqs, dtype=torch.int32, device=runtime.device
1788
+ )
1789
+ return runtime
1790
+
1791
+
1792
+def _make_hisparse_cache_handle(
1793
+ *,
1794
+ top_k: int = 4,
1795
+ device_buffer_size: int = 5,
1796
+ max_num_reqs: int = 2,
1797
+ row_width: int = 8,
1798
+ block_size: int = 64,
1799
+ max_swap_rows: int | None = None,
1800
+) -> HiSparseCacheHandle:
1801
+ runtime = _make_hisparse_runtime(
1802
+ top_k=top_k,
1803
+ device_buffer_size=device_buffer_size,
1804
+ max_num_reqs=max_num_reqs,
1805
+ row_width=row_width,
1806
+ block_size=block_size,
1807
+ max_swap_rows=max_swap_rows,
1808
)
1809
+ return HiSparseCacheHandle(runtime)
1810
1587
- assert split == (1, 2, 1, 5)
1811
1812
+@requires_hisparse_ops
1813
+def test_hisparse_uses_graph_stable_request_state_mapping():
1814
+ device = torch.device(DEVICE_TYPE)
1815
+ block_size, row_width = 64, 8
1816
+ runtime = _make_hisparse_runtime(
1817
+ top_k=1,
1818
+ device_buffer_size=2,
1819
+ max_num_reqs=2,
1820
+ row_width=row_width,
1821
+ block_size=block_size,
1822
+ )
1823
+ runtime.bind_source_cache(
1824
+ torch.arange(2 * block_size * row_width, dtype=torch.float32)
1825
+ .view(2, block_size, row_width)
1826
+ .pin_memory()
1827
+ )
1828
+ runtime.request_state_indices = torch.tensor([1], dtype=torch.int32, device=device)
1829
1590
-def test_flashmla_fp8_metadata_excludes_zero_token_decode_padding(monkeypatch):
1591
- monkeypatch.setattr(
1592
- "vllm.v1.attention.backends.mla.flashmla_sparse.get_mla_metadata",
1593
- lambda: (object(), None),
1830
+ cache_handle = HiSparseCacheHandle(runtime)
1831
+ cache_handle.all_context_pages_resident = False
1832
+ runtime.begin_forward()
1833
+ cache_handle.swap_in(
1834
+ req_id_per_token=torch.tensor([0], dtype=torch.int32, device=device),
1835
+ block_table=torch.tensor([[1]], dtype=torch.int32, device=device),
1836
+ logical_topk_indices=torch.tensor([[0]], dtype=torch.int32, device=device),
1837
+ block_size=block_size,
1838
)
1595
- builder = SimpleNamespace(
1596
- device=torch.device(DEVICE_TYPE),
1597
- dummy_block_table=torch.zeros(7, 1, device=DEVICE_TYPE),
1598
- max_model_len_tensor=torch.zeros(7, device=DEVICE_TYPE),
1839
+ torch.accelerator.synchronize()
1840
+
1841
+ assert (runtime.index_group.device_global_indices[0] == -1).all()
1842
+ assert (runtime.index_group.device_global_indices[1] == block_size).any()
1843
+
1844
+
1845
+@requires_hisparse_ops
1846
+def test_hisparse_maps_speculative_rows_through_request_state():
1847
+ """Multiple verification rows for one request share its persistent state."""
1848
+ device = torch.device(DEVICE_TYPE)
1849
+ block_size, row_width = 64, 8
1850
+ runtime = _make_hisparse_runtime(
1851
+ top_k=1,
1852
+ device_buffer_size=2,
1853
+ max_num_reqs=1,
1854
+ row_width=row_width,
1855
+ block_size=block_size,
1856
+ max_swap_rows=4,
1857
)
1600
- query_start_loc_cpu = torch.tensor([0, 110, 220, 330, 440, 550, 660, 660])
1601
- common_metadata = SimpleNamespace(
1602
- num_actual_tokens=660,
1603
- query_start_loc_cpu=query_start_loc_cpu,
1604
- seq_lens=torch.arange(7, device=DEVICE_TYPE),
1858
+ runtime.bind_source_cache(
1859
+ torch.zeros(1, block_size, row_width, dtype=torch.float32).pin_memory()
1860
)
1606
- metadata = FlashMLASparseMetadata(
1607
- num_reqs=7,
1608
- max_query_len=110,
1609
- max_seq_len=110,
1610
- num_actual_tokens=660,
1611
- query_start_loc=query_start_loc_cpu.to(DEVICE_TYPE),
1612
- slot_mapping=torch.arange(660, device=DEVICE_TYPE),
1613
- block_table=torch.zeros(7, 1, dtype=torch.int32, device=DEVICE_TYPE),
1614
- req_id_per_token=torch.zeros(660, dtype=torch.int32, device=DEVICE_TYPE),
1615
- num_decodes=7,
1616
- num_prefills=0,
1617
- num_decode_tokens=660,
1861
+
1862
+ cache_handle = HiSparseCacheHandle(runtime)
1863
+ resident_table = torch.tensor([[1]], dtype=torch.int32, device=device)
1864
+ cache_handle.bind_cache(
1865
+ runtime.hot.cache.view(torch.int8),
1866
+ byte_offset=0,
1867
+ block_stride=block_size * row_width * torch.float32.itemsize,
1868
+ num_blocks=2,
1869
+ block_size=block_size,
1870
+ block_table=resident_table,
1871
+ slot_mapping=torch.arange(4, dtype=torch.int64, device=device),
1872
+ )
1873
+ cache_handle.all_context_pages_resident = True
1874
+ runtime.begin_forward()
1875
+ physical = cache_handle.swap_in(
1876
+ req_id_per_token=torch.zeros(4, dtype=torch.int32, device=device),
1877
+ block_table=torch.zeros((1, 1), dtype=torch.int32, device=device),
1878
+ logical_topk_indices=torch.zeros((4, 1), dtype=torch.int32, device=device),
1879
+ block_size=block_size,
1880
)
1881
+ torch.accelerator.synchronize()
1882
1620
- fp8_metadata = FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode(
1621
- builder, common_metadata, metadata
1883
+ assert physical.tolist() == [[block_size]] * 4
1884
+
1885
+
1886
+@requires_hisparse_ops
1887
+def test_hisparse_resident_rows_bypass_hot_lru():
1888
+ device = torch.device(DEVICE_TYPE)
1889
+ block_size, row_width = 64, 8
1890
+ cache_handle = _make_hisparse_cache_handle(
1891
+ top_k=4,
1892
+ device_buffer_size=5,
1893
+ max_num_reqs=1,
1894
+ row_width=row_width,
1895
+ block_size=block_size,
1896
+ )
1897
+ raw = cache_handle.runtime.hot.cache.view(torch.int8)
1898
+ resident_table = torch.tensor([[1]], dtype=torch.int32, device=device)
1899
+ resident_slots = torch.tensor([block_size], dtype=torch.int64, device=device)
1900
+ cache_handle.bind_cache(
1901
+ raw,
1902
+ byte_offset=0,
1903
+ block_stride=block_size * row_width * torch.float32.itemsize,
1904
+ num_blocks=2,
1905
+ block_size=block_size,
1906
+ block_table=resident_table,
1907
+ slot_mapping=resident_slots,
1908
+ )
1909
+
1910
+ host = torch.randn(1, block_size, row_width).pin_memory()
1911
+ cache_handle.runtime.bind_source_cache(host)
1912
+ resident_rows = host[0].to(device).add(7)
1913
+ cache_handle.runtime.hot.cache[1].copy_(resident_rows)
1914
+ source_table = torch.tensor([[5]], dtype=torch.int32, device=device)
1915
+ request_ids = torch.zeros(1, dtype=torch.int32, device=device)
1916
+ topk = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32, device=device)
1917
+ lru_before = cache_handle.runtime.index_group.lru_slots.clone()
1918
+
1919
+ cache_handle.runtime.begin_forward()
1920
+ indices, valid_counts = cache_handle.swap_in(
1921
+ req_id_per_token=request_ids,
1922
+ block_table=source_table,
1923
+ logical_topk_indices=topk,
1924
+ block_size=block_size,
1925
+ return_valid_counts=True,
1926
)
1927
+ torch.accelerator.synchronize()
1928
1624
- assert fp8_metadata.num_decodes == 6
1625
- assert fp8_metadata.num_decode_tokens == 660
1626
- assert fp8_metadata.decode is not None
1627
- assert fp8_metadata.decode.decode_query_len == 110
1929
+ expected_indices = topk + block_size
1930
+ torch.testing.assert_close(indices, expected_indices)
1931
torch.testing.assert_close(
1629
- fp8_metadata.decode.seq_lens, torch.arange(6, device=DEVICE_TYPE)
1932
+ valid_counts, torch.tensor([topk.shape[1]], dtype=torch.int32, device=device)
1933
+ )
1934
+ torch.testing.assert_close(cache_handle.runtime.index_group.lru_slots, lru_before)
1935
+ torch.testing.assert_close(
1936
+ cache_handle.runtime.hot.attention_cache.view(-1, row_width)[
1937
+ indices.to(torch.long)
1938
+ ],
1939
+ resident_rows[topk.to(torch.long)],
1940
)
1941
1942
1633
-@pytest.mark.parametrize("use_mixed_batch", [False, True])
1634
-def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: bool):
1635
- num_decode_tokens = 2
1636
- num_batch_tokens = 5
1637
- q = torch.empty(num_decode_tokens, 2, 3, device=DEVICE_TYPE)
1638
- topk_indices = torch.empty(num_decode_tokens, 4, device=DEVICE_TYPE)
1639
- kernel_q_shapes = []
1640
-
1641
- def convert_indices(*args, **kwargs): # noqa: ARG001
1642
- assert not kwargs.get("HAS_PREFILL_WORKSPACE", False)
1643
- if not kwargs.get("return_valid_counts", False):
1644
- return topk_indices
1645
- valid_counts = torch.full(
1646
- (num_decode_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE
1943
+@requires_hisparse_ops
1944
+def test_hisparse_swap_in_preserves_rows_across_eviction():
1945
+ device = torch.device(DEVICE_TYPE)
1946
+ block_size = 64
1947
+ row_width = 64
1948
+ num_blocks = 64
1949
+ top_k = 128
1950
+ buf = 256
1951
+ num_reqs = 4
1952
+
1953
+ kv_pool = torch.randn(
1954
+ (num_blocks + 1, block_size, row_width), dtype=torch.float32
1955
+ ).pin_memory()
1956
+ flat_pool = kv_pool.reshape(-1, row_width)
1957
+
1958
+ runtime = _make_hisparse_runtime(
1959
+ top_k=top_k,
1960
+ device_buffer_size=buf,
1961
+ max_num_reqs=num_reqs,
1962
+ row_width=row_width,
1963
+ )
1964
+ runtime.bind_source_cache(kv_pool)
1965
+ cache = HiSparseCacheHandle(runtime)
1966
+ cache.all_context_pages_resident = False
1967
+ cache.runtime.begin_forward()
1968
+
1969
+ blocks_per_req = num_blocks // num_reqs
1970
+ block_table = torch.arange(
1971
+ 1, num_blocks + 1, dtype=torch.int32, device=device
1972
+ ).view(num_reqs, blocks_per_req)
1973
+ req_ids = torch.arange(num_reqs, dtype=torch.int32, device=device)
1974
+ seq_len = blocks_per_req * block_size
1975
+ base = torch.arange(top_k, dtype=torch.int32, device=device)
1976
+ for step in range(3):
1977
+ cache.runtime.begin_forward()
1978
+ topk = torch.stack(
1979
+ [(base + step * top_k + row * 17) % seq_len for row in range(num_reqs)]
1980
)
1648
- return topk_indices, valid_counts
1981
+ topk[:, -1] = -1
1982
+
1983
+ hot_indices, valid_counts = cache.swap_in(
1984
+ req_id_per_token=req_ids,
1985
+ block_table=block_table,
1986
+ logical_topk_indices=topk.clone(),
1987
+ block_size=block_size,
1988
+ return_valid_counts=True,
1989
+ )
1990
+ torch.accelerator.synchronize()
1991
1650
- monkeypatch.setattr(
1651
- "vllm.v1.attention.backends.mla.flashmla_sparse."
1652
- "triton_convert_req_index_to_global_index",
1653
- convert_indices,
1654
- )
1992
+ global_ref = _triton_convert_reference_impl(
1993
+ req_ids, block_table, topk, block_size, top_k
1994
+ )
1995
+ valid = hot_indices >= 0
1996
+ torch.testing.assert_close(
1997
+ valid_counts,
1998
+ (global_ref >= 0).sum(dim=1, dtype=torch.int32),
1999
+ rtol=0,
2000
+ atol=0,
2001
+ )
2002
+ gathered = cache.runtime.hot.attention_cache.reshape(-1, row_width)[
2003
+ hot_indices[valid].to(torch.long)
2004
+ ].cpu()
2005
+ expected = flat_pool[global_ref[valid].cpu().to(torch.long)]
2006
+ torch.testing.assert_close(gathered, expected)
2007
1656
- def run_kernel(**kwargs):
1657
- kernel_q_shapes.append(kwargs["q"].shape)
1658
- return kwargs["q"][..., :1], None
2008
1660
- if use_mixed_batch:
1661
- fp8_metadata = FlashMLASparseMetadata.FP8KernelMetadata(
1662
- scheduler_metadata=object(), # type: ignore[arg-type]
1663
- dummy_block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE),
1664
- cache_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE),
2009
+@requires_hisparse_ops
2010
+def test_hisparse_multi_step_swaps_match_independent():
2011
+ """The shared LRU must retain every speculative step until followers run."""
2012
+ device = torch.device(DEVICE_TYPE)
2013
+ torch.manual_seed(0)
2014
+ block_size = 64
2015
+ row_width = 64
2016
+ num_blocks = 64
2017
+ top_k = 128
2018
+ buf = 3 * top_k
2019
+ num_reqs = 4
2020
+
2021
+ kv_pool = torch.randn(
2022
+ (num_blocks, block_size, row_width), dtype=torch.float32
2023
+ ).pin_memory()
2024
+
2025
+ def make(
2026
+ index_group: hisparse_runtime.HiSparseIndexGroup | None = None,
2027
+ ) -> HiSparseCacheHandle:
2028
+ runtime = _make_hisparse_runtime(
2029
+ top_k=top_k,
2030
+ max_num_reqs=num_reqs,
2031
+ device_buffer_size=buf,
2032
+ row_width=row_width,
2033
+ block_size=block_size,
2034
+ max_swap_rows=2 * num_reqs,
2035
+ index_group=index_group,
2036
)
1666
- else:
1667
- FP8Meta = FlashMLASparseMetadata.FP8SeparatePrefillDecode
1668
- fp8_metadata = FP8Meta(
1669
- num_decodes=1,
1670
- num_prefills=1,
1671
- num_decode_tokens=num_decode_tokens,
1672
- num_prefill_tokens=num_batch_tokens - num_decode_tokens,
1673
- decode=FP8Meta.Decode(
1674
- seq_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE),
1675
- kernel_metadata=object(), # type: ignore[arg-type]
1676
- decode_query_len=num_decode_tokens,
1677
- ),
1678
- prefill=FP8Meta.Prefill(
1679
- request_ids=torch.empty(
1680
- num_batch_tokens, dtype=torch.int32, device=DEVICE_TYPE
1681
- ),
1682
- workspace_starts=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE),
1683
- chunks=[],
1684
- ),
2037
+ runtime.bind_source_cache(kv_pool)
2038
+ cache = HiSparseCacheHandle(runtime)
2039
+ cache.all_context_pages_resident = False
2040
+ return cache
2041
+
2042
+ blocks_per_req = num_blocks // num_reqs
2043
+ block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).view(
2044
+ num_reqs, blocks_per_req
2045
+ )
2046
+ req_ids = torch.arange(num_reqs, dtype=torch.int32, device=device)
2047
+ seq_len = blocks_per_req * block_size
2048
+ producer = make()
2049
+ shared = make(producer.runtime.index_group)
2050
+ indep = make()
2051
+ base = torch.arange(top_k, dtype=torch.int32, device=device)
2052
+ for iteration in range(2):
2053
+ producer.runtime.begin_forward()
2054
+ shared.runtime.begin_forward()
2055
+ indep.runtime.begin_forward()
2056
+ topks = [
2057
+ torch.stack(
2058
+ [
2059
+ (base + (iteration * 2 + step) * top_k + row * 17) % seq_len
2060
+ for row in range(num_reqs)
2061
+ ]
2062
+ )
2063
+ for step in range(2)
2064
+ ]
2065
+ kw = dict(
2066
+ req_id_per_token=req_ids,
2067
+ block_table=block_table,
2068
+ block_size=block_size,
2069
)
1686
- metadata = SimpleNamespace(
1687
- fp8_extra_metadata=fp8_metadata,
1688
- fp8_use_mixed_batch=use_mixed_batch,
1689
- num_actual_tokens=num_batch_tokens,
1690
- req_id_per_token=torch.empty(
1691
- num_batch_tokens, dtype=torch.int32, device=DEVICE_TYPE
1692
- ),
1693
- block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE),
1694
- block_size=64,
2070
+ producer_indices = [
2071
+ producer.swap_in(logical_topk_indices=topk.clone(), **kw) for topk in topks
2072
+ ]
2073
+ independent_indices = [
2074
+ indep.swap_in(logical_topk_indices=topk.clone(), **kw) for topk in topks
2075
+ ]
2076
+ follower_indices = [
2077
+ shared.swap_in(logical_topk_indices=topk.clone(), **kw) for topk in topks
2078
+ ]
2079
+ torch.accelerator.synchronize()
2080
+
2081
+ for idx_full, idx_indep, idx_shared in zip(
2082
+ producer_indices, independent_indices, follower_indices
2083
+ ):
2084
+ torch.testing.assert_close(idx_shared, idx_full, rtol=0, atol=0)
2085
+ torch.testing.assert_close(idx_indep, idx_full, rtol=0, atol=0)
2086
+ torch.testing.assert_close(shared.runtime.hot.cache, indep.runtime.hot.cache)
2087
+
2088
+ shared_rows = shared.runtime.hot.attention_cache.reshape(-1, row_width)
2089
+ host_rows = kv_pool.reshape(-1, row_width)
2090
+ for physical_rows, logical_rows in zip(follower_indices, topks):
2091
+ global_rows = _triton_convert_reference_impl(
2092
+ req_ids, block_table, logical_rows, block_size, top_k
2093
+ )
2094
+ valid = physical_rows >= 0
2095
+ torch.testing.assert_close(
2096
+ shared_rows[physical_rows[valid].to(torch.long)].cpu(),
2097
+ host_rows[global_rows[valid].cpu().to(torch.long)],
2098
+ )
2099
+
2100
+
2101
+@requires_hisparse_ops
2102
+def test_hisparse_multi_step_writes_request_major_output():
2103
+ device = torch.device(DEVICE_TYPE)
2104
+ block_size = 64
2105
+ row_width = 8
2106
+ num_reqs = 2
2107
+ query_len = 3
2108
+ top_k = 4
2109
+ cache = _make_hisparse_cache_handle(
2110
+ top_k=top_k,
2111
+ device_buffer_size=query_len * top_k,
2112
+ max_num_reqs=num_reqs,
2113
+ row_width=row_width,
2114
+ block_size=block_size,
2115
+ max_swap_rows=num_reqs * query_len,
2116
+ )
2117
+ host = torch.arange(
2118
+ (num_reqs + 1) * block_size * row_width, dtype=torch.float32
2119
+ ).view(num_reqs + 1, block_size, row_width)
2120
+ cache.runtime.bind_source_cache(host.pin_memory())
2121
+ block_table = torch.arange(1, num_reqs + 1, dtype=torch.int32, device=device).view(
2122
+ -1, 1
2123
+ )
2124
+ request_ids = torch.arange(num_reqs, dtype=torch.int32, device=device)
2125
+ logical = torch.tensor(
2126
+ [
2127
+ [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]],
2128
+ [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]],
2129
+ ],
2130
+ dtype=torch.int32,
2131
+ device=device,
2132
)
1696
- impl = SimpleNamespace(
1697
- kv_cache_dtype="fp8_ds_mla",
1698
- topk_indices_buffer=topk_indices,
1699
- num_heads=2,
1700
- kv_lora_rank=1,
1701
- dcp_world_size=1,
1702
- need_to_return_lse_for_decode=False,
1703
- _fp8_flash_mla_kernel=run_kernel,
2133
+ physical = torch.full_like(logical, -1)
2134
+ valid_counts = torch.zeros((num_reqs, query_len), dtype=torch.int32, device=device)
2135
+
2136
+ cache.runtime.begin_forward()
2137
+ for step in range(query_len):
2138
+ cache.swap_in(
2139
+ request_ids,
2140
+ block_table,
2141
+ logical[:, step],
2142
+ block_size=block_size,
2143
+ return_valid_counts=True,
2144
+ attention_indices_out=physical[:, step],
2145
+ valid_counts_out=valid_counts[:, step],
2146
+ )
2147
+ torch.accelerator.synchronize()
2148
+
2149
+ torch.testing.assert_close(
2150
+ valid_counts,
2151
+ torch.full_like(valid_counts, top_k),
2152
)
1705
- impl._forward_fp8_kv_mixed_batch = MethodType(
1706
- FlashMLASparseImpl._forward_fp8_kv_mixed_batch, impl
2153
+ global_rows = _triton_convert_reference_impl(
2154
+ torch.arange(num_reqs, device=device, dtype=torch.int32).repeat_interleave(
2155
+ query_len
2156
+ ),
2157
+ block_table,
2158
+ logical.view(num_reqs * query_len, top_k),
2159
+ block_size,
2160
+ top_k,
2161
)
1708
- impl._forward_fp8_kv_separate_prefill_decode = MethodType(
1709
- FlashMLASparseImpl._forward_fp8_kv_separate_prefill_decode, impl
2162
+ hot_rows = cache.runtime.hot.attention_cache.reshape(-1, row_width)
2163
+ host_rows = host.reshape(-1, row_width)
2164
+ torch.testing.assert_close(
2165
+ hot_rows[physical.view(num_reqs * query_len, top_k).to(torch.long)].cpu(),
2166
+ host_rows[global_rows.cpu().to(torch.long)],
2167
+ )
2168
+
2169
+
2170
+@requires_hisparse_ops
2171
+def test_hisparse_kv_update_writes_resident_and_staging_caches():
2172
+ device = torch.device(DEVICE_TYPE)
2173
+ block_size = 4
2174
+ row_width = 8
2175
+ cache_handle = _make_hisparse_cache_handle(
2176
+ max_num_reqs=1,
2177
+ row_width=row_width,
2178
+ block_size=block_size,
2179
+ )
2180
+ resident_slots = torch.tensor([4, 5, -1], dtype=torch.int64, device=device)
2181
+ cache_handle.bind_cache(
2182
+ cache_handle.runtime.hot.cache.view(torch.int8),
2183
+ byte_offset=0,
2184
+ block_stride=block_size * row_width * torch.float32.itemsize,
2185
+ num_blocks=cache_handle.runtime.hot.cache.shape[0],
2186
+ block_size=block_size,
2187
+ block_table=torch.tensor([[1]], dtype=torch.int32, device=device),
2188
+ slot_mapping=resident_slots,
2189
+ )
2190
+ cache_handle.mirror_staging_cache = torch.empty(
2191
+ (1, block_size, row_width), dtype=torch.float32, device=device
2192
+ )
2193
+ cache_handle.mirror_staging_slots = torch.arange(
2194
+ block_size, dtype=torch.int64, device=device
2195
+ )
2196
+ slots = torch.tensor([3, 7, -1], dtype=torch.int64, device=device)
2197
+ kv_c = torch.randn(8, row_width - 2, device=device)
2198
+ k_pe = torch.randn(8, 1, 2, device=device)
2199
+ cache_handle.num_actual_tokens = slots.numel()
2200
+ cache_handle.decode_batch = False
2201
+ cache_handle.host_mirror_required = True
2202
+ source_cache = torch.zeros_like(cache_handle.view.cache)
2203
+ impl = object.__new__(FlashMLASparseImpl)
2204
+ layer = SimpleNamespace(
2205
+ hisparse_cache=cache_handle,
2206
+ use_pcp=False,
2207
+ impl=impl,
2208
+ )
2209
+ mla_attention.MLAAttention.update_kv_cache(
2210
+ layer,
2211
+ kv_c,
2212
+ k_pe,
2213
+ source_cache,
2214
+ slots,
2215
+ SimpleNamespace(num_decode_tokens=0),
2216
+ "auto",
2217
+ torch.tensor(1.0, device=device),
2218
+ )
2219
+ torch.accelerator.synchronize()
2220
+
2221
+ expected = torch.cat([kv_c[:2], k_pe[:2, 0]], dim=-1).cpu()
2222
+ torch.testing.assert_close(
2223
+ cache_handle.view.cache.view(-1, row_width)[resident_slots[:2]],
2224
+ expected.to(device),
2225
+ )
2226
+ torch.testing.assert_close(source_cache, torch.zeros_like(source_cache))
2227
+ staged = cache_handle.mirror_staging_cache.view(-1, row_width)
2228
+ staged_expected = torch.cat([kv_c[:3], k_pe[:3, 0]], dim=-1)
2229
+ torch.testing.assert_close(staged[:3], staged_expected)
2230
+
2231
+
2232
+@requires_hisparse_ops
2233
+def test_hisparse_row_dma_copies_discontiguous_spans_across_layers():
2234
+ device = torch.device(DEVICE_TYPE)
2235
+ block_size = 4
2236
+ row_width = 8
2237
+ num_rows = 16
2238
+ resident_caches = tuple(
2239
+ (
2240
+ torch.arange(num_rows * row_width, dtype=torch.uint8, device=device)
2241
+ .add_(layer_index * 17)
2242
+ .view(num_rows // block_size, block_size, row_width)
2243
+ )
2244
+ for layer_index in range(2)
2245
+ )
2246
+ host_caches = tuple(
2247
+ torch.zeros((num_rows, row_width), dtype=torch.uint8, pin_memory=True)
2248
+ for _ in resident_caches
2249
+ )
2250
+ worker = object.__new__(HiSparseConnectorWorker)
2251
+ worker.is_host_writer = True
2252
+ worker.kernel_block_size = block_size
2253
+ worker.resident_caches = resident_caches
2254
+ worker.host_caches = host_caches
2255
+ worker.cache_handles = [
2256
+ SimpleNamespace(
2257
+ decode_batch=True,
2258
+ runtime=SimpleNamespace(resident_source_index=0),
2259
+ )
2260
+ for _ in resident_caches
2261
+ ]
2262
+ worker.dma_stream = torch.cuda.Stream(device=device)
2263
+ worker.host_write_event = torch.Event()
2264
+ worker._dma_free_descriptors = []
2265
+ worker._pending_dma_descriptors = deque()
2266
+ worker._pending_transfer_events = deque()
2267
+ worker._enqueued_transfer_ids = []
2268
+ worker._dma_submitted = False
2269
+ worker._set_row_mirrors(
2270
+ (
2271
+ SparseKVRowMirror((1,), 5, 2),
2272
+ SparseKVRowMirror((8,), 10, 4),
2273
+ )
2274
+ )
2275
+
2276
+ worker._enqueue_row_dma(range(2))
2277
+ current_stream().wait_event(worker.host_write_event)
2278
+ torch.accelerator.synchronize()
2279
+
2280
+ for resident, host in zip(resident_caches, host_caches):
2281
+ flat_resident = resident.view(num_rows, row_width).cpu()
2282
+ torch.testing.assert_close(host[5:7], flat_resident[1:3])
2283
+ torch.testing.assert_close(host[10:14], flat_resident[8:12])
2284
+
2285
+
2286
+@requires_hisparse_ops
2287
+def test_hisparse_remaps_strided_hma_rows_for_attention():
2288
+ device = torch.device(DEVICE_TYPE)
2289
+ block_size, row_width = 4, 8
2290
+ runtime = _make_hisparse_runtime(
2291
+ max_num_reqs=1,
2292
+ block_size=block_size,
2293
+ row_width=row_width,
2294
+ )
2295
+ page_elements = block_size * row_width
2296
+ stride_elements = 3 * page_elements
2297
+ byte_offset = page_elements * torch.float32.itemsize
2298
+ raw = torch.zeros(2 * stride_elements, dtype=torch.float32, device=device).view(
2299
+ torch.int8
2300
+ )
2301
+ runtime.bind_hot_cache(
2302
+ raw,
2303
+ byte_offset=byte_offset,
2304
+ block_stride=stride_elements * torch.float32.itemsize,
2305
+ num_blocks=2,
2306
+ block_size=block_size,
2307
+ block_table=torch.tensor([[0, 1]], dtype=torch.int32, device=device),
2308
+ )
2309
+
2310
+ source = (
2311
+ torch.arange(36 * row_width, dtype=torch.float32)
2312
+ .view(9, block_size, row_width)
2313
+ .pin_memory()
2314
+ )
2315
+ runtime.bind_source_cache(source)
2316
+ cache = HiSparseCacheHandle(runtime)
2317
+ cache.all_context_pages_resident = False
2318
+ runtime.begin_forward()
2319
+ attention_indices = cache.swap_in(
2320
+ req_id_per_token=torch.tensor([0], dtype=torch.int32, device=device),
2321
+ block_table=torch.arange(1, 9, dtype=torch.int32, device=device).view(1, 8),
2322
+ logical_topk_indices=torch.tensor(
2323
+ [[0, 1, 2, 3]], dtype=torch.int32, device=device
2324
+ ),
2325
+ block_size=block_size,
2326
+ )
2327
+
2328
+ assert runtime.hot.attention_cache is not None
2329
+ assert runtime.hot.attention_cache.is_contiguous()
2330
+ physical_indices = runtime.index_group.shared_topk.device_topk_rows[0]
2331
+ expected = (
2332
+ physical_indices // block_size * (stride_elements // row_width)
2333
+ + physical_indices % block_size
2334
+ )
2335
+ torch.testing.assert_close(attention_indices[0], expected)
2336
+
2337
+
2338
+@requires_hisparse_ops
2339
+def test_hisparse_gather_zeroes_unaligned_destination():
2340
+ """Invalid host rows are safely zeroed even at a byte-unaligned address."""
2341
+ device = torch.device(DEVICE_TYPE)
2342
+ row_bytes = 16
2343
+ host_cache = torch.ones(1, row_bytes, dtype=torch.uint8).pin_memory()
2344
+ storage = torch.full((2 * row_bytes + 1,), 7, dtype=torch.uint8, device=device)
2345
+ hot_cache = torch.as_strided(
2346
+ storage,
2347
+ size=(2, 1, row_bytes),
2348
+ stride=(row_bytes + 1, row_bytes, 1),
2349
+ )
2350
+ global_indices = torch.tensor([[1]], dtype=torch.int32, device=device)
2351
+ hot_indices = torch.tensor([[1]], dtype=torch.int32, device=device)
2352
+ miss_mask = torch.ones_like(global_indices)
2353
+
2354
+ torch.ops._C_cache_ops.hisparse_gather_plan(
2355
+ host_cache,
2356
+ hot_cache,
2357
+ global_indices,
2358
+ hot_indices,
2359
+ miss_mask,
2360
+ None,
2361
+ )
2362
+ torch.accelerator.synchronize()
2363
+
2364
+ torch.testing.assert_close(hot_cache[1, 0], torch.zeros_like(hot_cache[1, 0]))
2365
+ assert storage[row_bytes].item() == 7
2366
+
2367
+
2368
+@requires_hisparse_ops
2369
+def test_hisparse_gather_rejects_short_attention_block_stride():
2370
+ """Reject a stride that would let the gather kernel cross block bounds."""
2371
+ device = torch.device(DEVICE_TYPE)
2372
+ host_cache = torch.ones(1, 16, dtype=torch.uint8).pin_memory()
2373
+ hot_cache = torch.empty((1, 2, 16), dtype=torch.uint8, device=device)
2374
+ global_indices = torch.zeros((1, 1), dtype=torch.int32, device=device)
2375
+ hot_indices = torch.zeros_like(global_indices)
2376
+ miss_mask = torch.ones_like(global_indices)
2377
+ attention_indices = torch.empty_like(global_indices)
2378
+
2379
+ with pytest.raises(RuntimeError, match="stride must cover one hot block"):
2380
+ torch.ops._C_cache_ops.hisparse_gather_plan(
2381
+ host_cache,
2382
+ hot_cache,
2383
+ global_indices,
2384
+ hot_indices,
2385
+ miss_mask,
2386
+ None,
2387
+ attention_indices,
2388
+ 1,
2389
+ )
2390
+
2391
+
2392
+@requires_hisparse_ops
2393
+def test_hisparse_newest_write_and_recycled_slot_invalidation():
2394
+ """Newest writes clamp padding and invalidated rows are loaded again."""
2395
+ device = torch.device(DEVICE_TYPE)
2396
+ block_size = 4
2397
+ row_width = 8
2398
+ num_blocks = 8
2399
+
2400
+ kv_pool = (
2401
+ torch.arange(num_blocks * block_size * row_width, dtype=torch.float32)
2402
+ .view(num_blocks, block_size, row_width)
2403
+ .pin_memory()
2404
+ )
2405
+ flat_pool = kv_pool.reshape(-1, row_width)
2406
+
2407
+ cache_handle = _make_hisparse_cache_handle(block_size=block_size)
2408
+ cache_handle.runtime.eager_host_mirror = True
2409
+
2410
+ block_table = torch.tensor([[2, 0, 4]], dtype=torch.int32, device=device)
2411
+ resident_slots = torch.tensor([4, -1], dtype=torch.int64, device=device)
2412
+ cache_handle.bind_cache(
2413
+ cache_handle.runtime.hot.cache.view(torch.int8),
2414
+ byte_offset=0,
2415
+ block_stride=block_size * row_width * torch.float32.itemsize,
2416
+ num_blocks=cache_handle.runtime.hot.cache.shape[0],
2417
+ block_size=block_size,
2418
+ block_table=torch.tensor([[0, 0, 1]], dtype=torch.int32, device=device),
2419
+ slot_mapping=resident_slots,
2420
+ )
2421
+ req_ids = torch.tensor([0], dtype=torch.int32, device=device)
2422
+ newest_global = 4 * block_size
2423
+ padded_global = 1
2424
+ padded_row = flat_pool[padded_global].clone()
2425
+ slot_mapping = torch.tensor(
2426
+ [newest_global, padded_global], dtype=torch.int64, device=device
2427
+ )
2428
+ request_state_indices = torch.tensor([0, -1], dtype=torch.int32, device=device)
2429
+ cache_handle.runtime.request_state_indices = request_state_indices
2430
+
2431
+ kv_c = torch.randn(3, row_width - 2, device=device)
2432
+ k_pe = torch.randn(3, 1, 2, device=device)
2433
+ cache_handle.runtime.bind_source_cache(kv_pool)
2434
+ cache_handle.num_actual_tokens = slot_mapping.numel()
2435
+ ops.concat_and_cache_mla(
2436
+ kv_c[:1],
2437
+ k_pe[:1, 0],
2438
+ cache_handle.view.cache,
2439
+ resident_slots[:1],
2440
+ kv_cache_dtype="auto",
2441
+ scale=torch.tensor(1.0, device=device),
2442
+ )
2443
+ torch.accelerator.synchronize()
2444
+ expected_row = torch.cat([kv_c[0], k_pe[0, 0]]).cpu()
2445
+ flat_pool[newest_global].copy_(expected_row)
2446
+ torch.testing.assert_close(flat_pool[newest_global], expected_row)
2447
+ torch.testing.assert_close(flat_pool[padded_global], padded_row)
2448
+ flat_hot = cache_handle.runtime.hot.cache.view(-1, row_width)
2449
+ torch.testing.assert_close(flat_hot[resident_slots[0]].cpu(), expected_row)
2450
+
2451
+ topk = torch.tensor([[0, -1, -1, -1]], dtype=torch.int32, device=device)
2452
+ cache_handle.all_context_pages_resident = False
2453
+ cache_handle.runtime.begin_forward()
2454
+ hot_indices = cache_handle.swap_in(
2455
+ req_id_per_token=req_ids,
2456
+ block_table=block_table,
2457
+ logical_topk_indices=topk,
2458
+ block_size=block_size,
2459
+ )
2460
+ torch.accelerator.synchronize()
2461
+ stale_hot_slot = hot_indices[0, 0].item()
2462
+ stale_row = flat_hot[stale_hot_slot].clone()
2463
+
2464
+ flat_pool[8] += 1000
2465
+ cache_handle.runtime.invalidate_written_slots(
2466
+ torch.tensor([8], dtype=torch.int64, device=device), req_ids
2467
+ )
2468
+ cache_handle.runtime.begin_forward()
2469
+ hot_indices = cache_handle.swap_in(
2470
+ req_id_per_token=req_ids,
2471
+ block_table=block_table,
2472
+ logical_topk_indices=topk,
2473
+ block_size=block_size,
2474
+ )
2475
+ torch.accelerator.synchronize()
2476
+ idx = hot_indices.cpu().tolist()[0][0]
2477
+ assert not torch.equal(flat_hot[idx], stale_row)
2478
+ torch.testing.assert_close(flat_hot[idx].cpu(), flat_pool[8])
2479
+
2480
+
2481
+@requires_hisparse_ops
2482
+def test_hisparse_mixed_batch_bf16_row_split(
2483
+ default_vllm_config, dist_init, workspace_init
2484
+):
2485
+ """Host-resident mixed batch on the bf16 path is row-split.
2486
+
2487
+ Two long-context speculative-decode requests + one short local-prefill
2488
+ chunk: every decode step must be served from the bounded hot buffer before
2489
+ it is reused, while only the prefill rows' blocks are staged host->GPU.
2490
+ """
2491
+ ok, reason = flashmla.is_flashmla_sparse_supported()
2492
+ if not ok:
2493
+ pytest.skip(reason)
2494
+
2495
+ device = torch.device(DEVICE_TYPE)
2496
+ dtype = torch.bfloat16
2497
+ torch.manual_seed(0)
2498
+
2499
+ num_heads = 64
2500
+ kv_lora_rank = 512
2501
+ qk_nope_head_dim = 128
2502
+ qk_rope_head_dim = 64
2503
+ v_head_dim = 128
2504
+ head_size = kv_lora_rank + qk_rope_head_dim
2505
+ topk_tokens = 128
2506
+ block_size = 64
2507
+
2508
+ # Long decode contexts + a short prefill chunk (router shortcut shape).
2509
+ batch_spec = BatchSpec(seq_lens=[2048, 2048, 192], query_lens=[2, 2, 64])
2510
+ max_seqlen = max(batch_spec.seq_lens)
2511
+ total_cache_tokens = sum(batch_spec.seq_lens)
2512
+ total_tokens = batch_spec.compute_num_tokens()
2513
+
2514
+ vllm_config = create_vllm_config(
2515
+ model_name="deepseek-ai/DeepSeek-V2-Lite-Chat",
2516
+ tensor_parallel_size=1,
2517
+ max_model_len=max_seqlen,
2518
+ num_gpu_blocks=max(2048, cdiv(total_cache_tokens, block_size) + 1),
2519
+ block_size=block_size,
2520
+ hf_config_override={
2521
+ "index_topk": topk_tokens,
2522
+ "attn_module_list_cfg": [{"topk_tokens": topk_tokens}],
2523
+ },
2524
+ )
2525
+ vllm_config.attention_config.hisparse_config = HiSparseConfig(
2526
+ device_buffer_size=2 * topk_tokens,
2527
+ )
2528
+ vllm_config.speculative_config = SpeculativeConfig(
2529
+ method="ngram", num_speculative_tokens=1
2530
+ )
2531
+ model_config = vllm_config.model_config
2532
+ model_config.hf_text_config = SimpleNamespace(
2533
+ index_topk=topk_tokens,
2534
+ q_lora_rank=None,
2535
+ kv_lora_rank=kv_lora_rank,
2536
+ qk_nope_head_dim=qk_nope_head_dim,
2537
+ qk_rope_head_dim=qk_rope_head_dim,
2538
+ v_head_dim=v_head_dim,
2539
+ model_type="deepseek_v2",
2540
+ )
2541
+ model_config.dtype = dtype
2542
+ model_config.get_num_attention_heads = MethodType(
2543
+ lambda self, parallel_config: num_heads, model_config
2544
+ )
2545
+ model_config.get_num_kv_heads = MethodType(
2546
+ lambda self, parallel_config: 1, model_config
2547
+ )
2548
+ model_config.get_head_size = MethodType(lambda self: head_size, model_config)
2549
+ model_config.get_sliding_window = MethodType(lambda self: None, model_config)
2550
+ # create_hisparse_cache_handle sizes hot buffers per layer.
2551
+ model_config.get_num_layers = MethodType(
2552
+ lambda self, parallel_config: 1, model_config
2553
+ )
2554
+
2555
+ kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
2556
+ common_attn_metadata = create_common_attn_metadata(
2557
+ batch_spec,
2558
+ vllm_config.cache_config.block_size,
2559
+ device,
2560
+ arange_block_indices=True,
2561
+ )
2562
+ # Prepopulate every position of every sequence so the forward needs no
2563
+ # KV-cache update (the row split itself is what is under test).
2564
+ kv_c_contexts = [
2565
+ torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device)
2566
+ for s_len in batch_spec.seq_lens
2567
+ ]
2568
+ k_pe_contexts = [
2569
+ torch.rand(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device)
2570
+ for s_len in batch_spec.seq_lens
2571
+ ]
2572
+ kv_cache = create_and_prepopulate_kv_cache(
2573
+ kv_c_contexts=kv_c_contexts,
2574
+ k_pe_contexts=k_pe_contexts,
2575
+ block_size=vllm_config.cache_config.block_size,
2576
+ head_size=head_size,
2577
+ dtype=dtype,
2578
+ device=device,
2579
+ num_blocks=vllm_config.cache_config.num_gpu_blocks,
2580
+ common_attn_metadata=common_attn_metadata,
2581
+ randomize_blocks=False,
2582
+ kv_cache_dtype="auto",
2583
+ )
2584
+
2585
+ prefill_backend = get_mla_prefill_backend(vllm_config)(
2586
+ num_heads=num_heads,
2587
+ scale=1.0 / math.sqrt(head_size),
2588
+ kv_lora_rank=kv_lora_rank,
2589
+ qk_nope_head_dim=qk_nope_head_dim,
2590
+ qk_rope_head_dim=qk_rope_head_dim,
2591
+ v_head_dim=v_head_dim,
2592
+ vllm_config=vllm_config,
2593
+ )
2594
+ vllm_config.compilation_config.static_forward_context["placeholder"] = (
2595
+ SimpleNamespace(prefill_backend=prefill_backend)
2596
+ )
2597
+
2598
+ builder_cls = FlashMLASparseBackend.get_builder_cls()
2599
+ builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device)
2600
+ metadata = builder.build(
2601
+ common_prefix_len=0, common_attn_metadata=common_attn_metadata
2602
+ )
2603
+ assert isinstance(metadata.prefill, SparseMLAPrefillMetadata)
2604
+
2605
+ # Per-token sparse indices bounded by each token's position, with unique
2606
+ # offsets and -1 padding (same construction as the decode parity test).
2607
+ positions: list[int] = []
2608
+ for i in range(batch_spec.batch_size):
2609
+ ctx_len = batch_spec.seq_lens[i] - batch_spec.query_lens[i]
2610
+ positions.extend(ctx_len + q_idx for q_idx in range(batch_spec.query_lens[i]))
2611
+ sparse_indices = torch.empty(
2612
+ total_tokens, topk_tokens, dtype=torch.int32, device=device
2613
+ )
2614
+ for tok_idx in range(total_tokens):
2615
+ max_valid_idx = positions[tok_idx]
2616
+ offset = tok_idx * 7 # Prime number for varied offsets
2617
+ num_valid = min(topk_tokens // 2, max_valid_idx + 1)
2618
+ valid_range = torch.arange(num_valid, device=device, dtype=torch.int32)
2619
+ sparse_indices[tok_idx] = torch.cat(
2620
+ [
2621
+ (valid_range + offset) % (max_valid_idx + 1),
2622
+ torch.full(
2623
+ (topk_tokens - num_valid,), -1, device=device, dtype=torch.int32
2624
+ ),
2625
+ ]
2626
+ )
2627
+
2628
+ q = torch.rand(total_tokens, num_heads, head_size, dtype=dtype, device=device)
2629
+ mock_indexer = SimpleNamespace(
2630
+ topk_indices_buffer=sparse_indices, topk_tokens=topk_tokens
2631
+ )
2632
+
2633
+ impl_cls = FlashMLASparseBackend.get_impl_cls()
2634
+ with set_current_vllm_config(vllm_config):
2635
+ impl = impl_cls(
2636
+ num_heads=num_heads,
2637
+ head_size=head_size,
2638
+ scale=1.0 / math.sqrt(head_size),
2639
+ num_kv_heads=1,
2640
+ alibi_slopes=None,
2641
+ sliding_window=None,
2642
+ kv_cache_dtype="auto",
2643
+ logits_soft_cap=None,
2644
+ attn_type="decoder",
2645
+ kv_sharing_target_layer_name=None,
2646
+ q_lora_rank=None,
2647
+ kv_lora_rank=kv_lora_rank,
2648
+ qk_nope_head_dim=qk_nope_head_dim,
2649
+ qk_rope_head_dim=qk_rope_head_dim,
2650
+ qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,
2651
+ v_head_dim=v_head_dim,
2652
+ kv_b_proj=None,
2653
+ indexer=mock_indexer,
2654
+ )
2655
+ assert isinstance(impl.index_group, HiSparseMLAIndexGroup)
2656
+ cache_handle = impl.index_group.cache(impl.index_group_index)
2657
+ blocks_per_request = cdiv(cache_handle.runtime.region_stride, block_size)
2658
+ num_hot_blocks = cache_handle.runtime.max_num_reqs * blocks_per_request
2659
+ hot_cache = torch.zeros(
2660
+ num_hot_blocks * block_size * head_size,
2661
+ dtype=dtype,
2662
+ device=device,
2663
+ ).view(torch.int8)
2664
+ cache_handle.runtime.bind_hot_cache(
2665
+ hot_cache,
2666
+ byte_offset=0,
2667
+ block_stride=block_size * head_size * dtype.itemsize,
2668
+ num_blocks=num_hot_blocks,
2669
+ block_size=block_size,
2670
+ block_table=torch.arange(num_hot_blocks, dtype=torch.int32, device=device).view(
2671
+ cache_handle.runtime.max_num_reqs, blocks_per_request
2672
+ ),
2673
+ )
2674
+ cache_handle.runtime.request_state_indices = torch.arange(
2675
+ cache_handle.runtime.max_num_reqs, dtype=torch.int32, device=device
2676
+ )
2677
+ impl.prepare_for_batch(metadata)
2678
+
2679
+ # Device-resident reference: the whole batch converted against the full
2680
+ # block table and run as one kernel call over the GPU cache.
2681
+ ref_topk, ref_topk_length = triton_convert_req_index_to_global_index(
2682
+ metadata.req_id_per_token,
2683
+ metadata.block_table,
2684
+ sparse_indices,
2685
+ BLOCK_SIZE=metadata.block_size,
2686
+ NUM_TOPK_TOKENS=topk_tokens,
2687
+ return_valid_counts=True,
2688
+ )
2689
+ reference, _ = impl._bf16_flash_mla_kernel(q, kv_cache, ref_topk, ref_topk_length)
2690
+
2691
+ # Host-resident pool with identical contents.
2692
+ kv_pool = kv_cache.squeeze(1).cpu().pin_memory()
2693
+ cache_handle.runtime.bind_source_cache(kv_pool)
2694
+ cache_handle.source_block_table = metadata.block_table
2695
+
2696
+ staging_calls = []
2697
+ original_gather = cache_handle.runtime.gather_prefill_cache
2698
+
2699
+ def spy_gather(self, kv, plan, resident_cache=None):
2700
+ staged = original_gather(kv, plan, resident_cache)
2701
+ staging_calls.append((plan, staged.shape))
2702
+ return staged
2703
+
2704
+ cache_handle.runtime.gather_prefill_cache = MethodType(
2705
+ spy_gather, cache_handle.runtime
2706
+ )
2707
+
2708
+ backend_output, _ = impl._forward_bf16_kv(
2709
+ q, kv_pool, sparse_indices, metadata, q.shape[1]
2710
+ )
2711
+ torch.accelerator.synchronize()
2712
+
2713
+ # Only the prefill rows' blocks were staged: the decode rows' 2048-token
2714
+ # contexts (32 blocks each) must stay off the staging gather.
2715
+ assert len(staging_calls) == 1
2716
+ plan, staged_shape = staging_calls[0]
2717
+ assert plan is metadata.prefill.host_staging_plan
2718
+ prefill_blocks = cdiv(batch_spec.seq_lens[-1], block_size)
2719
+ assert staged_shape[0] <= prefill_blocks + 1 # +1: block-0 tail padding
2720
+
2721
+ torch.testing.assert_close(backend_output, reference, rtol=0.01, atol=0.01)
2722
+
2723
+
2724
+def test_hisparse_prefill_staging_remap():
2725
+ """Compacted staging references the same host rows as direct indexing."""
2726
+ block_size = 4
2727
+ block_table = torch.tensor(
2728
+ [[5, 2, -1, -1], [2, 7, 3, -1], [9, -1, -1, -1]], dtype=torch.int32
2729
+ )
2730
+
2731
+ new_bt, row_ids = hisparse_prefill_staging_remap(block_table, block_size)
2732
+
2733
+ assert new_bt.shape == block_table.shape
2734
+ assert row_ids.dtype == torch.int32 and new_bt.dtype == torch.int32
2735
+ n_unique = len({0, 2, 3, 5, 7, 9}) # -1 clamps to block 0
2736
+ flat_rows = row_ids.flatten()
2737
+ assert (flat_rows >= 0).sum() == n_unique * block_size
2738
+ for i in range(block_table.shape[0]):
2739
+ for j in range(block_table.shape[1]):
2740
+ orig = max(int(block_table[i, j]), 0)
2741
+ staged = int(new_bt[i, j])
2742
+ for k in range(block_size):
2743
+ assert int(flat_rows[staged * block_size + k]) == orig * block_size + k
2744
+
2745
+
2746
+def test_hisparse_prefill_staging_plan_masks_unused_blocks():
2747
+ block_table = torch.tensor([[5, 2, 7], [9, 3, 4]], dtype=torch.int32)
2748
+ plan = build_hisparse_prefill_staging_plan(
2749
+ block_table,
2750
+ seq_lens=torch.tensor([5, 8], dtype=torch.int32),
2751
+ block_size=4,
2752
+ staging_block_capacity=4,
2753
+ )
2754
+
2755
+ assert (plan.block_table[:, 2] == 0).all()
2756
+ valid_rows = plan.row_ids[plan.row_ids >= 0]
2757
+ assert set((valid_rows[::4] // 4).tolist()) == {0, 2, 3, 5, 9}
2758
+ for row, used in enumerate((2, 2)):
2759
+ for column in range(used):
2760
+ staged_block = int(plan.block_table[row, column])
2761
+ host_block = int(block_table[row, column])
2762
+ assert int(valid_rows[staged_block * 4]) == host_block * 4
2763
+
2764
+
2765
+def test_hisparse_prefill_staging_plan_resolves_resident_sources():
2766
+ """Per-page resident hits become device sources; misses stay host DMAs."""
2767
+ block_size = 4
2768
+ resident_block_size = 2
2769
+ block_table = torch.tensor([[5, 2, 0], [9, 3, 0]], dtype=torch.int32)
2770
+ seq_lens = torch.tensor([5, 8], dtype=torch.int32)
2771
+ plan = build_hisparse_prefill_staging_plan(
2772
+ block_table, seq_lens, block_size, staging_block_capacity=4
2773
+ )
2774
+ # Two resident pages per host block; 0 entries are null (not resident).
2775
+ resident_table = torch.tensor(
2776
+ [[11, 12, 0, 13, 0, 0], [21, 0, 22, 23, 0, 0]], dtype=torch.int32
2777
+ )
2778
+
2779
+ plan.ensure_gpu_sources(resident_table, resident_block_size)
2780
+
2781
+ assert plan.gpu_row_ids is not None
2782
+ unique_hosts = (plan.row_ids[0].view(-1, block_size)[:, 0] // block_size).tolist()
2783
+ gpu_rows = plan.gpu_row_ids[0].view(-1, block_size)
2784
+ reps = {5: (0, 0), 2: (0, 1), 9: (1, 0), 3: (1, 1)}
2785
+ for u, host_id in enumerate(unique_hosts):
2786
+ for t in range(block_size):
2787
+ if host_id < 0:
2788
+ assert int(gpu_rows[u, t]) == -1
2789
+ continue
2790
+ if host_id == 0:
2791
+ assert int(gpu_rows[u, t]) == -1
2792
+ continue
2793
+ row, col = reps[host_id]
2794
+ page = t // resident_block_size
2795
+ res_block = int(resident_table[row, col * 2 + page])
2796
+ expected = (
2797
+ res_block * resident_block_size + t % resident_block_size
2798
+ if res_block > 0
2799
+ else -1
2800
+ )
2801
+ assert int(gpu_rows[u, t]) == expected
2802
+ valid_rows = plan.row_ids >= 0
2803
+ torch.testing.assert_close(
2804
+ plan.miss_mask,
2805
+ ((plan.gpu_row_ids < 0) & valid_rows).int(),
2806
+ check_dtype=False,
2807
+ )
2808
+
2809
+ plan.ensure_gpu_sources(torch.zeros_like(resident_table), resident_block_size)
2810
+
2811
+ assert plan.gpu_row_ids is not None
2812
+ assert (plan.gpu_row_ids == -1).all()
2813
+
2814
+
2815
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
2816
+def test_hisparse_gather_prefill_cache_prefers_resident_rows():
2817
+ """Staged rows come from the resident cache when a shadow page exists."""
2818
+ if not _has_hisparse_ops():
2819
+ pytest.skip("hisparse CUDA ops unavailable")
2820
+ device = torch.device("cuda")
2821
+ block_size, resident_block_size, row_width = 4, 2, 16
2822
+ block_table = torch.tensor([[5, 2, 0], [9, 3, 0]], dtype=torch.int32, device=device)
2823
+ seq_lens = torch.tensor([5, 8], dtype=torch.int32, device=device)
2824
+ plan = build_hisparse_prefill_staging_plan(
2825
+ block_table, seq_lens, block_size, staging_block_capacity=4
2826
+ )
2827
+ resident_table = torch.tensor(
2828
+ [[11, 12, 0, 13, 0, 0], [21, 0, 22, 23, 0, 0]],
2829
+ dtype=torch.int32,
2830
+ device=device,
2831
+ )
2832
+ plan.ensure_gpu_sources(resident_table, resident_block_size)
2833
+
2834
+ num_host_blocks, num_res_blocks = 10, 24
2835
+ host_cache = (
2836
+ torch.arange(num_host_blocks * block_size * row_width, dtype=torch.float32)
2837
+ .view(num_host_blocks, block_size, row_width)
2838
+ .pin_memory()
2839
+ )
2840
+ resident_cache = (
2841
+ (
2842
+ -torch.arange(
2843
+ num_res_blocks * resident_block_size * row_width, dtype=torch.float32
2844
+ )
2845
+ - 1.0
2846
+ )
2847
+ .view(num_res_blocks, resident_block_size, row_width)
2848
+ .to(device)
2849
+ )
2850
+
2851
+ staged = HiSparseRuntime.gather_prefill_cache(
2852
+ None, host_cache, plan, resident_cache=resident_cache
2853
+ )
2854
+
2855
+ staged_flat = staged.view(-1, row_width).cpu()
2856
+ host_flat = host_cache.view(-1, row_width)
2857
+ resident_flat = resident_cache.reshape(-1, row_width).cpu()
2858
+ gpu_rows = plan.gpu_row_ids[0].cpu()
2859
+ host_rows = plan.row_ids[0].cpu()
2860
+ for i in range(staged_flat.shape[0]):
2861
+ if int(host_rows[i]) < 0:
2862
+ assert int(plan.miss_mask[0, i]) == 0
2863
+ continue
2864
+ expected = (
2865
+ resident_flat[int(gpu_rows[i])]
2866
+ if int(gpu_rows[i]) >= 0
2867
+ else host_flat[int(host_rows[i])]
2868
+ )
2869
+ torch.testing.assert_close(staged_flat[i], expected)
2870
+
2871
+
2872
+def test_hisparse_fp8_decode_resolves_steps_then_runs_batched_attention(monkeypatch):
2873
+ """FP8 decode must use active dimensions when common metadata is padded."""
2874
+ device = torch.device("cpu")
2875
+ num_decodes = 2
2876
+ query_len = 3
2877
+ num_tokens = num_decodes * query_len
2878
+ q = torch.randn(num_tokens, 2, 4, device=device)
2879
+ topk = (
2880
+ torch.arange(num_tokens, dtype=torch.int32, device=device)
2881
+ .view(-1, 1)
2882
+ .expand(-1, 4)
2883
+ )
2884
+ steps: list[int] = []
2885
+ kernel_shapes: list[torch.Size] = []
2886
+
2887
+ def swap_in(req_ids, *, logical_topk_indices, **kwargs): # noqa: ARG001
2888
+ steps.append(int(logical_topk_indices[0, 0]))
2889
+ output = kwargs["attention_indices_out"]
2890
+ if output is not None:
2891
+ output.copy_(logical_topk_indices + 10)
2892
+ return topk[:num_decodes]
2893
+
2894
+ def run_kernel(self, *, q, **kwargs): # noqa: ARG001
2895
+ kernel_shapes.append(q.shape)
2896
+ return q[..., :1], None
2897
+
2898
+ runtime = SimpleNamespace(
2899
+ hot=SimpleNamespace(attention_cache=torch.empty(1, device=device))
2900
+ )
2901
+ leader_cache = SimpleNamespace(
2902
+ runtime=runtime,
2903
+ source_block_table=torch.empty(
2904
+ num_decodes, 1, dtype=torch.int32, device=device
2905
+ ),
2906
+ swap_in=MagicMock(side_effect=swap_in),
2907
+ )
2908
+ follower_cache = SimpleNamespace(
2909
+ runtime=runtime,
2910
+ source_block_table=torch.empty(
2911
+ num_decodes, 1, dtype=torch.int32, device=device
2912
+ ),
2913
+ swap_in=MagicMock(side_effect=swap_in),
2914
+ )
2915
+ index_group = object.__new__(HiSparseMLAIndexGroup)
2916
+ index_group.caches = [leader_cache, follower_cache]
2917
+ index_group.physical_topk_indices = torch.empty(
2918
+ (num_tokens + 1, 4), dtype=torch.int32, device=device
2919
+ )
2920
+ index_group.valid_topk_counts = torch.empty(
2921
+ num_tokens + 1, dtype=torch.int32, device=device
2922
+ )
2923
+ index_group.request_ids = torch.arange(
2924
+ num_decodes, dtype=torch.int32, device=device
2925
+ )
2926
+ impl = SimpleNamespace(
2927
+ kv_lora_rank=1,
2928
+ index_group=index_group,
2929
+ index_group_index=0,
2930
+ )
2931
+ impl._fp8_flash_mla_kernel = MethodType(run_kernel, impl)
2932
+ metadata = SimpleNamespace(
2933
+ num_decodes=num_tokens,
2934
+ num_decode_tokens=num_tokens,
2935
+ decode_max_query_len=1,
2936
+ query_start_loc=torch.arange(
2937
+ 0,
2938
+ num_tokens + 1,
2939
+ query_len,
2940
+ dtype=torch.int32,
2941
+ device=device,
2942
+ ),
2943
+ block_table=torch.empty(num_decodes, 1, dtype=torch.int32, device=device),
2944
+ block_size=64,
2945
+ )
2946
+
2947
+ output = FlashMLASparseImpl._host_backed_fp8_decode(
2948
+ impl,
2949
+ q,
2950
+ topk,
2951
+ metadata,
2952
+ SimpleNamespace(),
2953
+ num_decodes,
2954
+ query_len,
2955
+ )
2956
+
2957
+ assert steps == [0, 1, 2]
2958
+ assert kernel_shapes == [(num_decodes, query_len, 2, 4)]
2959
+ assert output.shape == (num_tokens, 2, 1)
2960
+ torch.testing.assert_close(
2961
+ index_group.physical_topk_indices[:num_tokens], topk + 10
2962
+ )
2963
+
2964
+ follower_indices = index_group.convert_decode_logical_to_physical_topk(
2965
+ 1,
2966
+ topk,
2967
+ metadata,
2968
+ return_valid_counts=False,
2969
+ num_decodes=num_decodes,
2970
+ decode_query_len=query_len,
2971
+ )
2972
+ assert steps == [0, 1, 2, 0, 1, 2]
2973
+ torch.testing.assert_close(follower_indices, topk + 10)
2974
+
2975
+
2976
+def test_hisparse_varlen_decode_uses_canonical_request_rows():
2977
+ """Graph metadata's query upper bound must preserve packed token rows."""
2978
+ num_decodes = 2
2979
+ topk = torch.arange(18, dtype=torch.int32).view(9, 2)
2980
+
2981
+ def swap_in(_request_ids, *, logical_topk_indices, **_kwargs):
2982
+ return logical_topk_indices + 10
2983
+
2984
+ cache = SimpleNamespace(
2985
+ source_block_table=torch.zeros((num_decodes, 1), dtype=torch.int32),
2986
+ swap_in=MagicMock(side_effect=swap_in),
2987
+ )
2988
+ index_group = object.__new__(HiSparseMLAIndexGroup)
2989
+ index_group.caches = [cache]
2990
+ index_group.physical_topk_indices = torch.empty((10, 2), dtype=torch.int32)
2991
+ index_group.valid_topk_counts = torch.empty(10, dtype=torch.int32)
2992
+ index_group.request_ids = torch.arange(num_decodes, dtype=torch.int32)
2993
+ metadata = SimpleNamespace(
2994
+ num_decodes=num_decodes,
2995
+ num_decode_tokens=topk.shape[0],
2996
+ decode_max_query_len=15,
2997
+ query_start_loc=torch.tensor([0, 8, 9], dtype=torch.int32),
2998
+ block_size=64,
2999
+ )
3000
+
3001
+ physical_topk = index_group.convert_decode_logical_to_physical_topk(
3002
+ 0,
3003
+ topk,
3004
+ metadata,
3005
+ return_valid_counts=False,
3006
+ )
3007
+
3008
+ torch.testing.assert_close(physical_topk, topk + 10)
3009
+ request_rows = cache.swap_in.call_args_list[0].args[0]
3010
+ torch.testing.assert_close(request_rows, index_group.request_ids)
3011
+
3012
+
3013
+def test_flashinfer_hisparse_decode_runs_batched_attention():
3014
+ device = torch.device("cpu")
3015
+ num_tokens = 6
3016
+ q = torch.randn(num_tokens, 2, 4, device=device)
3017
+ topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=device)
3018
+ physical_topk = topk + 10
3019
+ valid_counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=device)
3020
+ kernel_shapes: list[torch.Size] = []
3021
+
3022
+ def convert_decode(self, *args, **kwargs): # noqa: ARG001
3023
+ return physical_topk, valid_counts
3024
+
3025
+ def prepare_kernel(self, *args, **kwargs): # noqa: ARG001
3026
+ pass
3027
+
3028
+ def run_kernel(self, q, cache, indices, counts): # noqa: ARG001
3029
+ kernel_shapes.append(q.shape)
3030
+ return q[..., :1], None
3031
+
3032
+ cache_handle = SimpleNamespace(
3033
+ runtime=SimpleNamespace(
3034
+ hot=SimpleNamespace(attention_cache=torch.empty(1, device=device))
3035
+ ),
3036
+ )
3037
+ index_group = object.__new__(HiSparseMLAIndexGroup)
3038
+ index_group.caches = [cache_handle]
3039
+ index_group.convert_decode_logical_to_physical_topk = MethodType(
3040
+ convert_decode, index_group
3041
+ )
3042
+ impl = object.__new__(FlashInferMLASparseImpl)
3043
+ impl.topk_indices_buffer = topk
3044
+ impl.index_group = index_group
3045
+ impl.index_group_index = 0
3046
+ impl._prepare_mqa_kernel = MethodType(prepare_kernel, impl)
3047
+ impl._run_mqa_kernel = MethodType(run_kernel, impl)
3048
+ metadata = SimpleNamespace(num_decode_tokens=num_tokens)
3049
+
3050
+ output, lse = FlashInferMLASparseImpl.forward_mqa(
3051
+ impl,
3052
+ q,
3053
+ torch.empty(1, device=device),
3054
+ metadata,
3055
+ SimpleNamespace(),
3056
+ )
3057
+
3058
+ assert kernel_shapes == [q.shape]
3059
+ assert output.shape == (num_tokens, 2, 1)
3060
+ assert lse is None
3061
+
3062
+
3063
+def test_flashattn_hisparse_decode_uses_index_group():
3064
+ num_tokens = 4
3065
+ q_nope = torch.empty(num_tokens, 2, 3, device=DEVICE_TYPE)
3066
+ q_rope = torch.empty(num_tokens, 2, 1, device=DEVICE_TYPE)
3067
+ topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=DEVICE_TYPE)
3068
+ physical = topk.clone()
3069
+ counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE)
3070
+ index_group = object.__new__(HiSparseMLAIndexGroup)
3071
+ index_group.convert_decode_logical_to_physical_topk = MagicMock(
3072
+ return_value=(physical, counts)
3073
+ )
3074
+ index_group.physical_kv_cache = MagicMock(
3075
+ return_value=torch.empty(1, device=DEVICE_TYPE)
3076
+ )
3077
+ impl = object.__new__(FlashAttnMLASparseImpl)
3078
+ impl.topk_indices_buffer = topk
3079
+ impl.index_group = index_group
3080
+ impl.index_group_index = 0
3081
+ impl._run_mqa_kernel = MagicMock(return_value=q_nope[..., :1])
3082
+ metadata = SimpleNamespace(num_decode_tokens=num_tokens, block_size=64)
3083
+
3084
+ output, lse = FlashAttnMLASparseImpl.forward_mqa(
3085
+ impl,
3086
+ (q_nope, q_rope),
3087
+ torch.empty(1, device=DEVICE_TYPE),
3088
+ metadata,
3089
+ SimpleNamespace(),
3090
+ )
3091
+
3092
+ assert output.shape == (num_tokens, 2, 1)
3093
+ assert lse is None
3094
+ index_group.convert_decode_logical_to_physical_topk.assert_called_once()
3095
+ impl._run_mqa_kernel.assert_called_once()
3096
+
3097
+
3098
+def test_flashinfer_sm120_hisparse_decode_uses_index_group():
3099
+ num_tokens = 4
3100
+ q = torch.empty(num_tokens, 2, 4, device=DEVICE_TYPE)
3101
+ topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=DEVICE_TYPE)
3102
+ physical = topk.clone()
3103
+ index_group = object.__new__(HiSparseMLAIndexGroup)
3104
+ index_group.convert_decode_logical_to_physical_topk = MagicMock(
3105
+ return_value=physical
3106
+ )
3107
+ index_group.physical_kv_cache = MagicMock(
3108
+ return_value=torch.empty(1, device=DEVICE_TYPE)
3109
+ )
3110
+ impl = object.__new__(FlashInferMLASparseSM120Impl)
3111
+ impl.topk_indices_buffer = topk
3112
+ impl.index_group = index_group
3113
+ impl.index_group_index = 0
3114
+ impl._run_mqa_kernel = MagicMock(return_value=q[..., :1])
3115
+ metadata = SimpleNamespace(
3116
+ num_decode_tokens=num_tokens,
3117
+ topk_tokens=topk.shape[1],
3118
+ )
3119
+
3120
+ output, lse = FlashInferMLASparseSM120Impl.forward_mqa(
3121
+ impl,
3122
+ q,
3123
+ torch.empty(1, device=DEVICE_TYPE),
3124
+ metadata,
3125
+ SimpleNamespace(),
3126
+ )
3127
+
3128
+ assert output.shape == (num_tokens, 2, 1)
3129
+ assert lse is None
3130
+ index_group.convert_decode_logical_to_physical_topk.assert_called_once()
3131
+ impl._run_mqa_kernel.assert_called_once()
3132
+
3133
+
3134
+def test_hisparse_resident_prefill_uses_attention_block_stride():
3135
+ expected = torch.tensor([[19]], dtype=torch.int32)
3136
+ cache_handle = SimpleNamespace(
3137
+ all_context_pages_resident=True,
3138
+ view=SimpleNamespace(block_size=64, attention_block_stride=832),
3139
+ block_table=torch.tensor([[3]], dtype=torch.int32),
3140
+ )
3141
+ index_group = object.__new__(HiSparseMLAIndexGroup)
3142
+ index_group.caches = [cache_handle]
3143
+ index_group.physical_topk_indices = torch.empty((1, 1), dtype=torch.int32)
3144
+ index_group._convert_once = MagicMock(return_value=expected)
3145
+ topk = torch.zeros((2, 1), dtype=torch.int32)
3146
+ metadata = SimpleNamespace(
3147
+ req_id_per_token=torch.zeros(2, dtype=torch.int32),
3148
+ )
3149
+
3150
+ result = index_group.convert_logical_to_physical_topk(
3151
+ 0,
3152
+ topk,
3153
+ metadata,
3154
+ block_stride_rows=None,
3155
+ return_valid_counts=False,
3156
+ )
3157
+
3158
+ assert result is expected
3159
+ assert index_group._convert_once.call_args.kwargs["block_stride_rows"] == 832
3160
+
3161
+
3162
+def test_hisparse_shared_sparse_builder_routes_multi_token_chunks_to_prefill():
3163
+ """Default sparse backends must route multi-token chunks to prefill."""
3164
+ builder = object.__new__(FlashInferMLASparseMetadataBuilder)
3165
+ builder.vllm_config = SimpleNamespace(
3166
+ attention_config=SimpleNamespace(hisparse_config=object()),
3167
+ speculative_config=SimpleNamespace(
3168
+ num_speculative_tokens=4,
3169
+ parallel_drafting=False,
3170
+ ),
3171
+ parallel_config=SimpleNamespace(decode_context_parallel_size=1),
3172
+ )
3173
+
3174
+ builder._init_reorder_batch_threshold(
3175
+ 1024,
3176
+ supports_spec_as_decode=True,
3177
+ supports_dcp_with_varlen=True,
3178
+ )
3179
+
3180
+ assert builder.reorder_batch_threshold == 1
3181
+
3182
+
3183
+def test_hisparse_flashmla_reorders_full_speculative_window_as_decode():
3184
+ builder = object.__new__(FlashMLASparseMetadataBuilder)
3185
+ builder.vllm_config = SimpleNamespace(
3186
+ attention_config=SimpleNamespace(hisparse_config=object()),
3187
+ speculative_config=SimpleNamespace(
3188
+ num_speculative_tokens=3,
3189
+ parallel_drafting=False,
3190
+ ),
3191
+ parallel_config=SimpleNamespace(decode_context_parallel_size=1),
3192
+ )
3193
+
3194
+ builder._init_reorder_batch_threshold(
3195
+ 256,
3196
+ supports_spec_as_decode=True,
3197
+ supports_dcp_with_varlen=True,
3198
+ )
3199
+
3200
+ assert builder.reorder_batch_threshold == 4
3201
+
3202
+
3203
+def test_hisparse_mixed_single_token_batch_is_not_decode_only():
3204
+ cache = SimpleNamespace(
3205
+ runtime=SimpleNamespace(
3206
+ begin_forward=lambda: None,
3207
+ eager_host_mirror=True,
3208
+ ),
3209
+ )
3210
+ metadata = SimpleNamespace(
3211
+ num_decode_tokens=1,
3212
+ num_actual_tokens=2,
3213
+ max_query_len=1,
3214
+ num_reqs=2,
3215
+ req_id_per_token=torch.arange(2),
3216
+ )
3217
+
3218
+ HiSparseCacheHandle._prepare_for_batch(cache, metadata)
3219
+
3220
+ assert not cache.decode_batch
3221
+
3222
+
3223
+def test_flashmla_cache_dtype_aliases_use_ds_layout():
3224
+ # kv-cache dtype aliases are canonicalized to fp8_ds_mla before the layer
3225
+ # stores kv_cache_dtype, so they cannot bypass the gate.
3226
+ for alias in ("fp8", "fp8_e4m3"):
3227
+ assert (
3228
+ _canonicalize_sparse_mla_kv_cache_dtype(FlashMLASparseBackend, alias)
3229
+ == "fp8_ds_mla"
3230
+ )
3231
+
3232
+
3233
+def test_flashmla_common_metadata_requires_uniform_decodes():
3234
+ common_metadata = SimpleNamespace(
3235
+ max_query_len=3,
3236
+ num_reqs=3,
3237
+ num_actual_tokens=6,
3238
+ query_start_loc_cpu=torch.tensor([0, 1, 3, 6]),
3239
+ is_prefilling=None,
3240
+ )
3241
+
3242
+ split = split_decodes_and_prefills(
3243
+ common_metadata,
3244
+ decode_threshold=128,
3245
+ require_uniform=FlashMLASparseMetadataBuilder.require_uniform_decodes,
3246
+ )
3247
+
3248
+ assert split == (1, 2, 1, 5)
3249
+
3250
+
3251
+def test_flashmla_fp8_metadata_excludes_zero_token_decode_padding(monkeypatch):
3252
+ monkeypatch.setattr(
3253
+ "vllm.v1.attention.backends.mla.flashmla_sparse.get_mla_metadata",
3254
+ lambda: (object(), None),
3255
+ )
3256
+ builder = SimpleNamespace(
3257
+ device=torch.device(DEVICE_TYPE),
3258
+ dummy_block_table=torch.zeros(7, 1, device=DEVICE_TYPE),
3259
+ max_model_len_tensor=torch.zeros(7, device=DEVICE_TYPE),
3260
+ )
3261
+ query_start_loc_cpu = torch.tensor([0, 110, 220, 330, 440, 550, 660, 660])
3262
+ common_metadata = SimpleNamespace(
3263
+ num_actual_tokens=660,
3264
+ query_start_loc_cpu=query_start_loc_cpu,
3265
+ seq_lens=torch.arange(7, device=DEVICE_TYPE),
3266
+ )
3267
+ metadata = FlashMLASparseMetadata(
3268
+ num_reqs=7,
3269
+ max_query_len=110,
3270
+ max_seq_len=110,
3271
+ num_actual_tokens=660,
3272
+ query_start_loc=query_start_loc_cpu.to(DEVICE_TYPE),
3273
+ slot_mapping=torch.arange(660, device=DEVICE_TYPE),
3274
+ block_table=torch.zeros(7, 1, dtype=torch.int32, device=DEVICE_TYPE),
3275
+ req_id_per_token=torch.zeros(660, dtype=torch.int32, device=DEVICE_TYPE),
3276
+ num_decodes=7,
3277
+ num_prefills=0,
3278
+ num_decode_tokens=660,
3279
+ )
3280
+
3281
+ fp8_metadata = FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode(
3282
+ builder, common_metadata, metadata
3283
+ )
3284
+
3285
+ assert fp8_metadata.num_decodes == 6
3286
+ assert fp8_metadata.num_decode_tokens == 660
3287
+ assert fp8_metadata.decode is not None
3288
+ assert fp8_metadata.decode.decode_query_len == 110
3289
+ torch.testing.assert_close(
3290
+ fp8_metadata.decode.seq_lens, torch.arange(6, device=DEVICE_TYPE)
3291
+ )
3292
+
3293
+
3294
+@pytest.mark.parametrize("use_mixed_batch", [False, True])
3295
+def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: bool):
3296
+ num_decode_tokens = 2
3297
+ num_batch_tokens = 5
3298
+ q = torch.empty(num_decode_tokens, 2, 3, device=DEVICE_TYPE)
3299
+ topk_indices = torch.empty(num_decode_tokens, 4, device=DEVICE_TYPE)
3300
+ kernel_q_shapes = []
3301
+
3302
+ def convert_indices(*args, **kwargs): # noqa: ARG001
3303
+ assert not kwargs.get("HAS_PREFILL_WORKSPACE", False)
3304
+ if not kwargs.get("return_valid_counts", False):
3305
+ return topk_indices
3306
+ valid_counts = torch.full(
3307
+ (num_decode_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE
3308
+ )
3309
+ return topk_indices, valid_counts
3310
+
3311
+ monkeypatch.setattr(
3312
+ "vllm.v1.attention.backends.mla.flashmla_sparse."
3313
+ "triton_convert_req_index_to_global_index",
3314
+ convert_indices,
3315
+ )
3316
+
3317
+ def run_kernel(**kwargs):
3318
+ kernel_q_shapes.append(kwargs["q"].shape)
3319
+ return kwargs["q"][..., :1], None
3320
+
3321
+ if use_mixed_batch:
3322
+ fp8_metadata = FlashMLASparseMetadata.FP8KernelMetadata(
3323
+ scheduler_metadata=object(), # type: ignore[arg-type]
3324
+ dummy_block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE),
3325
+ cache_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE),
3326
+ )
3327
+ else:
3328
+ FP8Meta = FlashMLASparseMetadata.FP8SeparatePrefillDecode
3329
+ fp8_metadata = FP8Meta(
3330
+ num_decodes=1,
3331
+ num_prefills=1,
3332
+ num_decode_tokens=num_decode_tokens,
3333
+ num_prefill_tokens=num_batch_tokens - num_decode_tokens,
3334
+ decode=FP8Meta.Decode(
3335
+ seq_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE),
3336
+ kernel_metadata=object(), # type: ignore[arg-type]
3337
+ decode_query_len=num_decode_tokens,
3338
+ ),
3339
+ prefill=FP8Meta.Prefill(
3340
+ request_ids=torch.empty(
3341
+ num_batch_tokens, dtype=torch.int32, device=DEVICE_TYPE
3342
+ ),
3343
+ workspace_starts=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE),
3344
+ chunks=[],
3345
+ ),
3346
+ )
3347
+ metadata = SimpleNamespace(
3348
+ fp8_extra_metadata=fp8_metadata,
3349
+ fp8_use_mixed_batch=use_mixed_batch,
3350
+ num_actual_tokens=num_batch_tokens,
3351
+ num_decode_tokens=num_decode_tokens,
3352
+ req_id_per_token=torch.empty(
3353
+ num_batch_tokens, dtype=torch.int32, device=DEVICE_TYPE
3354
+ ),
3355
+ block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE),
3356
+ block_size=64,
3357
+ )
3358
+ impl = SimpleNamespace(
3359
+ kv_cache_dtype="fp8_ds_mla",
3360
+ topk_indices_buffer=topk_indices,
3361
+ num_heads=2,
3362
+ kv_lora_rank=1,
3363
+ index_group=None,
3364
+ index_group_index=0,
3365
+ dcp_world_size=1,
3366
+ need_to_return_lse_for_decode=False,
3367
+ _fp8_flash_mla_kernel=run_kernel,
3368
+ _convert_logical_to_physical_topk=(
3369
+ lambda indices, metadata, **kwargs: convert_indices(
3370
+ metadata.req_id_per_token[: indices.shape[0]],
3371
+ metadata.block_table,
3372
+ indices,
3373
+ return_valid_counts=kwargs["return_valid_counts"],
3374
+ )
3375
+ ),
3376
+ )
3377
+ impl._forward_fp8_kv_mixed_batch = MethodType(
3378
+ FlashMLASparseImpl._forward_fp8_kv_mixed_batch, impl
3379
+ )
3380
+ impl._forward_fp8_kv_separate_prefill_decode = MethodType(
3381
+ FlashMLASparseImpl._forward_fp8_kv_separate_prefill_decode, impl
3382
)
3383
3384
output, lse = FlashMLASparseImpl.forward_mqa(
@@ -1856,6 +3528,8 @@ def test_fp8_mixed_batch_dcp_neutralizes_empty_rows(monkeypatch):
3528
cp_kv_cache_interleave_size=1,
3529
)
3530
impl = SimpleNamespace(
3531
+ index_group=None,
3532
+ index_group_index=0,
3533
dcp_world_size=2,
3534
dcp_rank=0,
3535
need_to_return_lse_for_decode=True,
@@ -1876,6 +3550,117 @@ def test_fp8_mixed_batch_dcp_neutralizes_empty_rows(monkeypatch):
3550
assert not lse.isnan().any()
3551
3552
3553
+def test_hisparse_prefill_reuses_builder_staging_plan():
3554
+ """Every layer must reuse the batch plan instead of synchronizing to dedupe."""
3555
+ plan = SimpleNamespace(
3556
+ block_table=torch.tensor([[0]], dtype=torch.int32),
3557
+ ensure_gpu_sources=MagicMock(),
3558
+ )
3559
+ staged = torch.empty((1, 1, 8))
3560
+ calls = []
3561
+
3562
+ def gather(kv_cache, staging_plan, resident_cache=None):
3563
+ calls.append((kv_cache, staging_plan, resident_cache))
3564
+ return staged
3565
+
3566
+ resident_cache = torch.empty((1, 1, 8))
3567
+ resident_block_table = torch.tensor([[1]], dtype=torch.int32)
3568
+ cache = SimpleNamespace(
3569
+ runtime=SimpleNamespace(
3570
+ gather_prefill_cache=gather,
3571
+ ),
3572
+ view=SimpleNamespace(cache=resident_cache, block_size=1),
3573
+ block_table=resident_block_table,
3574
+ )
3575
+ index_group = object.__new__(HiSparseMLAIndexGroup)
3576
+ index_group.caches = [cache]
3577
+ source = torch.empty((1, 1, 8))
3578
+ metadata = SimpleNamespace(
3579
+ num_decodes=0,
3580
+ num_decode_tokens=0,
3581
+ seq_lens=torch.tensor([1], dtype=torch.int32),
3582
+ prefill=SimpleNamespace(host_staging_plan=plan),
3583
+ req_id_per_token=torch.tensor([0], dtype=torch.int32),
3584
+ )
3585
+
3586
+ result, block_table, request_ids = index_group.stage_prefill_rows(
3587
+ 0, source, metadata
3588
+ )
3589
+
3590
+ assert result is staged
3591
+ assert block_table is plan.block_table
3592
+ torch.testing.assert_close(request_ids, metadata.req_id_per_token)
3593
+ plan.ensure_gpu_sources.assert_called_once()
3594
+ args = plan.ensure_gpu_sources.call_args.args
3595
+ torch.testing.assert_close(args[0], resident_block_table)
3596
+ assert args[1] == 1
3597
+ assert calls == [(source, plan, resident_cache)]
3598
+
3599
+
3600
+def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch):
3601
+ compute_stream = object()
3602
+ prefill_stream = MagicMock()
3603
+ gather = MagicMock()
3604
+ monkeypatch.setattr(index_group_module, "current_stream", lambda: compute_stream)
3605
+ monkeypatch.setattr(
3606
+ index_group_module.ops, "cp_gather_and_upconvert_fp8_kv_cache", gather
3607
+ )
3608
+
3609
+ resident_cache = torch.empty((2, 4, 656), dtype=torch.uint8)
3610
+ ready = MagicMock()
3611
+ cache = SimpleNamespace(
3612
+ view=SimpleNamespace(cache=resident_cache, block_size=4),
3613
+ block_table=torch.tensor([[2, 3]], dtype=torch.int32),
3614
+ )
3615
+ index_group = object.__new__(HiSparseMLAIndexGroup)
3616
+ index_group.caches = [cache]
3617
+ index_group.prefill_stream = prefill_stream
3618
+ index_group.prefill_ready_events = [ready]
3619
+ source = torch.empty((3, 4, 656), dtype=torch.uint8)
3620
+ dst = torch.empty((8, 576), dtype=torch.bfloat16)
3621
+ block_table = torch.tensor([[0, 1]], dtype=torch.int32)
3622
+ workspace_starts = torch.tensor([0], dtype=torch.int32)
3623
+ plan = SimpleNamespace(
3624
+ block_table=torch.tensor([[0, 1]], dtype=torch.int32),
3625
+ row_ids=torch.arange(8, dtype=torch.int32).view(1, -1),
3626
+ gpu_row_ids=torch.full((1, 8), -1, dtype=torch.int32),
3627
+ ensure_gpu_sources=MagicMock(),
3628
+ )
3629
+ metadata = SimpleNamespace(
3630
+ prefill=SimpleNamespace(host_staging_plan=plan),
3631
+ num_decodes=0,
3632
+ )
3633
+
3634
+ result = index_group.gather_fp8_prefill(
3635
+ 0,
3636
+ source,
3637
+ dst,
3638
+ block_table,
3639
+ workspace_starts,
3640
+ 1,
3641
+ metadata,
3642
+ 0,
3643
+ )
3644
+
3645
+ prefill_stream.wait_stream.assert_called_once_with(compute_stream)
3646
+ ready.record.assert_called_once_with(prefill_stream)
3647
+ assert result is ready
3648
+ gather.assert_called_once()
3649
+ args = gather.call_args
3650
+ assert args.args[0] is resident_cache
3651
+ assert args.args[1] is dst
3652
+ torch.testing.assert_close(args.args[2], plan.block_table)
3653
+ assert args.args[3] is workspace_starts
3654
+ assert args.args[4] == 1
3655
+ assert args.kwargs["host_cache"].data_ptr() == source.data_ptr()
3656
+ assert args.kwargs["host_row_ids"] is plan.row_ids
3657
+ assert args.kwargs["device_row_ids"] is plan.gpu_row_ids
3658
+ plan.ensure_gpu_sources.assert_called_once()
3659
+ ensure_args = plan.ensure_gpu_sources.call_args.args
3660
+ torch.testing.assert_close(ensure_args[0], cache.block_table)
3661
+ assert ensure_args[1] == 4
3662
+
3663
+
3664
def test_sparse_impl_observes_repointed_indexer_buffer():
3665
"""The MTP proposer repoints the draft's indexer at the target model's buffer
3666
after the backend impl is built, so the impl must resolve the buffer per read.
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+6
-67
@@ -41,7 +41,9 @@ def test_connector_without_divergent_hit_support_uses_common_lookup():
41
manager = MagicMock()
42
manager.get_computed_blocks.return_value = (common_blocks, 0, 0)
43
scheduler = SimpleNamespace(
44
- connector=SimpleNamespace(supports_divergent_local_hybrid_hits=False),
44
+ connector=SimpleNamespace(
45
+ supports_divergent_local_hybrid_hits=False,
46
+ ),
47
kv_cache_manager=manager,
48
)
49
@@ -61,7 +63,9 @@ def test_capable_connector_uses_divergent_partial_hit_lookup():
63
True,
64
)
65
scheduler = SimpleNamespace(
64
- connector=SimpleNamespace(supports_divergent_local_hybrid_hits=True),
66
+ connector=SimpleNamespace(
67
+ supports_divergent_local_hybrid_hits=True,
68
+ ),
69
kv_cache_manager=manager,
70
)
71
@@ -922,71 +926,6 @@ def test_connector_finish_registers_partial_tail_before_cleanup():
926
scheduler.kv_cache_manager.remove_skipped_blocks.assert_called_once()
927
928
925
-def test_block_pool_touch_pins_released_cow_target():
926
- """The connector can rescue an offered CoW target after its step-scoped
927
- retention is released by using the bound BlockPool's touch method."""
928
- hash_block_size = 2
929
- block_size = 2 * hash_block_size
930
- kv_cache_config = KVCacheConfig(
931
- num_blocks=24,
932
- kv_cache_tensors=[],
933
- kv_cache_groups=[
934
- KVCacheGroupSpec(
935
- ["full"],
936
- FullAttentionSpec(
937
- block_size=hash_block_size,
938
- num_kv_heads=1,
939
- head_size=1,
940
- dtype=torch.float32,
941
- ),
942
- ),
943
- KVCacheGroupSpec(
944
- ["mamba"],
945
- MambaSpec(
946
- block_size=block_size,
947
- shapes=(1, 1),
948
- dtypes=(torch.float32,),
949
- mamba_cache_mode="align",
950
- ),
951
- ),
952
- ],
953
- )
954
- manager = make_kv_cache_manager(
955
- kv_cache_config=kv_cache_config,
956
- max_model_len=8192,
957
- enable_caching=True,
958
- hash_block_size=hash_block_size,
959
- )
960
- req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
961
- computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
962
- assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
963
- req0.num_computed_tokens = 6
964
- req0.append_output_token_ids([3])
965
- assert manager.allocate_slots(req0, 1) is not None
966
-
967
- # Retention released before the drain (defer_block_free=False ordering).
968
- _copies, retained = manager.take_kv_cache_block_copies()
969
- manager.block_pool.free_blocks(retained)
970
-
971
- offloads = drain_boundary_state_offloads(manager)
972
- ((_group_id, block_id, boundary_tokens),) = offloads["0"]
973
- assert boundary_tokens == 6
974
- cow_block = manager.block_pool.blocks[block_id]
975
- assert cow_block.ref_cnt == 0
976
- assert block_id in _free_block_ids(manager)
977
-
978
- manager.block_pool.touch([cow_block])
979
- assert cow_block.ref_cnt == 1
980
- assert block_id not in _free_block_ids(manager)
981
-
982
- # The connector-pinned block is out of the free queue: draining every free block
983
- # neither trips the allocator's ref_cnt assert nor hands it out.
984
- new_blocks = manager.block_pool.get_new_blocks(
985
- manager.block_pool.get_num_free_blocks()
986
- )
987
- assert block_id not in {b.block_id for b in new_blocks}
988
-
989
-
929
def test_boundary_state_offload_dropped_when_request_freed_before_drain():
930
"""A hand-off recorded in the same scheduling pass as the request's death
931
must not be drained: its release hook has already run, so draining would
tests/v1/core/test_async_scheduler.py
+27
-1
@@ -13,7 +13,7 @@ from vllm.v1.request import RequestStatus
13
from vllm.v1.structured_output import StructuredOutputGrammar
14
from vllm.v1.utils import ConstantList
15
16
-from .utils import create_requests, create_scheduler
16
+from .utils import create_requests, create_scheduler, mock_kv
17
18
pytestmark = pytest.mark.cpu_test
19
@@ -145,6 +145,32 @@ def test_abort():
145
assert req.num_output_tokens == abort_order_copy.index(i)
146
147
148
+def test_connector_metadata_precedes_async_placeholder_advance(monkeypatch):
149
+ """Mirroring must see earlier unresolved outputs, excluding the current step."""
150
+ scheduler = create_scheduler(
151
+ async_scheduling=True,
152
+ use_kv_connector=mock_kv(matched_tokens=0, is_async=False),
153
+ )
154
+ (request,) = create_requests(num_requests=1, num_tokens=4, max_tokens=4)
155
+ scheduler.add_request(request)
156
+
157
+ assert scheduler.connector is not None
158
+ build_connector_meta = scheduler.connector.build_connector_meta
159
+ observed_placeholders = []
160
+
161
+ def record_placeholders(scheduler_output):
162
+ observed_placeholders.append(request.num_output_placeholders)
163
+ return build_connector_meta(scheduler_output)
164
+
165
+ monkeypatch.setattr(
166
+ scheduler.connector, "build_connector_meta", record_placeholders
167
+ )
168
+ scheduler.schedule()
169
+ scheduler.schedule()
170
+
171
+ assert observed_placeholders == [0, 1]
172
+
173
+
174
def test_preempt():
175
scheduler = create_scheduler(async_scheduling=True)
176
requests = create_requests(num_requests=10, max_tokens=20)
tests/v1/core/test_contiguous_kv_packing.py
+1
@@ -75,6 +75,7 @@ def _mock_vllm_config(layout: str | None):
75
config.cache_config = CacheConfig()
76
config.cache_config.num_gpu_blocks_override = None
77
config.cache_config.kv_cache_layout = layout
78
+ config.attention_config.hisparse_config = None
79
return config
80
81
tests/v1/core/test_kv_cache_utils.py
+144
@@ -22,6 +22,7 @@ from vllm.config import (
22
SchedulerConfig,
23
VllmConfig,
24
)
25
+from vllm.config.attention import HiSparseConfig
26
from vllm.config.kv_events import KVEventsConfig
27
from vllm.lora.request import LoRARequest
28
from vllm.multimodal.inputs import (
@@ -52,10 +53,16 @@ from vllm.v1.core.kv_cache_utils import (
53
make_block_hash_with_group_id,
54
tensor_data,
55
)
56
+from vllm.v1.hisparse.layout import (
57
+ create_hisparse_layout,
58
+ get_hisparse_gpu_memory_usage,
59
+)
60
from vllm.v1.kv_cache_interface import (
61
ChunkedLocalAttentionSpec,
62
FullAttentionSpec,
63
HiddenStateCacheSpec,
64
+ HiSparseHotSpec,
65
+ HiSparseResidentSpec,
66
KpoolTailSpec,
67
KVCacheConfig,
68
KVCacheGroupSpec,
@@ -68,18 +75,139 @@ from vllm.v1.kv_cache_interface import (
75
SinkFullAttentionSpec,
76
SlidingWindowMLASpec,
77
SlidingWindowSpec,
78
+ SparseCacheRole,
79
UniformTypeKVCacheSpecs,
80
get_kv_cache_spec_kind,
81
get_kv_cache_spec_sliding_window,
82
is_full_attention_spec,
83
iter_layer_specs,
84
)
85
+from vllm.v1.kv_cache_layout import KVCacheLayout
86
from vllm.v1.metrics.stats import CachingMetrics, PrefixCacheStats
87
from vllm.v1.request import Request
88
89
pytestmark = pytest.mark.cpu_test
90
91
92
+@pytest.mark.parametrize("gpu_block_size", [32, 64])
93
+def test_hisparse_hma_uses_resolved_gpu_block_size(monkeypatch, gpu_block_size):
94
+ specs = {
95
+ "model.layers.0.self_attn": MLAAttentionSpec(
96
+ block_size=gpu_block_size,
97
+ num_kv_heads=1,
98
+ head_size=576,
99
+ dtype=torch.bfloat16,
100
+ is_index_group_leader=True,
101
+ ),
102
+ "model.layers.0.self_attn.indexer": MLAAttentionSpec(
103
+ block_size=gpu_block_size,
104
+ num_kv_heads=1,
105
+ head_size=128,
106
+ dtype=torch.bfloat16,
107
+ cache_role=SparseCacheRole.INDEXER,
108
+ ),
109
+ }
110
+ group_spec = UniformTypeKVCacheSpecs.from_specs(specs)
111
+ assert group_spec is not None
112
+ group = KVCacheGroupSpec(list(specs), group_spec)
113
+ config = SimpleNamespace(
114
+ attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()),
115
+ model_config=SimpleNamespace(
116
+ hf_config=SimpleNamespace(index_topk=128),
117
+ max_model_len=gpu_block_size,
118
+ ),
119
+ parallel_config=SimpleNamespace(decode_context_parallel_size=1),
120
+ cache_config=SimpleNamespace(
121
+ num_gpu_blocks_override=7,
122
+ prefix_cache_retention_interval=None,
123
+ get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC,
124
+ ),
125
+ )
126
+ indexer_spec = specs["model.layers.0.self_attn.indexer"]
127
+ assert get_hisparse_gpu_memory_usage(config, [group]) == (
128
+ indexer_spec.max_memory_usage_bytes(config)
129
+ )
130
+
131
+ monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30)
132
+ cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
133
+ config, [group], available_memory=2**30
134
+ )
135
+ assert cache_config.num_blocks == 7
136
+ assert cache_config.hisparse_host_num_blocks is not None
137
+ assert cache_config.hisparse_host_num_blocks > 7
138
+
139
+ host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups
140
+ assert host_group.host_resident
141
+ assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups])
142
+ host_layers = set(host_group.layer_names)
143
+ for tensor in cache_config.kv_cache_tensors:
144
+ assert all(
145
+ (name in host_layers) == tensor.host_resident for name in tensor.layers
146
+ )
147
+ host_tensor = next(t for t in cache_config.kv_cache_tensors if t.host_resident)
148
+ host_spec = host_group.kv_cache_spec.kv_cache_specs[host_tensor.layers[0]]
149
+ assert host_tensor.block_stride == host_spec.page_size_bytes
150
+ assert host_tensor.layer_stride == (
151
+ host_spec.page_size_bytes * cache_config.hisparse_host_num_blocks
152
+ )
153
+ assert host_group.kv_cache_spec.block_size == gpu_block_size
154
+ assert indexer_group.kv_cache_spec.block_size == gpu_block_size
155
+ host_specs = host_group.kv_cache_spec.kv_cache_specs
156
+ gpu_indexer_specs = indexer_group.kv_cache_spec.kv_cache_specs
157
+ assert set(host_specs) == {"model.layers.0.self_attn"}
158
+ assert set(gpu_indexer_specs) == {"model.layers.0.self_attn.indexer"}
159
+ assert indexer_group.kv_cache_spec.prefix_cacheable
160
+ assert host_group.enable_kv_transfer
161
+ auxiliary_specs = [group.kv_cache_spec for group in auxiliary_groups]
162
+ assert any(isinstance(spec, HiSparseResidentSpec) for spec in auxiliary_specs)
163
+ assert any(isinstance(spec, HiSparseHotSpec) for spec in auxiliary_specs)
164
+ assert all(
165
+ spec.block_size == gpu_block_size
166
+ for spec in auxiliary_specs
167
+ if isinstance(spec, (HiSparseResidentSpec, HiSparseHotSpec))
168
+ )
169
+ scheduler_block_size, hash_block_size = kv_cache_utils.resolve_kv_cache_block_sizes(
170
+ cache_config,
171
+ SimpleNamespace(
172
+ cache_config=SimpleNamespace(
173
+ block_size=16,
174
+ enable_prefix_caching=True,
175
+ prefix_match_unit=None,
176
+ ),
177
+ parallel_config=SimpleNamespace(decode_context_parallel_size=1),
178
+ kv_transfer_config=object(),
179
+ ),
180
+ )
181
+ assert scheduler_block_size == hash_block_size == gpu_block_size
182
+
183
+
184
+def test_hisparse_rejects_deepseek_v4():
185
+ full_specs = {
186
+ "model.layers.0.attn": MLAAttentionSpec(
187
+ block_size=256,
188
+ num_kv_heads=1,
189
+ head_size=512,
190
+ dtype=torch.uint8,
191
+ model_version="deepseek_v4",
192
+ )
193
+ }
194
+ full_uniform = UniformTypeKVCacheSpecs.from_specs(full_specs)
195
+ assert full_uniform is not None
196
+ group = KVCacheGroupSpec(list(full_specs), full_uniform)
197
+ config = SimpleNamespace(
198
+ attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()),
199
+ model_config=SimpleNamespace(hf_config=SimpleNamespace(index_topk=512)),
200
+ cache_config=SimpleNamespace(num_gpu_blocks_override=7),
201
+ )
202
+
203
+ with pytest.raises(ValueError, match="does not support DeepSeek V4"):
204
+ create_hisparse_layout(
205
+ config,
206
+ [group],
207
+ host_budget=2**30,
208
+ )
209
+
210
+
211
@pytest.fixture(autouse=True)
212
def _auto_init_hash_fn(request):
213
hash_fn: Callable
@@ -3848,6 +3976,22 @@ def test_iter_layer_specs_returns_group_members():
3976
assert list(iter_layer_specs(wrapped)) == [full, mla]
3977
3978
3979
+def test_wrapped_mamba_group_requires_block_zeroing():
3980
+ mamba = MambaSpec(
3981
+ block_size=4,
3982
+ shapes=((4, 1),),
3983
+ dtypes=(torch.float32,),
3984
+ )
3985
+ wrapped = UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs={"mamba": mamba})
3986
+ config = KVCacheConfig(
3987
+ num_blocks=4,
3988
+ kv_cache_tensors=[],
3989
+ kv_cache_groups=[KVCacheGroupSpec(["mamba"], wrapped)],
3990
+ )
3991
+
3992
+ assert config.needs_kv_cache_zeroing
3993
+
3994
+
3995
def _spec_decode_grouping_config(method="dspark", model_type=None):
3996
"""Grouping config with an EAGLE-family speculative method enabled."""
3997
return SimpleNamespace(
tests/v1/core/test_prefix_caching.py
+1096
-1
@@ -7,6 +7,7 @@ from collections.abc import Callable
7
from dataclasses import replace
8
from math import lcm
9
from types import SimpleNamespace
10
+from unittest.mock import MagicMock
11
12
import pytest
13
import torch
@@ -14,11 +15,21 @@ import torch
15
import vllm.v1.core.kv_cache_manager as kv_cache_manager
16
import vllm.v1.core.kv_cache_utils as kv_cache_utils
17
from vllm.distributed.kv_events import (
18
+ MEDIUM_CPU,
19
MEDIUM_GPU,
20
AllBlocksCleared,
21
BlockRemoved,
22
BlockStored,
23
)
24
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
25
+ HiSparseConnector,
26
+ HiSparseConnectorScheduler,
27
+ HiSparseConnectorWorkerMetadata,
28
+)
29
+from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector
30
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
31
+ OffloadingConnector,
32
+)
33
from vllm.lora.request import LoRARequest
34
from vllm.multimodal.inputs import (
35
MultiModalFeatureSpec,
@@ -40,16 +51,23 @@ from vllm.v1.core.kv_cache_utils import (
51
init_none_hash,
52
make_block_hash_with_group_id,
53
)
54
+from vllm.v1.core.sched.output import SchedulerOutput
55
from vllm.v1.core.sched.scheduler import Scheduler
56
+from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
57
from vllm.v1.kv_cache_interface import (
58
FullAttentionSpec,
59
+ HiSparseHotSpec,
60
+ HiSparseResidentSpec,
61
KVCacheConfig,
62
+ KVCacheGroupRole,
63
KVCacheGroupSpec,
64
KVCacheSpecKind,
65
MambaSpec,
66
MLAAttentionSpec,
67
SlidingWindowSpec,
68
)
69
+from vllm.v1.outputs import KVConnectorOutput
70
+from vllm.v1.request import RequestStatus
71
72
pytestmark = pytest.mark.cpu_test
73
@@ -118,7 +136,10 @@ def make_kv_cache_manager(kv_cache_config: KVCacheConfig, **kwargs) -> KVCacheMa
136
kv_cache_config,
137
prefix_cache_retention_interval=kwargs.pop("retention_interval"),
138
)
121
- return KVCacheManager(kv_cache_config, **kwargs)
139
+ manager = KVCacheManager(kv_cache_config, **kwargs)
140
+ if kv_cache_config.hisparse_host_num_blocks is not None:
141
+ get_hisparse_coordinator(manager)
142
+ return manager
143
144
145
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:
160
)
161
162
163
+HISPARSE_BLOCK_SIZE = 16
164
+
165
+
166
+def make_hisparse_kv_cache_config(
167
+ num_blocks: int,
168
+ host_num_blocks: int,
169
+ *,
170
+ transfer_device_cache: bool = False,
171
+) -> KVCacheConfig:
172
+ source_spec = FullAttentionSpec(
173
+ block_size=HISPARSE_BLOCK_SIZE,
174
+ num_kv_heads=1,
175
+ head_size=1,
176
+ dtype=torch.float32,
177
+ )
178
+ groups = [
179
+ KVCacheGroupSpec(
180
+ ["source"],
181
+ source_spec,
182
+ role=KVCacheGroupRole.HISPARSE_SOURCE,
183
+ host_resident=True,
184
+ ),
185
+ KVCacheGroupSpec(
186
+ ["indexer"],
187
+ source_spec,
188
+ enable_kv_transfer=transfer_device_cache,
189
+ role=KVCacheGroupRole.HISPARSE_INDEXER,
190
+ ),
191
+ ]
192
+ groups.append(
193
+ KVCacheGroupSpec(
194
+ ["resident"],
195
+ HiSparseResidentSpec(
196
+ block_size=HISPARSE_BLOCK_SIZE,
197
+ page_size=HISPARSE_BLOCK_SIZE * 4,
198
+ ),
199
+ enable_kv_transfer=transfer_device_cache,
200
+ )
201
+ )
202
+ groups.append(
203
+ KVCacheGroupSpec(
204
+ ["hot"],
205
+ HiSparseHotSpec(
206
+ block_size=HISPARSE_BLOCK_SIZE,
207
+ page_size=HISPARSE_BLOCK_SIZE * 4,
208
+ blocks_per_request=2,
209
+ ),
210
+ enable_kv_transfer=False,
211
+ )
212
+ )
213
+ return KVCacheConfig(
214
+ num_blocks=num_blocks,
215
+ hisparse_host_num_blocks=host_num_blocks,
216
+ kv_cache_tensors=[],
217
+ kv_cache_groups=groups,
218
+ )
219
+
220
+
221
+def make_hisparse_kv_cache_manager(
222
+ num_blocks: int,
223
+ host_num_blocks: int,
224
+ *,
225
+ max_model_len: int = 128,
226
+ max_in_flight_tokens: int | None = None,
227
+ enable_caching: bool = False,
228
+ **config_kwargs,
229
+) -> KVCacheManager:
230
+ config = make_hisparse_kv_cache_config(num_blocks, host_num_blocks, **config_kwargs)
231
+ return make_kv_cache_manager(
232
+ config,
233
+ max_model_len=max_model_len,
234
+ max_in_flight_tokens=max_in_flight_tokens,
235
+ enable_caching=enable_caching,
236
+ hash_block_size=HISPARSE_BLOCK_SIZE,
237
+ )
238
+
239
+
240
+def test_hisparse_builds_dma_row_mirrors_across_pages():
241
+ manager = make_hisparse_kv_cache_manager(32, 16)
242
+ request = make_request(
243
+ "request",
244
+ list(range(2 * HISPARSE_BLOCK_SIZE)),
245
+ HISPARSE_BLOCK_SIZE,
246
+ sha256,
247
+ )
248
+ assert manager.allocate_slots(request, num_new_tokens=32) is not None
249
+ coordinator = get_hisparse_coordinator(manager)
250
+ resident_blocks = coordinator.resident_managers[0].req_to_blocks[request.request_id]
251
+ host_blocks = coordinator.host_manager.req_to_blocks[request.request_id]
252
+
253
+ mirrors = coordinator.build_row_mirrors([(request.request_id, 14, 4)])
254
+
255
+ assert [mirror.source_starts for mirror in mirrors] == [
256
+ (resident_blocks[0].block_id * HISPARSE_BLOCK_SIZE + 14,),
257
+ (resident_blocks[1].block_id * HISPARSE_BLOCK_SIZE,),
258
+ ]
259
+ assert [mirror.destination_start for mirror in mirrors] == [
260
+ host_blocks[0].block_id * HISPARSE_BLOCK_SIZE + 14,
261
+ host_blocks[1].block_id * HISPARSE_BLOCK_SIZE,
262
+ ]
263
+ assert [mirror.num_rows for mirror in mirrors] == [2, 2]
264
+
265
+
266
+def test_hisparse_async_speculation_mirrors_uncertain_position_range():
267
+ """Unresolved drafts must not leave gaps in the eager host mirror."""
268
+ coordinator = MagicMock(host_group_id=0)
269
+ coordinator.take_block_table_updates.return_value = {}
270
+ coordinator.build_offload_command.return_value = None
271
+ coordinator.build_row_mirrors.return_value = ()
272
+ scheduler = HiSparseConnectorScheduler(
273
+ async_speculative=True,
274
+ draft_kv_lookahead=4,
275
+ )
276
+ scheduler.bind_coordinator(coordinator)
277
+ request = SimpleNamespace(request_id="request", num_output_placeholders=3)
278
+ connector = object.__new__(HiSparseConnector)
279
+ connector.connector_scheduler = scheduler
280
+ connector.update_state_after_alloc(request, None, 0)
281
+ scheduler_output = SimpleNamespace(
282
+ block_table_updates=None,
283
+ kv_cache_block_copies=None,
284
+ scheduled_new_reqs=[],
285
+ scheduled_cached_reqs=SimpleNamespace(
286
+ new_block_ids=[],
287
+ req_ids=["request"],
288
+ num_computed_tokens=[103],
289
+ ),
290
+ num_scheduled_tokens={"request": 4},
291
+ )
292
+
293
+ scheduler.build_connector_meta(scheduler_output)
294
+
295
+ coordinator.build_row_mirrors.assert_called_once_with((("request", 100, 11),))
296
+
297
+ # Acceptance arrives between steps; the connector must read the live count.
298
+ request.num_output_placeholders = 1
299
+ scheduler.build_connector_meta(scheduler_output)
300
+ coordinator.build_row_mirrors.assert_called_with((("request", 102, 9),))
301
+
302
+ connector.request_finished_all_groups(request, ())
303
+ assert not scheduler.requests
304
+
305
+
306
+def test_hisparse_reports_when_context_is_fully_resident():
307
+ """A page the pool reuses under a host-reading request stops being resident."""
308
+ manager = make_hisparse_kv_cache_manager(32, 16)
309
+ tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
310
+ request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256)
311
+ assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None
312
+ coordinator = get_hisparse_coordinator(manager)
313
+ scheduled = ((request.request_id, len(tokens) - 1, 1),)
314
+ assert coordinator.all_context_pages_resident(scheduled)
315
+
316
+ for hot_manager in coordinator.hot_managers:
317
+ hot_manager.require_hot(request.request_id)
318
+ assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None
319
+ _publish_hisparse_pages(manager)
320
+ pool = manager.block_pool
321
+ pool.get_new_blocks(pool.get_num_free_blocks())
322
+
323
+ assert not coordinator.all_context_pages_resident(scheduled)
324
+ assert coordinator.take_block_table_updates().keys() == {request.request_id}
325
+
326
+
327
+def test_hisparse_host_prefix_can_be_completed_by_indexer_offload():
328
+ """Keep indexer-only imports host-backed in the resident block table."""
329
+ manager = make_hisparse_kv_cache_manager(
330
+ 32,
331
+ 16,
332
+ enable_caching=True,
333
+ )
334
+ tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
335
+ original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
336
+ assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None
337
+ spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers
338
+ spill_counts = {spill.transfer_id: 1 for spill in spills}
339
+ get_hisparse_coordinator(manager).update_spills(spill_counts, spill_counts)
340
+ _, indexer_blocks, _, _ = manager.get_blocks(original.request_id).blocks
341
+ evicted_indexer_id = indexer_blocks[2].block_id
342
+ manager.free(original)
343
+ manager.block_pool.evict_blocks({evicted_indexer_id})
344
+
345
+ resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
346
+ blocks, num_local, _ = manager.get_computed_blocks(resumed)
347
+ max_completion = HISPARSE_BLOCK_SIZE
348
+ assert num_local == 2 * HISPARSE_BLOCK_SIZE
349
+ assert [len(group_blocks) for group_blocks in blocks.blocks] == [3, 2, 0, 0]
350
+
351
+ allocated = manager.allocate_slots(
352
+ resumed,
353
+ num_new_tokens=1,
354
+ num_new_computed_tokens=num_local,
355
+ new_computed_blocks=blocks,
356
+ num_external_computed_tokens=max_completion,
357
+ )
358
+
359
+ assert allocated is not None
360
+ source, indexer, resident, hot = manager.get_blocks(resumed.request_id).blocks
361
+ assert len(source) == len(indexer) == len(resident) == 4
362
+ assert len(hot) == 2
363
+ # The local prefix is adopted from shadow pages (GPU-resident), while the
364
+ # externally imported page stays host-backed until its tail allocation.
365
+ assert not any(block.is_null for block in resident[:2])
366
+ assert resident[2].is_null
367
+ assert not resident[3].is_null
368
+
369
+
370
+def test_hisparse_indexer_offload_is_capped_by_missing_host_prefix():
371
+ manager = make_hisparse_kv_cache_manager(
372
+ 32,
373
+ 16,
374
+ enable_caching=True,
375
+ )
376
+ tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
377
+ original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
378
+ assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None
379
+ spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers
380
+ spill_counts = {spill.transfer_id: 1 for spill in spills}
381
+ get_hisparse_coordinator(manager).update_spills(spill_counts, spill_counts)
382
+ host_blocks, _, _, _ = manager.get_blocks(original.request_id).blocks
383
+ evicted_host_id = host_blocks[0].block_id
384
+ manager.free(original)
385
+ get_hisparse_coordinator(manager).host_manager.block_pool.evict_blocks(
386
+ {evicted_host_id}
387
+ )
388
+
389
+ resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
390
+ _, num_local, _ = manager.get_computed_blocks(resumed)
391
+ connector = object.__new__(OffloadingConnector)
392
+ connector._kv_cache_config = manager.kv_cache_config
393
+ connector.bind_kv_cache_manager(manager)
394
+ assert num_local == 0
395
+ assert connector._max_loadable_tokens(resumed, num_local) == 0
396
+
397
+
398
+@pytest.mark.parametrize("nested", [False, True])
399
+@pytest.mark.parametrize(
400
+ "available_tokens,imported_tokens,host_blocks",
401
+ [(0, 0, 2), (16, 16, 2), (64, 32, 3)],
402
+)
403
+def test_connector_completes_partial_prefix_without_importing_missing_host_kv(
404
+ nested, available_tokens, imported_tokens, host_blocks
405
+):
406
+ manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
407
+ tokens = list(range(64))
408
+ original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
409
+ assert manager.allocate_slots(original, num_new_tokens=64) is not None
410
+ _publish_hisparse_pages(manager)
411
+ host, indexer, _, _ = manager.get_blocks(original.request_id).blocks
412
+ missing_indexer = indexer[1].block_id
413
+ missing_host = host[0].block_id
414
+ manager.free(original)
415
+ manager.block_pool.evict_blocks({missing_indexer})
416
+
417
+ defer = True
418
+
419
+ def lookup(request, num_computed_tokens, max_num_new_tokens=None):
420
+ nonlocal defer
421
+ if defer:
422
+ defer = False
423
+ return None, False
424
+ available = available_tokens
425
+ if max_num_new_tokens is not None:
426
+ available = min(available, max_num_new_tokens)
427
+ return available, available > 0
428
+
429
+ connector = object.__new__(OffloadingConnector)
430
+ connector._kv_cache_config = manager.kv_cache_config
431
+ connector.connector_scheduler = SimpleNamespace(get_num_new_matched_tokens=lookup)
432
+ if nested:
433
+ for _ in range(2):
434
+ outer = object.__new__(MultiConnector)
435
+ outer._connectors = (connector,)
436
+ outer._requests_to_connector = {}
437
+ connector = outer
438
+ connector.bind_kv_cache_manager(manager)
439
+ scheduler = SimpleNamespace(connector=connector, kv_cache_manager=manager)
440
+ resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
441
+
442
+ blocks, local, _, diverged = Scheduler._get_local_prefix_cache_hit(
443
+ scheduler, resumed
444
+ )
445
+ assert local == 16 and not diverged
446
+ assert connector.get_num_new_matched_tokens(resumed, local) == (None, False)
447
+
448
+ blocks, local, _, diverged = Scheduler._get_local_prefix_cache_hit(
449
+ scheduler, resumed
450
+ )
451
+ external, asynchronous = connector.get_num_new_matched_tokens(resumed, local)
452
+ assert external == imported_tokens
453
+ assert asynchronous == (external > 0)
454
+ assert [len(group) for group in blocks.blocks] == [3, 1, 0, 0]
455
+ assert (
456
+ manager.allocate_slots(
457
+ resumed,
458
+ int(external == 0),
459
+ delay_cache_blocks=True,
460
+ num_new_computed_tokens=local,
461
+ new_computed_blocks=blocks,
462
+ num_external_computed_tokens=external,
463
+ )
464
+ is not None
465
+ )
466
+ allocated = manager.get_blocks(resumed.request_id).blocks
467
+ assert len(allocated[0]) == host_blocks
468
+ assert allocated[2][1].is_null == (external > 0)
469
+ manager.free(resumed)
470
+
471
+ # A surviving indexer offload must not make absent host KV look computed.
472
+ get_hisparse_coordinator(manager).host_manager.block_pool.evict_blocks(
473
+ {missing_host}
474
+ )
475
+ retry = make_request("retry", tokens, HISPARSE_BLOCK_SIZE, sha256)
476
+ _, local, _, _ = Scheduler._get_local_prefix_cache_hit(scheduler, retry)
477
+ assert local == 0
478
+ assert connector.get_num_new_matched_tokens(retry, local) == (0, False)
479
+
480
+
481
+def allocate_external_prefix(
482
+ manager: KVCacheManager, request: Request, num_tokens: int
483
+) -> KVCacheBlocks | None:
484
+ return manager.allocate_slots(
485
+ request,
486
+ num_new_tokens=0,
487
+ num_external_computed_tokens=num_tokens,
488
+ delay_cache_blocks=True,
489
+ full_sequence_must_fit=True,
490
+ )
491
+
492
+
493
+@pytest.mark.parametrize("enable_caching", [False, True])
494
+def test_hisparse_keeps_resident_pages_until_hot_buffer_is_allocated(enable_caching):
495
+ """Later admissions must not remove a scheduled request's only readable tier."""
496
+ manager = make_hisparse_kv_cache_manager(
497
+ 32, 32, max_model_len=256, enable_caching=enable_caching
498
+ )
499
+ first = make_request("first", list(range(160)), HISPARSE_BLOCK_SIZE, sha256)
500
+ assert manager.allocate_slots(first, num_new_tokens=128) is not None
501
+ _publish_hisparse_pages(manager)
502
+ first.num_computed_tokens = 128
503
+ pool = manager.block_pool
504
+ held = pool.get_new_blocks(pool.get_num_free_blocks() - 3)
505
+
506
+ assert manager.allocate_slots(first, num_new_tokens=1) is not None
507
+ second = make_request("second", [1000] * 16, HISPARSE_BLOCK_SIZE, sha256)
508
+ manager.allocate_slots(second, num_new_tokens=16)
509
+ coordinator = get_hisparse_coordinator(manager)
510
+ scheduled = (("first", 128, 1),)
511
+ assert coordinator.all_context_pages_resident(scheduled)
512
+ assert all(not m.has_hot("first") for m in coordinator.hot_managers)
513
+
514
+ # Once another allocation finishes, A can acquire its hot buffer and
515
+ # safely release its clean pages for other requests to reuse.
516
+ pool.free_blocks(held)
517
+ assert manager.allocate_slots(first, num_new_tokens=1) is not None
518
+ assert all(m.has_hot("first") for m in coordinator.hot_managers)
519
+ pool.get_new_blocks(pool.get_num_free_blocks())
520
+ assert not coordinator.all_context_pages_resident(scheduled)
521
+ assert "first" in coordinator.take_block_table_updates()
522
+
523
+
524
+def test_hisparse_full_pool_keeps_pages_pinned_until_preemption():
525
+ """Without room for a hot buffer, admission must defer instead of losing KV."""
526
+ manager = make_hisparse_kv_cache_manager(18, 18, max_model_len=160)
527
+ first = make_request("first", list(range(128)), HISPARSE_BLOCK_SIZE, sha256)
528
+ assert manager.allocate_slots(first, num_new_tokens=128) is not None
529
+ _publish_hisparse_pages(manager)
530
+ assert manager.block_pool.get_num_free_blocks() == 1
531
+ first.num_computed_tokens = 128
532
+ assert manager.allocate_slots(first, num_new_tokens=1) is None
533
+ assert get_hisparse_coordinator(manager).all_context_pages_resident(
534
+ (("first", 127, 1),)
535
+ )
536
+
537
+ manager.free(first)
538
+ second = make_request("second", [1000] * 16, HISPARSE_BLOCK_SIZE, sha256)
539
+ assert manager.allocate_slots(second, num_new_tokens=16) is not None
540
+
541
+
542
+def test_hisparse_materializes_prefix_without_allocating_hot_blocks():
543
+ """A host prefix becomes visible only when every page is durable.
544
+
545
+ Completing a later page first must not expose a prefix with a hole.
546
+ """
547
+ manager = make_hisparse_kv_cache_manager(
548
+ 32,
549
+ 16,
550
+ enable_caching=True,
551
+ )
552
+ tokens = list(range(2 * HISPARSE_BLOCK_SIZE))
553
+ request = make_request("resident", tokens, HISPARSE_BLOCK_SIZE, sha256)
554
+ assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None
555
+ blocks = manager.get_block_ids(request.request_id)
556
+ assert len(blocks[2]) == 2
557
+ assert blocks[3] == []
558
+
559
+ spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers
560
+ assert len(spills) == 2
561
+ assert all(spill.after_forward for spill in spills)
562
+ resident_blocks = manager.get_blocks(request.request_id).blocks[2]
563
+ assert [block.ref_cnt for block in resident_blocks] == [2, 2]
564
+ spill_counts = {spill.transfer_id: 1 for spill in spills}
565
+ duplicate = make_request(
566
+ "duplicate",
567
+ list(range(3 * HISPARSE_BLOCK_SIZE)),
568
+ HISPARSE_BLOCK_SIZE,
569
+ sha256,
570
+ )
571
+ _, num_computed, _ = manager.get_computed_blocks(duplicate)
572
+ assert num_computed == 0
573
+
574
+ get_hisparse_coordinator(manager).update_spills(spill_counts, {})
575
+ assert [block.ref_cnt for block in resident_blocks] == [2, 2]
576
+ _, num_computed, _ = manager.get_computed_blocks(duplicate)
577
+ assert num_computed == 0
578
+
579
+ get_hisparse_coordinator(manager).update_spills({}, {spills[1].transfer_id: 1})
580
+ assert [block.ref_cnt for block in resident_blocks] == [2, 1]
581
+ _, num_computed, _ = manager.get_computed_blocks(duplicate)
582
+ assert num_computed == 0
583
+
584
+ get_hisparse_coordinator(manager).update_spills({}, {spills[0].transfer_id: 1})
585
+ _, num_computed, _ = manager.get_computed_blocks(duplicate)
586
+ assert num_computed == 2 * HISPARSE_BLOCK_SIZE
587
+ blocks = manager.get_block_ids(request.request_id)
588
+ assert len(blocks[2]) == 2
589
+ assert blocks[3] == []
590
+
591
+
592
+def test_hisparse_materialization_respects_per_step_spill_budget():
593
+ """Prefix publication must not bypass the configured spill batch limit."""
594
+ manager = make_hisparse_kv_cache_manager(
595
+ 32,
596
+ 16,
597
+ enable_caching=True,
598
+ )
599
+ coordinator = get_hisparse_coordinator(manager)
600
+ coordinator.max_spill_pages = 1
601
+ tokens = list(range(2 * HISPARSE_BLOCK_SIZE))
602
+ request = make_request("bounded", tokens, HISPARSE_BLOCK_SIZE, sha256)
603
+
604
+ assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None
605
+ first = coordinator.build_offload_command().page_transfers
606
+ assert len(first) == 1
607
+
608
+ coordinator.plan_prefix_materialization(request.request_id, len(tokens))
609
+ second = coordinator.build_offload_command().page_transfers
610
+ assert len(second) == 1
611
+ assert first[0].transfer_id != second[0].transfer_id
612
+
613
+
614
+def test_hisparse_deferred_free_retains_host_blocks_until_fence():
615
+ """Detaching a request must not make in-flight host pages reusable."""
616
+ manager = make_hisparse_kv_cache_manager(16, 16)
617
+ host_pool = get_hisparse_coordinator(manager).get_host_block_pool()
618
+ assert host_pool is not None
619
+ host_free = host_pool.get_num_free_blocks()
620
+ device_free = manager.block_pool.get_num_free_blocks()
621
+ request = make_request("free", list(range(16)), HISPARSE_BLOCK_SIZE, sha256)
622
+ assert manager.allocate_slots(request, 16, delay_cache_blocks=True) is not None
623
+ blocks = manager.pop_blocks_for_free(request)
624
+ assert host_pool.get_num_free_blocks() < host_free
625
+ assert blocks and any(host_pool.blocks[b.block_id] is b for b in blocks)
626
+ manager.block_pool.free_blocks(reversed(blocks))
627
+ assert host_pool.get_num_free_blocks() == host_free
628
+ assert manager.block_pool.get_num_free_blocks() == device_free
629
+
630
+
631
+def test_hisparse_host_cow_copy_is_drained_without_a_gpu_pool():
632
+ """Host-only copy-on-write work must reach the worker copy queue."""
633
+ manager = make_hisparse_kv_cache_manager(16, 16)
634
+ source_manager = manager.coordinator.single_type_managers[0]
635
+ source_block = source_manager.block_pool.get_new_blocks(1)[0]
636
+ source_manager.req_to_blocks["cow"] = [source_block]
637
+ source_manager._partial_hit_reqs["cow"] = (0, source_block)
638
+
639
+ new_blocks = source_manager.allocate_new_blocks(
640
+ "cow", HISPARSE_BLOCK_SIZE, HISPARSE_BLOCK_SIZE
641
+ )
642
+ new_block_ids = manager.take_new_block_ids()
643
+ coordinator = get_hisparse_coordinator(manager)
644
+ assert manager.take_kv_cache_block_copies() == ([], [])
645
+ connector = HiSparseConnectorScheduler(async_speculative=False)
646
+ connector.bind_coordinator(coordinator)
647
+ output = SchedulerOutput.make_empty()
648
+ copies = connector.build_connector_meta(output).host_block_copies
649
+ assert output.has_sync_kv_loads
650
+
651
+ assert new_blocks and new_block_ids == []
652
+ assert len(copies) == 1
653
+ assert copies[0].src_block_id == source_block.block_id
654
+ assert copies[0].dst_block_id == new_blocks[0].block_id
655
+ source_manager.free("cow")
656
+ assert source_block.ref_cnt > 0 and new_blocks[0].ref_cnt > 0
657
+ assert coordinator.has_pending_work()
658
+ connector.update_connector_output(
659
+ KVConnectorOutput(
660
+ kv_connector_worker_meta=HiSparseConnectorWorkerMetadata(
661
+ {}, {}, completed_host_copy_dst_ids=(copies[0].dst_block_id,)
662
+ )
663
+ )
664
+ )
665
+ assert source_block.ref_cnt == new_blocks[0].ref_cnt == 0
666
+ assert not coordinator.has_pending_work()
667
+
668
+
669
+def test_hisparse_inflight_host_import_reserves_remaining_gpu_pages():
670
+ """An in-flight host import must reserve its unwritten resident pages."""
671
+ manager = make_hisparse_kv_cache_manager(
672
+ 11,
673
+ 16,
674
+ transfer_device_cache=True,
675
+ )
676
+ request = make_request(
677
+ "partial-import",
678
+ list(range(6 * HISPARSE_BLOCK_SIZE)),
679
+ HISPARSE_BLOCK_SIZE,
680
+ sha256,
681
+ )
682
+ imported_tokens = 4 * HISPARSE_BLOCK_SIZE
683
+
684
+ assert allocate_external_prefix(manager, request, imported_tokens) is not None
685
+ required = manager.coordinator.get_num_blocks_to_allocate(
686
+ request_id=request.request_id,
687
+ num_tokens=request.num_tokens,
688
+ new_computed_blocks=manager.empty_kv_cache_blocks.blocks,
689
+ num_encoder_tokens=0,
690
+ total_computed_tokens=imported_tokens,
691
+ num_local_computed_tokens=0,
692
+ num_tokens_main_model=request.num_tokens,
693
+ apply_admission_cap=True,
694
+ )
695
+
696
+ assert required == 4
697
+
698
+
699
+@pytest.mark.parametrize(
700
+ "full_sequence_must_fit,host_num_blocks,admitted",
701
+ [(True, 2, False), (True, 3, True), (False, 2, False), (False, 3, True)],
702
+)
703
+def test_hisparse_async_admission_requires_only_import_destinations(
704
+ full_sequence_must_fit, host_num_blocks, admitted, tmp_path
705
+):
706
+ """Imports need host destinations, but do not reserve future prefill pages."""
707
+ from .utils import create_scheduler, mock_kv
708
+
709
+ (tmp_path / "config.json").write_text(
710
+ '{"architectures": ["OPTForCausalLM"], "model_type": "opt"}'
711
+ )
712
+ scheduler = create_scheduler(
713
+ model=str(tmp_path),
714
+ skip_tokenizer_init=True,
715
+ max_model_len=128,
716
+ use_kv_connector=mock_kv(matched_tokens=HISPARSE_BLOCK_SIZE, is_async=True),
717
+ )
718
+ manager = make_hisparse_kv_cache_manager(32, host_num_blocks)
719
+ scheduler.kv_cache_manager = manager
720
+ scheduler.kv_cache_config = manager.kv_cache_config
721
+ scheduler.scheduler_reserve_full_isl = full_sequence_must_fit
722
+ inflight = make_request(
723
+ "inflight", list(range(3 * HISPARSE_BLOCK_SIZE)), HISPARSE_BLOCK_SIZE, sha256
724
+ )
725
+ scheduler.add_request(inflight)
726
+ scheduler.schedule()
727
+ assert inflight in scheduler._inflight_prefills
728
+
729
+ request = make_request(
730
+ "waiting", list(range(4 * HISPARSE_BLOCK_SIZE)), HISPARSE_BLOCK_SIZE, sha256
731
+ )
732
+ scheduler.add_request(request)
733
+ scheduler.schedule()
734
+
735
+ assert (request in scheduler._inflight_prefills) == admitted
736
+ assert bool(manager.get_blocks(request.request_id).blocks[0]) == admitted
737
+
738
+
739
+@pytest.mark.parametrize(
740
+ "evictable,local_tokens,free_blocks,admitted",
741
+ [
742
+ (False, HISPARSE_BLOCK_SIZE, 1, True),
743
+ (True, HISPARSE_BLOCK_SIZE, 1, False),
744
+ (False, HISPARSE_BLOCK_SIZE - 1, 1, False),
745
+ (False, HISPARSE_BLOCK_SIZE - 1, 2, True),
746
+ ],
747
+)
748
+def test_hisparse_import_capacity_includes_hits_and_cow(
749
+ evictable, local_tokens, free_blocks, admitted
750
+):
751
+ """Touching an evictable hit and copying a partial hit consume capacity."""
752
+ manager = make_hisparse_kv_cache_manager(32, free_blocks + 2 - int(evictable))
753
+ coordinator = get_hisparse_coordinator(manager)
754
+ host_pool = coordinator.get_host_block_pool()
755
+ assert host_pool is not None
756
+ hit = host_pool.get_new_blocks(1)[0]
757
+ if evictable:
758
+ host_pool.free_blocks([hit])
759
+ assert host_pool.get_num_free_blocks() == free_blocks
760
+ assert (
761
+ coordinator.host_manager.get_num_blocks_to_allocate(
762
+ "waiting", 64, [hit], 32, local_tokens, 64
763
+ )
764
+ <= host_pool.get_num_free_blocks()
765
+ ) == admitted
766
+
767
+
768
+@pytest.mark.parametrize("enable_caching", [False, True])
769
+def test_hisparse_host_exhaustion_keeps_gpu_pages_readable(enable_caching):
770
+ """Pages without a host destination must survive GPU cache reclamation."""
771
+ manager = make_hisparse_kv_cache_manager(32, 2, enable_caching=enable_caching)
772
+ request = make_request("resident", list(range(64)), HISPARSE_BLOCK_SIZE, sha256)
773
+ assert manager.allocate_slots(request, 64) is not None
774
+ coordinator = get_hisparse_coordinator(manager)
775
+ host = coordinator.host_manager
776
+ assert host is not None
777
+ assert [b.is_null for b in host.req_to_blocks[request.request_id]] == [
778
+ False,
779
+ True,
780
+ True,
781
+ True,
782
+ ]
783
+ _publish_hisparse_pages(manager)
784
+ for hot in coordinator.hot_managers:
785
+ hot.require_hot(request.request_id)
786
+ hot.allocate_new_blocks(request.request_id, 64, 64)
787
+ coordinator.update_residency(request.request_id)
788
+ for resident in coordinator.resident_managers:
789
+ for block in resident.req_to_blocks[request.request_id][1:]:
790
+ assert not block.is_null and block.ref_cnt > 0
791
+ assert coordinator.request_states[request.request_id].valid_pages == {0}
792
+ pool = coordinator.get_host_block_pool()
793
+ assert pool is not None
794
+ manager.free(request)
795
+ assert pool.get_num_free_blocks() == 1
796
+
797
+
798
+def test_hisparse_host_cow_takes_priority_over_new_pages():
799
+ """The last host block must hold the private tail, not a fresh page."""
800
+ manager = make_hisparse_kv_cache_manager(32, 3)
801
+ coordinator = get_hisparse_coordinator(manager)
802
+ host = coordinator.host_manager
803
+ assert host is not None
804
+ source = host.block_pool.get_new_blocks(1)[0]
805
+ host.req_to_blocks["cow"] = [source]
806
+ host._partial_hit_reqs["cow"] = (0, source)
807
+
808
+ host.allocate_new_blocks("cow", 32, 32)
809
+
810
+ private, missing = host.req_to_blocks["cow"]
811
+ assert not private.is_null and private is not source
812
+ assert missing.is_null
813
+ assert host.take_host_cow_copies() == [(source, private)]
814
+
815
+
816
+@pytest.mark.parametrize("failed", [False, True])
817
+def test_host_receive_completion_without_spill_metadata(failed):
818
+ """Only successful external receives make host pages readable."""
819
+ manager = make_hisparse_kv_cache_manager(32, 16)
820
+ request = make_request("import", list(range(32)), HISPARSE_BLOCK_SIZE, sha256)
821
+ assert allocate_external_prefix(manager, request, 32) is not None
822
+ request.num_computed_tokens = 32
823
+ request.status = RequestStatus.WAITING_FOR_REMOTE_KVS
824
+ scheduler = HiSparseConnectorScheduler(async_speculative=False)
825
+ scheduler.bind_coordinator(get_hisparse_coordinator(manager))
826
+ scheduler.requests[request.request_id] = request
827
+ connector = object.__new__(HiSparseConnector)
828
+ connector.connector_scheduler = scheduler
829
+ composite = object.__new__(MultiConnector)
830
+ composite._connectors = (connector,)
831
+ composite.update_connector_output(
832
+ KVConnectorOutput(
833
+ finished_recving={request.request_id},
834
+ failed_recving={request.request_id} if failed else set(),
835
+ )
836
+ )
837
+ state = get_hisparse_coordinator(manager).request_states.get(request.request_id)
838
+ if failed:
839
+ assert state is None
840
+ else:
841
+ assert state is not None and state.valid_pages == {0, 1}
842
+
843
+
844
+def test_hisparse_recomputation_does_not_complete_failed_host_import():
845
+ """Allocating recomputation is not evidence that imported host KV is valid."""
846
+ manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
847
+ request = make_request(
848
+ "partial-failure", list(range(48)), HISPARSE_BLOCK_SIZE, sha256
849
+ )
850
+ assert allocate_external_prefix(manager, request, 32) is not None
851
+ # A partial load failure leaves only the first page computed. The next
852
+ # allocation reaches the original import boundary before the worker runs.
853
+ request.num_computed_tokens = 16
854
+ manager.cache_blocks(request, 16)
855
+ assert manager.allocate_slots(request, 16) is not None
856
+ state = get_hisparse_coordinator(manager).request_states.get(request.request_id)
857
+ assert state is None or 1 not in state.valid_pages
858
+
859
+
860
+def test_hisparse_host_import_ignores_unsealed_tail():
861
+ """A partial imported page must not become readable from stale host data."""
862
+ manager = make_hisparse_kv_cache_manager(16, 16)
863
+ coordinator = get_hisparse_coordinator(manager)
864
+
865
+ coordinator.complete_host_import("partial", HISPARSE_BLOCK_SIZE + 1)
866
+
867
+ state = coordinator.request_states["partial"]
868
+ assert state.valid_pages == {0}
869
+ assert state.ready_prefix_pages == 1
870
+
871
+
872
+def test_hisparse_resident_request_can_grow_without_hot_capacity():
873
+ """Running resident history must not be mistaken for a new host-prefix hit."""
874
+ manager = make_hisparse_kv_cache_manager(8, 16)
875
+ request = make_request("resident", list(range(48)), HISPARSE_BLOCK_SIZE, sha256)
876
+ assert manager.allocate_slots(request, num_new_tokens=32) is not None
877
+ request.num_computed_tokens = 32
878
+
879
+ # Three free blocks fit the next indexer/resident pages, but not a hot region.
880
+ assert manager.block_pool.get_num_free_blocks() == 3
881
+ assert manager.allocate_slots(request, num_new_tokens=1) is not None
882
+ assert manager.get_block_ids(request.request_id)[3] == []
883
+
884
+
885
+def test_hisparse_capacity_query_does_not_require_hot_blocks():
886
+ """A read-only capacity query must not mutate hot-block requirements."""
887
+ manager = make_hisparse_kv_cache_manager(16, 16)
888
+ coordinator = get_hisparse_coordinator(manager)
889
+ host_pool = coordinator.get_host_block_pool()
890
+ assert host_pool is not None
891
+ computed: tuple[list[KVCacheBlock], ...] = (
892
+ [host_pool.blocks[1]],
893
+ [],
894
+ [],
895
+ [],
896
+ )
897
+
898
+ manager.coordinator.get_num_blocks_to_allocate(
899
+ request_id="query-only",
900
+ num_tokens=HISPARSE_BLOCK_SIZE,
901
+ new_computed_blocks=computed,
902
+ num_encoder_tokens=0,
903
+ total_computed_tokens=HISPARSE_BLOCK_SIZE,
904
+ num_local_computed_tokens=HISPARSE_BLOCK_SIZE,
905
+ num_tokens_main_model=HISPARSE_BLOCK_SIZE,
906
+ )
907
+
908
+ assert all(
909
+ "query-only" not in hot_manager.hot_required
910
+ for hot_manager in coordinator.hot_managers
911
+ )
912
+
913
+
914
+def _publish_hisparse_pages(manager: KVCacheManager) -> None:
915
+ """Ack all planned spills so host pages publish to the prefix cache."""
916
+ command = get_hisparse_coordinator(manager).build_offload_command()
917
+ counts = {spill.transfer_id: 1 for spill in command.page_transfers}
918
+ get_hisparse_coordinator(manager).update_spills(counts, counts)
919
+
920
+
921
+def test_hisparse_events_report_host_and_device_placement():
922
+ """Host publication and eviction must not advertise GPU-resident KV."""
923
+ manager = make_kv_cache_manager(
924
+ make_hisparse_kv_cache_config(32, 16),
925
+ max_model_len=128,
926
+ enable_caching=True,
927
+ hash_block_size=HISPARSE_BLOCK_SIZE,
928
+ enable_kv_cache_events=True,
929
+ )
930
+ request = make_request("request", list(range(64)), HISPARSE_BLOCK_SIZE, sha256)
931
+ assert manager.allocate_slots(request, 64) is not None
932
+ _publish_hisparse_pages(manager)
933
+
934
+ stored = [
935
+ event for event in manager.take_events() if isinstance(event, BlockStored)
936
+ ]
937
+ assert {(event.group_idx, event.medium) for event in stored} == {
938
+ (0, MEDIUM_CPU),
939
+ (1, MEDIUM_GPU),
940
+ }
941
+
942
+ host_block = manager.get_blocks(request.request_id).blocks[0][0]
943
+ get_hisparse_coordinator(manager).host_manager.block_pool.evict_blocks(
944
+ {host_block.block_id}
945
+ )
946
+ [removed] = manager.take_events()
947
+ assert isinstance(removed, BlockRemoved)
948
+ assert removed.medium == MEDIUM_CPU
949
+ assert removed.group_idx == 0
950
+
951
+
952
+def test_mamba_boundary_handoffs_do_not_pin_obsolete_blocks():
953
+ """Draining unused connector offers must not retain old recurrent states."""
954
+ manager = make_kv_cache_manager(
955
+ KVCacheConfig(
956
+ num_blocks=16,
957
+ kv_cache_tensors=[],
958
+ kv_cache_groups=[
959
+ KVCacheGroupSpec(
960
+ ["mamba"],
961
+ MambaSpec(
962
+ block_size=16,
963
+ shapes=((1,),),
964
+ dtypes=(torch.float32,),
965
+ mamba_cache_mode="align",
966
+ ),
967
+ )
968
+ ],
969
+ ),
970
+ max_model_len=128,
971
+ enable_caching=True,
972
+ hash_block_size=16,
973
+ )
974
+ request = make_request("request", list(range(64)), 16, sha256)
975
+ old_blocks = []
976
+ for end in (16, 32, 48):
977
+ assert manager.allocate_slots(request, 16) is not None
978
+ request.num_computed_tokens = end
979
+ offers = manager.take_boundary_state_offloads()
980
+ [(_, block_id, boundary)] = offers[request.request_id]
981
+ assert boundary == end
982
+ old_blocks.append(manager.block_pool.blocks[block_id])
983
+ manager.remove_skipped_blocks(request.request_id, end, 64)
984
+ manager.new_step_starts()
985
+
986
+ assert all(block.ref_cnt == 0 for block in old_blocks[:-1])
987
+
988
+
989
+def test_hisparse_prefix_hit_adopts_gpu_shadow_pages():
990
+ """A host prefix hit must come back GPU-resident while shadows survive."""
991
+ manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
992
+ tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
993
+ original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
994
+ assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None
995
+ _publish_hisparse_pages(manager)
996
+ original_resident_ids = [
997
+ block.block_id for block in manager.get_blocks("original").blocks[2]
998
+ ]
999
+ manager.free(original)
1000
+
1001
+ resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
1002
+ computed, num_computed, _ = manager.get_computed_blocks(resumed)
1003
+ assert num_computed == 3 * HISPARSE_BLOCK_SIZE
1004
+ assert manager.allocate_slots(
1005
+ resumed,
1006
+ num_new_tokens=len(tokens) - num_computed,
1007
+ num_new_computed_tokens=num_computed,
1008
+ new_computed_blocks=computed,
1009
+ )
1010
+ resident_blocks = manager.get_blocks("resumed").blocks[2]
1011
+ assert [block.block_id for block in resident_blocks[:3]] == (
1012
+ original_resident_ids[:3]
1013
+ )
1014
+ assert not any(block.is_null for block in resident_blocks[:3])
1015
+ assert get_hisparse_coordinator(manager).all_context_pages_resident(
1016
+ ((resumed.request_id, num_computed, len(tokens) - num_computed),)
1017
+ )
1018
+
1019
+
1020
+def test_hisparse_host_backed_request_accepts_local_prefix_hit():
1021
+ """Local group-completion hits need no external host import allocation."""
1022
+ manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
1023
+ tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
1024
+ original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
1025
+ assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None
1026
+ _publish_hisparse_pages(manager)
1027
+ manager.free(original)
1028
+
1029
+ resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
1030
+ computed, num_computed, _ = manager.get_computed_blocks(resumed)
1031
+ assert num_computed == 3 * HISPARSE_BLOCK_SIZE
1032
+ assert manager.allocate_slots(
1033
+ resumed,
1034
+ num_new_tokens=len(tokens) - num_computed,
1035
+ num_new_computed_tokens=num_computed,
1036
+ new_computed_blocks=computed,
1037
+ num_external_computed_tokens=0,
1038
+ )
1039
+
1040
+
1041
+def test_hisparse_finished_request_leaves_free_copies():
1042
+ """Clean pages of a finished request are free-queue blocks, not pins."""
1043
+ manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
1044
+ tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
1045
+ original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
1046
+ assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None
1047
+ _publish_hisparse_pages(manager)
1048
+ manager.free(original)
1049
+
1050
+ coordinator = get_hisparse_coordinator(manager)
1051
+ assert coordinator.copies
1052
+ pool = manager.block_pool
1053
+ assert pool.get_num_free_blocks() == pool.num_gpu_blocks - 1
1054
+ pool.get_new_blocks(pool.get_num_free_blocks())
1055
+ assert not coordinator.copies
1056
+ assert not coordinator.spills_to_send
1057
+
1058
+
1059
+def test_hisparse_reset_prefix_cache_drops_copies():
1060
+ """Reset must forget GPU copies along with the host hashes they mirror."""
1061
+ manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
1062
+ tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
1063
+ request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256)
1064
+ assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None
1065
+ _publish_hisparse_pages(manager)
1066
+ manager.free(request)
1067
+
1068
+ coordinator = get_hisparse_coordinator(manager)
1069
+ assert coordinator.copies
1070
+ assert manager.reset_prefix_cache()
1071
+ assert not coordinator.copies
1072
+
1073
+
1074
+def test_hisparse_reused_copy_is_not_adopted():
1075
+ """A host prefix hit must not adopt a GPU copy the pool has handed out."""
1076
+ manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
1077
+ tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
1078
+ original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
1079
+ assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None
1080
+ _publish_hisparse_pages(manager)
1081
+ indexer_blocks = list(manager.get_blocks("original").blocks[1])
1082
+ manager.free(original)
1083
+ pool = manager.block_pool
1084
+ pool.touch(indexer_blocks)
1085
+ reused = pool.get_new_blocks(pool.get_num_free_blocks())
1086
+ pool.free_blocks(reused)
1087
+ pool.free_blocks(indexer_blocks)
1088
+
1089
+ resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
1090
+ computed, num_computed, _ = manager.get_computed_blocks(resumed)
1091
+ assert num_computed == 3 * HISPARSE_BLOCK_SIZE
1092
+ assert manager.allocate_slots(
1093
+ resumed,
1094
+ num_new_tokens=len(tokens) - num_computed,
1095
+ num_new_computed_tokens=num_computed,
1096
+ new_computed_blocks=computed,
1097
+ )
1098
+ resident_blocks = manager.get_blocks("resumed").blocks[2]
1099
+ assert all(block.is_null for block in resident_blocks[:3])
1100
+
1101
+
1102
+def test_hisparse_external_import_uses_hard_gpu_footprint():
1103
+ """An external prefix larger than resident capacity must remain admissible.
1104
+
1105
+ This guards the admission boundary where indexer + one writable resident
1106
+ page + the fixed hot allocation fit, but indexer + full resident history do
1107
+ not. Requiring one more resident page would silently strand long P/D inputs.
1108
+ """
1109
+ num_prompt_blocks = 4
1110
+ num_prompt_tokens = num_prompt_blocks * HISPARSE_BLOCK_SIZE
1111
+ manager = make_hisparse_kv_cache_manager(
1112
+ 8,
1113
+ 9,
1114
+ transfer_device_cache=True,
1115
+ )
1116
+ request = make_request(
1117
+ "host-import",
1118
+ list(range(num_prompt_tokens)),
1119
+ HISPARSE_BLOCK_SIZE,
1120
+ sha256,
1121
+ )
1122
+
1123
+ allocated = allocate_external_prefix(manager, request, num_prompt_tokens)
1124
+
1125
+ assert allocated is not None
1126
+ source, indexer, resident, hot = manager.get_blocks(request.request_id).blocks
1127
+ assert len(source) == num_prompt_blocks
1128
+ assert len(indexer) == num_prompt_blocks
1129
+ assert len(resident) == num_prompt_blocks
1130
+ assert all(block.is_null for block in resident)
1131
+ assert len(hot) == 2
1132
+ assert manager.block_pool.get_num_free_blocks() == 1
1133
+
1134
+
1135
+def test_hisparse_external_import_survives_capacity_retry():
1136
+ manager = make_hisparse_kv_cache_manager(
1137
+ 9,
1138
+ 7,
1139
+ transfer_device_cache=True,
1140
+ )
1141
+ tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
1142
+ first = make_request("first", tokens, HISPARSE_BLOCK_SIZE, sha256)
1143
+ second = make_request("second", tokens, HISPARSE_BLOCK_SIZE, sha256)
1144
+
1145
+ assert allocate_external_prefix(manager, first, len(tokens)) is not None
1146
+
1147
+ assert allocate_external_prefix(manager, second, len(tokens)) is None
1148
+
1149
+ manager.free(first)
1150
+ assert allocate_external_prefix(manager, second, len(tokens)) is not None
1151
+ _, _, resident, hot = manager.get_blocks(second.request_id).blocks
1152
+ assert all(block.is_null for block in resident)
1153
+ assert len(hot) == 2
1154
+
1155
+
1156
def make_kv_cache_config_hybrid_model(
1157
block_size: int,
1158
num_blocks: int,
@@ -3874,6 +4888,42 @@ def test_can_fit_full_sequence_swa_cap_admits_long_prompt():
4888
)
4889
4890
4891
+@pytest.mark.parametrize(("reserved_blocks", "fits"), [(4, True), (10, False)])
4892
+def test_full_sequence_admission_reserves_device_blocks_for_immediate_chunk(
4893
+ reserved_blocks, fits
4894
+):
4895
+ """Reservations constrain the current allocation, not the entire prompt."""
4896
+ block_size = 16
4897
+ config = KVCacheConfig(
4898
+ num_blocks=11, # Ten usable blocks after the null block.
4899
+ kv_cache_tensors=[],
4900
+ kv_cache_groups=[
4901
+ KVCacheGroupSpec(
4902
+ ["layer"],
4903
+ FullAttentionSpec(
4904
+ block_size=block_size,
4905
+ num_kv_heads=1,
4906
+ head_size=1,
4907
+ dtype=torch.float32,
4908
+ ),
4909
+ )
4910
+ ],
4911
+ )
4912
+ manager = make_kv_cache_manager(
4913
+ config, max_model_len=128, enable_caching=True, hash_block_size=block_size
4914
+ )
4915
+ request = make_request("chunk", list(range(128)), block_size, sha256)
4916
+
4917
+ blocks = manager.allocate_slots(
4918
+ request,
4919
+ block_size,
4920
+ full_sequence_must_fit=True,
4921
+ reserved_blocks=reserved_blocks,
4922
+ )
4923
+
4924
+ assert (blocks is not None) is fits
4925
+
4926
+
4927
def test_can_fit_full_sequence_full_attention_still_gates_oversized():
4928
"""The cap only loosens the SWA group; a prompt that exceeds the
4929
full-attention pool capacity must still be rejected."""
@@ -3926,6 +4976,30 @@ def test_can_fit_full_sequence_full_attention_still_gates_oversized():
4976
assert manager.allocate_slots(req, block_size, full_sequence_must_fit=True) is None
4977
4978
4979
+def test_can_fit_full_sequence_hisparse_caps_resident_pages():
4980
+ manager = make_hisparse_kv_cache_manager(
4981
+ num_blocks=11,
4982
+ host_num_blocks=9,
4983
+ max_model_len=8 * HISPARSE_BLOCK_SIZE,
4984
+ max_in_flight_tokens=2 * HISPARSE_BLOCK_SIZE,
4985
+ )
4986
+ request = make_request(
4987
+ "long-hisparse",
4988
+ list(range(8 * HISPARSE_BLOCK_SIZE)),
4989
+ HISPARSE_BLOCK_SIZE,
4990
+ sha256,
4991
+ )
4992
+
4993
+ assert (
4994
+ manager.allocate_slots(
4995
+ request,
4996
+ num_new_tokens=HISPARSE_BLOCK_SIZE,
4997
+ full_sequence_must_fit=True,
4998
+ )
4999
+ is not None
5000
+ )
5001
+
5002
+
5003
def test_cache_hit_local_and_external():
5004
# Regression test for #33775: when a request hits the local prefix cache
5005
# 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():
5845
assert last_req_hit(retention=0, pin=False) == 0
5846
# retention=0 with the pin keeps the junction window -> reuse restored.
5847
assert last_req_hit(retention=0, pin=True) == 4 * block_size
5848
+
5849
+
5850
+def test_device_eviction_keeps_host_prefix():
5851
+ """Device eviction must not target a same-numbered host block."""
5852
+ manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
5853
+ request = make_request("producer", list(range(64)), HISPARSE_BLOCK_SIZE, sha256)
5854
+ assert manager.allocate_slots(request, 64) is not None
5855
+ _publish_hisparse_pages(manager)
5856
+ host, indexer, _, _ = manager.get_blocks(request.request_id).blocks
5857
+ device_block = indexer[0]
5858
+ host_block = host[0]
5859
+ assert device_block.block_id == host_block.block_id
5860
+ manager.free(request)
5861
+ assert host_block.block_hash is not None and device_block.block_hash is not None
5862
+ manager.evict_blocks({device_block.block_id})
5863
+ assert device_block.block_hash is None, (
5864
+ "Device eviction left the indexer prefix cached"
5865
+ )
5866
+ assert host_block.block_hash is not None, (
5867
+ "Device eviction invalidated unrelated host KV"
5868
+ )
tests/v1/core/test_scheduler.py
+69
@@ -19,7 +19,12 @@ from vllm.config import (
19
SpeculativeConfig,
20
VllmConfig,
21
)
22
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
23
+ HiSparseConnector,
24
+ HiSparseConnectorScheduler,
25
+)
26
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
27
+from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector
28
from vllm.multimodal.inputs import (
29
MultiModalFeatureSpec,
30
MultiModalKwargsItem,
@@ -34,6 +39,7 @@ from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
39
from vllm.v1.core.sched.scheduler import Scheduler
40
from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs
41
from vllm.v1.engine import FinishReason
42
+from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
43
from vllm.v1.kv_cache_interface import (
44
FullAttentionSpec,
45
KVCacheConfig,
@@ -148,6 +154,30 @@ def test_get_num_unfinished_requests():
154
assert scheduler.get_num_unfinished_requests() == len(requests) - i - 1
155
156
157
+def _bind_hisparse_connector(scheduler):
158
+ connector = object.__new__(HiSparseConnector)
159
+ connector.connector_scheduler = HiSparseConnectorScheduler(async_speculative=False)
160
+ from tests.v1.core.test_prefix_caching import make_hisparse_kv_cache_manager
161
+
162
+ coordinator = get_hisparse_coordinator(make_hisparse_kv_cache_manager(16, 16))
163
+ connector.connector_scheduler.bind_coordinator(coordinator)
164
+ composite = object.__new__(MultiConnector)
165
+ composite._connectors = (connector,)
166
+ scheduler.connector = composite
167
+ return coordinator
168
+
169
+
170
+def test_pending_hisparse_spill_keeps_scheduler_alive():
171
+ """A final host spill must complete after the last request finishes."""
172
+ scheduler = create_scheduler()
173
+ coordinator = _bind_hisparse_connector(scheduler)
174
+ pending = coordinator.pending_spills
175
+ pending[0] = Mock()
176
+ assert scheduler.has_requests()
177
+ pending.clear()
178
+ assert not scheduler.has_requests()
179
+
180
+
181
@pytest.mark.parametrize(
182
"enable_prefix_caching, prompt_logprobs",
183
[
@@ -1181,6 +1211,45 @@ def test_preempt_during_execution():
1211
assert requests[1].output_token_ids[0] == 42
1212
1213
1214
+def test_pending_hisparse_reclamation_defers_preemption(monkeypatch):
1215
+ """Pending reclamation must stall rather than preempt its owning request."""
1216
+ scheduler = create_scheduler(enable_prefix_caching=False)
1217
+ request = create_requests(num_requests=1, num_tokens=16, block_size=16)[0]
1218
+ scheduler.add_request(request)
1219
+ scheduler_output = scheduler.schedule()
1220
+ scheduler.update_from_output(
1221
+ scheduler_output,
1222
+ ModelRunnerOutput(
1223
+ req_ids=[request.request_id],
1224
+ req_id_to_index={request.request_id: 0},
1225
+ sampled_token_ids=[[0]],
1226
+ logprobs=None,
1227
+ prompt_logprobs_dict={},
1228
+ pooler_output=[],
1229
+ ),
1230
+ )
1231
+
1232
+ monkeypatch.setattr(
1233
+ scheduler.kv_cache_manager, "allocate_slots", Mock(return_value=None)
1234
+ )
1235
+ coordinator = _bind_hisparse_connector(scheduler)
1236
+ monkeypatch.setattr(
1237
+ scheduler.connector, "build_connector_meta", Mock(return_value=None)
1238
+ )
1239
+ monkeypatch.setattr(
1240
+ coordinator,
1241
+ "has_pending_reclamation",
1242
+ Mock(return_value=True),
1243
+ )
1244
+
1245
+ deferred_output = scheduler.schedule()
1246
+
1247
+ assert not deferred_output.num_scheduled_tokens
1248
+ assert scheduler.running == [request]
1249
+ assert request.status == RequestStatus.RUNNING
1250
+ assert request.num_preemptions == 0
1251
+
1252
+
1253
def test_prefix_cache_query_not_inflated_by_connector_defer():
1254
"""The GPU prefix-cache query is recorded at admission, so a request the
1255
connector defers several times is counted once, not once per retry."""
tests/v1/e2e/general/test_hisparse.py
new
+225
@@ -0,0 +1,225 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import prometheus_client
5
+import pytest
6
+import torch
7
+
8
+from tests.conftest import VllmRunner
9
+from tests.utils import fork_new_process_for_each_test
10
+from vllm.config import AttentionConfig, HiSparseConfig, KVTransferConfig
11
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
12
+ HiSparseConnector,
13
+)
14
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
15
+ HiSparseConnectorWorker,
16
+)
17
+from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
18
+ MultiConnector,
19
+)
20
+from vllm.platforms import current_platform
21
+from vllm.utils.gpu_sync_debug import gpu_sync_allowed
22
+from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
23
+
24
+MODEL = "deepseek-ai/DeepSeek-V3.2"
25
+
26
+HF_OVERRIDES = {
27
+ "num_hidden_layers": 8,
28
+ "hidden_size": 256,
29
+ "intermediate_size": 512,
30
+ "num_attention_heads": 8,
31
+ "num_key_value_heads": 1,
32
+ "n_routed_experts": 8,
33
+ "num_experts_per_tok": 2,
34
+ "index_topk": 128,
35
+}
36
+
37
+
38
+def _shrink_config(config):
39
+ overrides = HF_OVERRIDES.copy()
40
+ if any("MTP" in architecture for architecture in config.architectures):
41
+ overrides.pop("num_hidden_layers")
42
+ config.update(overrides)
43
+ return config
44
+
45
+
46
+def _num_hisparse_spills(runner: VllmRunner) -> int:
47
+ client = runner.llm.llm_engine.engine_core
48
+ coordinator = get_hisparse_coordinator(
49
+ client.engine_core.scheduler.kv_cache_manager
50
+ )
51
+ return coordinator.next_spill_id
52
+
53
+
54
+def _offload_load_bytes() -> float:
55
+ return sum(
56
+ sample.value
57
+ for metric in prometheus_client.REGISTRY.collect()
58
+ if metric.name == "vllm:kv_offload_load_bytes"
59
+ for sample in metric.samples
60
+ if sample.name == "vllm:kv_offload_load_bytes_total"
61
+ )
62
+
63
+
64
+def _get_hisparse_worker(runner: VllmRunner) -> HiSparseConnectorWorker:
65
+ engine_core = runner.llm.llm_engine.engine_core.engine_core
66
+ model_runner = engine_core.model_executor.driver_worker.worker.model_runner
67
+ connector = model_runner.kv_connector.kv_connector
68
+ if isinstance(connector, MultiConnector):
69
+ connector = next(
70
+ child
71
+ for child in connector.sub_connectors
72
+ if isinstance(child, HiSparseConnector)
73
+ )
74
+ assert isinstance(connector, HiSparseConnector)
75
+ assert connector.connector_worker is not None
76
+ return connector.connector_worker
77
+
78
+
79
+@pytest.mark.skipif(
80
+ not current_platform.is_cuda(), reason="HiSparse requires NVIDIA CUDA"
81
+)
82
+@pytest.mark.parametrize(
83
+ "with_offloading", [False, True], ids=["standalone", "offload"]
84
+)
85
+@fork_new_process_for_each_test
86
+def test_hisparse_spill_and_prefix_restore(
87
+ monkeypatch: pytest.MonkeyPatch,
88
+ vllm_runner: type[VllmRunner],
89
+ with_offloading: bool,
90
+):
91
+ """Spilled prefixes restore and FULL-graph decode writes reach host KV.
92
+
93
+ A regression omitted attention metadata before FULL graph replay, so decode
94
+ completed normally while its newly written KV rows were never copied to host.
95
+ """
96
+ capability = current_platform.get_device_capability()
97
+ if capability is None or capability.major < 9:
98
+ pytest.skip("Sparse MLA requires Hopper or newer")
99
+
100
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
101
+ monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip")
102
+
103
+ target = [1000 + i % 64 for i in range(257)]
104
+ pressure = [
105
+ [2000 + request_idx * 128 + i % 64 for i in range(257)]
106
+ for request_idx in range(4)
107
+ ]
108
+ hisparse_connector = {
109
+ "kv_connector": "HiSparseConnector",
110
+ "kv_role": "kv_both",
111
+ "kv_connector_extra_config": {"host_pool_gib": 1},
112
+ }
113
+ kv_transfer_config = KVTransferConfig(**hisparse_connector)
114
+ if with_offloading:
115
+ kv_transfer_config = KVTransferConfig(
116
+ kv_connector="MultiConnector",
117
+ kv_role="kv_both",
118
+ kv_connector_extra_config={
119
+ "connectors": [
120
+ hisparse_connector,
121
+ {
122
+ "kv_connector": "OffloadingConnector",
123
+ "kv_role": "kv_both",
124
+ "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30},
125
+ },
126
+ ]
127
+ },
128
+ )
129
+
130
+ with vllm_runner(
131
+ MODEL,
132
+ load_format="dummy",
133
+ hf_overrides=_shrink_config,
134
+ attention_config=AttentionConfig(
135
+ hisparse_config=HiSparseConfig(
136
+ device_buffer_size=512,
137
+ )
138
+ ),
139
+ kv_transfer_config=kv_transfer_config,
140
+ block_size=64,
141
+ max_model_len=320,
142
+ max_num_batched_tokens=1024,
143
+ max_num_seqs=4,
144
+ num_gpu_blocks_override=128,
145
+ disable_log_stats=False,
146
+ enable_chunked_prefill=True,
147
+ enable_prefix_caching=True,
148
+ speculative_config={"method": "mtp", "num_speculative_tokens": 3},
149
+ compilation_config={
150
+ "cudagraph_mode": "FULL_AND_PIECEWISE",
151
+ "cudagraph_capture_sizes": [1, 4],
152
+ },
153
+ ) as runner:
154
+ worker = _get_hisparse_worker(runner)
155
+ engine_core = runner.llm.llm_engine.engine_core.engine_core
156
+ model_runner = engine_core.model_executor.driver_worker.worker.model_runner
157
+ full_graph_calls = 0
158
+ full_replay_pending = False
159
+ original_run_fullgraph = model_runner.cudagraph_manager.run_fullgraph
160
+
161
+ def record_fullgraph(batch_desc):
162
+ nonlocal full_graph_calls, full_replay_pending
163
+ full_graph_calls += 1
164
+ full_replay_pending = True
165
+ return original_run_fullgraph(batch_desc)
166
+
167
+ monkeypatch.setattr(
168
+ model_runner.cudagraph_manager, "run_fullgraph", record_fullgraph
169
+ )
170
+
171
+ original_finish_forward = worker.finish_forward
172
+ verified_full_replay_rows = 0
173
+ copied_nonzero_kv = False
174
+
175
+ def finish_forward():
176
+ nonlocal full_replay_pending, verified_full_replay_rows
177
+ nonlocal copied_nonzero_kv
178
+ try:
179
+ original_finish_forward()
180
+ if not full_replay_pending:
181
+ return
182
+ # Synchronize only the test's inspection, not finish_forward itself.
183
+ with gpu_sync_allowed():
184
+ torch.accelerator.synchronize()
185
+ assert worker._row_mirrors
186
+ for layer_index, cache in enumerate(worker.cache_handles):
187
+ source_index = cache.runtime.resident_source_index
188
+ for mirror in worker._row_mirrors:
189
+ source_row = mirror.source_starts[source_index]
190
+ destination_row = mirror.destination_start
191
+ num_rows = mirror.num_rows
192
+ source_block, source_offset = divmod(
193
+ source_row, worker.kernel_block_size
194
+ )
195
+ with gpu_sync_allowed():
196
+ gpu_rows = worker.resident_caches[layer_index][
197
+ source_block, source_offset : source_offset + num_rows
198
+ ].cpu()
199
+ host_rows = worker.host_caches[layer_index][
200
+ destination_row : destination_row + num_rows
201
+ ]
202
+ assert torch.equal(host_rows, gpu_rows)
203
+ copied_nonzero_kv |= bool(torch.count_nonzero(gpu_rows).item())
204
+ verified_full_replay_rows += num_rows
205
+ finally:
206
+ full_replay_pending = False
207
+
208
+ monkeypatch.setattr(worker, "finish_forward", finish_forward)
209
+
210
+ expected = runner.generate_greedy([target], max_tokens=8)
211
+
212
+ assert full_graph_calls > 0
213
+ assert verified_full_replay_rows > 0
214
+ assert copied_nonzero_kv
215
+
216
+ runner.generate_greedy(pressure, max_tokens=8)
217
+
218
+ assert _num_hisparse_spills(runner) > 0
219
+ load_bytes = _offload_load_bytes()
220
+ actual = runner.generate_greedy([target], max_tokens=8)
221
+ runner.generate_greedy([[42]], max_tokens=1)
222
+
223
+ assert actual == expected
224
+ if with_offloading:
225
+ assert _offload_load_bytes() > load_bytes
tests/v1/kv_connector/unit/offloading_connector/test_config.py
+31
-1
@@ -21,6 +21,7 @@ from vllm.v1.kv_cache_interface import (
21
FullAttentionSpec,
22
HiddenStateCacheSpec,
23
KVCacheConfig,
24
+ KVCacheGroupRole,
25
KVCacheGroupSpec,
26
KVCacheSpec,
27
KVCacheTensor,
@@ -336,6 +337,36 @@ def test_worker_kv_bytes_preserves_tensor_layout(packed: bool):
337
assert offloading_config.cache.blocks_per_chunk == 2
338
339
340
+def test_hisparse_offloads_only_indexer_group():
341
+ source = KVCacheGroupSpec(
342
+ ["source"],
343
+ _full_attention_spec(),
344
+ host_resident=True,
345
+ )
346
+ indexer = KVCacheGroupSpec(
347
+ ["indexer"],
348
+ _full_attention_spec(),
349
+ role=KVCacheGroupRole.HISPARSE_INDEXER,
350
+ )
351
+ kv_cache_config = KVCacheConfig(
352
+ num_blocks=4,
353
+ kv_cache_tensors=[],
354
+ kv_cache_groups=[source, indexer],
355
+ hisparse_host_num_blocks=4,
356
+ )
357
+ config = _make_vllm_config()
358
+
359
+ offloading_config = build_offloading_config(config, kv_cache_config)
360
+ scheduler_config = SchedulerOffloadConfig.from_spec(
361
+ MockOffloadingSpec(offloading_config), config, kv_cache_config
362
+ )
363
+
364
+ assert [
365
+ (group.group_id, group.layer_names) for group in offloading_config.groups
366
+ ] == [(1, ("indexer",))]
367
+ assert [group.group_idx for group in scheduler_config.kv_group_configs] == [1]
368
+
369
+
370
def test_zero_blocks_skips_tensor_layout_validation():
371
kv_cache_config = _make_sizing_kv_cache_config(packed=False)
372
kv_cache_config.num_blocks = 0
@@ -792,7 +823,6 @@ def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]):
823
False,
824
id="uniform-uncertifiable-inner",
825
),
795
- pytest.param([], False, id="attention-free"),
826
],
827
)
828
def test_canonical_layout_gate(kv_cache_groups, certified):
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+45
-2
@@ -98,6 +98,14 @@ def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state):
98
)
99
manager = CPUOffloadingManager(num_chunks=100)
100
spec = SimpleNamespace(
101
+ config=SimpleNamespace(
102
+ groups=tuple(
103
+ SimpleNamespace(
104
+ group_id=i, tokens_per_block=group.kv_cache_spec.block_size
105
+ )
106
+ for i, group in enumerate(groups)
107
+ )
108
+ ),
109
tokens_per_block=(256, 64, 8),
110
tokens_per_hash=8,
111
blocks_per_chunk=4,
@@ -303,7 +311,24 @@ def test_aligned_boundary_store_uses_exact_source_with_partial_tail():
311
assert src_spec.block_indices == [0, 0]
312
313
306
-def test_aligned_boundary_store_flushes_before_cow_destination_reuse():
314
+def test_aligned_boundary_store_maps_sparse_group_id_to_dense_transfer_slot():
315
+ scheduler = _make_partial_tail_scheduler()
316
+ indexer_config = scheduler.config.kv_group_configs[1]
317
+ scheduler.config = scheduler.config._replace(kv_group_configs=(indexer_config,))
318
+ _make_partial_tail_request(scheduler)
319
+ scheduler.manager.prepare_store.side_effect = lambda keys, req_context: (
320
+ generate_store_output(keys)
321
+ )
322
+
323
+ jobs = scheduler._build_aligned_boundary_store_jobs({"req": [(1, 98, 16)]})
324
+
325
+ [job] = jobs.values()
326
+ assert job.src_spec.group_sizes == [1]
327
+ assert job.src_spec.block_indices == [0]
328
+
329
+
330
+@pytest.mark.parametrize("cow_reuse", [False, True])
331
+def test_aligned_boundary_store_flushes_before_block_reuse(cow_reuse):
332
scheduler = _make_partial_tail_scheduler()
333
_make_partial_tail_request(scheduler)
334
scheduler.manager.prepare_store.side_effect = lambda keys, req_context: (
@@ -320,7 +345,11 @@ def test_aligned_boundary_store_flushes_before_cow_destination_reuse():
345
[job_id] = meta.store_jobs
346
347
output = SchedulerOutput.make_empty()
323
- output.kv_cache_block_copies = [KVCacheBlockCopy(98, 99)]
348
+ if cow_reuse:
349
+ output.kv_cache_block_copies = [KVCacheBlockCopy(98, 99)]
350
+ else:
351
+ output.scheduled_cached_reqs.req_ids = ["req"]
352
+ output.scheduled_cached_reqs.new_block_ids = [([], [99])]
353
meta = scheduler.build_connector_meta(output)
354
355
assert meta.jobs_to_flush == {job_id}
@@ -384,6 +413,20 @@ def test_partial_lookup_returns_exact_boundary_and_group_load_keys():
413
assert req_status.partial_tail_boundary is None
414
415
416
+def test_lookup_cap_stops_at_authoritative_prefix_boundary():
417
+ scheduler = _make_partial_tail_scheduler()
418
+ request = _make_partial_tail_request(scheduler)
419
+ request.skip_reading_prefix_cache = False
420
+ scheduler.manager.lookup.return_value = LookupResult.HIT
421
+
422
+ tokens, load_async = scheduler.get_num_new_matched_tokens(
423
+ request, 0, max_num_new_tokens=20
424
+ )
425
+
426
+ assert (tokens, load_async) == (20, True)
427
+ assert scheduler._req_status["req"].partial_tail_boundary == 20
428
+
429
+
430
def test_recurrent_group_unhashed_block_does_not_truncate_load_boundary():
431
"""A recurrent group may hold unhashed blocks inside the computed prefix.
432
tests/v1/kv_connector/unit/test_hisparse_connector.py
new
+116
@@ -0,0 +1,116 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from types import SimpleNamespace
4
+from unittest.mock import MagicMock
5
+
6
+import pytest
7
+import torch
8
+
9
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
10
+from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import (
11
+ ExampleConnector,
12
+)
13
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
14
+ HiSparseConnector,
15
+ HiSparseConnectorMetadata,
16
+ HiSparseConnectorScheduler,
17
+)
18
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
19
+ HiSparseConnectorWorker,
20
+)
21
+from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector
22
+from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
23
+from vllm.v1.hisparse.runtime import HiSparseCacheHandle
24
+from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector
25
+
26
+
27
+@pytest.mark.parametrize("nested", [False, True])
28
+def test_cache_manager_binding_preserves_hisparse_and_legacy_pool_hooks(nested):
29
+ """Composite binding must reach HiSparse and existing pool-only connectors."""
30
+ hisparse = object.__new__(HiSparseConnector)
31
+ hisparse._role = KVConnectorRole.SCHEDULER
32
+ hisparse.connector_scheduler = HiSparseConnectorScheduler(async_speculative=False)
33
+ legacy = object.__new__(ExampleConnector)
34
+ legacy.bind_gpu_block_pool = MagicMock()
35
+ connector = object.__new__(MultiConnector)
36
+ connector._connectors = [hisparse, legacy]
37
+ if nested:
38
+ parent = object.__new__(MultiConnector)
39
+ parent._connectors = [connector]
40
+ connector = parent
41
+ from tests.v1.core.test_prefix_caching import make_hisparse_kv_cache_manager
42
+
43
+ manager = make_hisparse_kv_cache_manager(16, 16)
44
+
45
+ connector.bind_kv_cache_manager(manager)
46
+
47
+ assert hisparse.connector_scheduler.coordinator is get_hisparse_coordinator(manager)
48
+ legacy.bind_gpu_block_pool.assert_called_once_with(manager.block_pool)
49
+
50
+
51
+def test_hisparse_requires_block_outermost_device_layout():
52
+ assert HiSparseConnector.get_required_kvcache_layout(MagicMock()) == "BLHNC"
53
+
54
+
55
+def test_no_forward_enqueues_deferred_hisparse_transfers():
56
+ """A zero-token step must still enqueue deferred post-forward transfers."""
57
+ connector = object.__new__(ActiveKVConnector)
58
+ connector._disabled = False
59
+ connector.pre_forward = MagicMock()
60
+ connector.finish_forward = MagicMock()
61
+ connector.post_forward = MagicMock(return_value=None)
62
+
63
+ scheduler_output = SimpleNamespace(finished_req_ids=set())
64
+ connector.no_forward(scheduler_output)
65
+
66
+ connector.pre_forward.assert_called_once_with(scheduler_output)
67
+ connector.finish_forward.assert_called_once_with()
68
+
69
+
70
+def test_full_graph_step_prepares_host_mirror_outside_model():
71
+ """Graph replay must restore host-mirror state cleared at step start."""
72
+ runtime = SimpleNamespace(
73
+ is_group_leader=True,
74
+ eager_host_mirror=True,
75
+ begin_forward=MagicMock(),
76
+ invalidate_written_slots=MagicMock(),
77
+ )
78
+ handle = HiSparseCacheHandle(runtime)
79
+ handle.mirror_slot_mapping = torch.tensor([4, 5])
80
+ worker = object.__new__(HiSparseConnectorWorker)
81
+ worker.cache_layer_names = ["layer"]
82
+ worker.cache_handles = [handle]
83
+ worker._group_leaders = (("layer", handle),)
84
+ worker._per_layer_mirrored = set()
85
+ worker._submitted_mirror_layers = set()
86
+ worker.is_host_writer = True
87
+ worker._enqueue_row_dma = MagicMock()
88
+ worker.start_step = MagicMock(
89
+ side_effect=lambda *_args, **_kwargs: worker._clear_forward_mirror_state()
90
+ )
91
+
92
+ connector = object.__new__(HiSparseConnector)
93
+ connector.connector_worker = worker
94
+ connector._get_connector_metadata = MagicMock(
95
+ return_value=HiSparseConnectorMetadata(None, (), (), {}, True)
96
+ )
97
+ req_id_per_token = torch.tensor([0, 1], dtype=torch.int32)
98
+ attn_metadata = SimpleNamespace(
99
+ num_actual_tokens=2,
100
+ num_decode_tokens=2,
101
+ num_reqs=2,
102
+ max_query_len=1,
103
+ req_id_per_token=req_id_per_token,
104
+ )
105
+
106
+ connector.start_load_kv(
107
+ SimpleNamespace(),
108
+ request_state_indices=None,
109
+ request_ids=[],
110
+ attn_metadata={"layer": attn_metadata},
111
+ )
112
+
113
+ worker._enqueue_host_mirror()
114
+
115
+ worker._enqueue_row_dma.assert_called_once_with((0,), ready_event=None)
116
+ runtime.invalidate_written_slots.assert_called_once()
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+85
@@ -3297,7 +3297,9 @@ def _make_bare_worker(
3297
# path; everything flows through the coordinator).
3298
from vllm.v1.kv_cache_interface import (
3299
FullAttentionSpec,
3300
+ KVCacheConfig,
3301
KVCacheGroupSpec,
3302
+ KVCacheTensor,
3303
)
3304
3305
worker.disk_offload_buffer_budget_bytes = None
@@ -3311,6 +3313,18 @@ def _make_bare_worker(
3313
block_size=block_size, num_kv_heads=8, head_size=64, dtype=None
3314
)
3315
group = KVCacheGroupSpec(["layer0", "__cross_layer__"], spec)
3316
+ worker._kv_cache_config = KVCacheConfig(
3317
+ num_blocks=num_gpu_blocks,
3318
+ kv_cache_tensors=[
3319
+ KVCacheTensor(
3320
+ size=0,
3321
+ layers=["layer0", "__cross_layer__"],
3322
+ layer_stride=0,
3323
+ block_stride=0,
3324
+ )
3325
+ ],
3326
+ kv_cache_groups=[group],
3327
+ )
3328
worker._kv_cache_groups = [group]
3329
worker.pcp_size = 1
3330
worker.dcp_size = 1
@@ -4059,6 +4073,77 @@ def test_register_kv_caches_shared_storage(layout: KVCacheLayout):
4073
worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes)
4074
4075
4076
+def test_register_kv_caches_uses_transfer_group_memory_domain():
4077
+ """Derived GPU caches must not change host-source transfer addresses."""
4078
+ from vllm.v1.kv_cache_interface import (
4079
+ KVCacheConfig,
4080
+ KVCacheGroupRole,
4081
+ KVCacheGroupSpec,
4082
+ KVCacheTensor,
4083
+ )
4084
+
4085
+ host_num_blocks = 4
4086
+ gpu_num_blocks = 2
4087
+ page_size = 32
4088
+ source_spec = FullAttentionSpec(
4089
+ block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float16
4090
+ )
4091
+ indexer_spec = FullAttentionSpec(
4092
+ block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float16
4093
+ )
4094
+ source_group = KVCacheGroupSpec(
4095
+ ["source"],
4096
+ source_spec,
4097
+ host_resident=True,
4098
+ )
4099
+ indexer_group = KVCacheGroupSpec(
4100
+ ["indexer"],
4101
+ indexer_spec,
4102
+ enable_kv_transfer=False,
4103
+ role=KVCacheGroupRole.HISPARSE_INDEXER,
4104
+ )
4105
+ config = KVCacheConfig(
4106
+ num_blocks=gpu_num_blocks,
4107
+ kv_cache_tensors=[
4108
+ KVCacheTensor(
4109
+ size=host_num_blocks * page_size,
4110
+ layers=["source"],
4111
+ layer_stride=host_num_blocks * page_size,
4112
+ block_stride=page_size,
4113
+ host_resident=True,
4114
+ ),
4115
+ KVCacheTensor(
4116
+ size=gpu_num_blocks * page_size,
4117
+ layers=["indexer"],
4118
+ layer_stride=gpu_num_blocks * page_size,
4119
+ block_stride=page_size,
4120
+ ),
4121
+ ],
4122
+ kv_cache_groups=[source_group, indexer_group],
4123
+ hisparse_host_num_blocks=host_num_blocks,
4124
+ )
4125
+ worker = _make_bare_worker(num_gpu_blocks=gpu_num_blocks)
4126
+ worker._kv_cache_config = config
4127
+ worker._kv_cache_groups = list(config.transfer_groups)
4128
+ worker.token_dbs = [
4129
+ ChunkedTokenDatabase(KeyMetadata("test-model", 0, 0, 0, 0), block_size=16)
4130
+ ]
4131
+ source = torch.zeros(host_num_blocks, page_size, dtype=torch.uint8)
4132
+ indexer = torch.zeros(gpu_num_blocks, page_size, dtype=torch.uint8)
4133
+
4134
+ _register_with_mocked_threads(
4135
+ worker,
4136
+ {"source": source, "indexer": indexer},
4137
+ )
4138
+
4139
+ db = worker.token_dbs[0]
4140
+ assert db.kv_caches_base_addr == [source.data_ptr()]
4141
+ assert db.block_len == [page_size]
4142
+ worker.store.register_buffer.assert_called_once_with(
4143
+ source.data_ptr(), source.nbytes
4144
+ )
4145
+
4146
+
4147
def test_register_kv_caches_separate_head_groups():
4148
# LHBNC gives each head group its own region; the K/V split doubles heads.
4149
layout = KVCacheLayout.LHBNC
tests/v1/kv_connector/unit/test_multi_connector.py
+6
-6
@@ -304,10 +304,10 @@ def test_multi_example_connector_consistency():
304
events = get_connector_events()
305
storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"])
306
storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"])
307
- # Initial events bind the block pool, query completion counts, and exchange
307
+ # Initial events bind the cache manager, query completion counts, and exchange
308
# handshake metadata before the request is enqueued.
309
assert storage1_scheduler_events[:7] == [
310
- "bind_gpu_block_pool",
310
+ "bind_kv_cache_manager",
311
"get_finished_count",
312
"set_xfer_handshake_metadata_pp_aware",
313
"on_new_request",
@@ -315,8 +315,8 @@ def test_multi_example_connector_consistency():
315
"update_state_after_alloc num_blocks=[7] 0",
316
"build_connector_meta",
317
]
318
- # First three events are from initialization. During generate(), layer hooks
319
- # run before the deferred load is started after the forward pass.
318
+ # First three events are from initialization. Layer hooks run before the
319
+ # deferred load starts after the forward pass.
320
expected_worker_prefix = [
321
"register_kv_caches",
322
"set_host_xfer_buffer_ops",
@@ -328,12 +328,12 @@ def test_multi_example_connector_consistency():
328
]
329
for connector_name in ("storage1-WORKER", "storage2-WORKER"):
330
worker_events = events[connector_name]
331
- assert worker_events[:7] == expected_worker_prefix
331
+ assert worker_events[: len(expected_worker_prefix)] == expected_worker_prefix
332
assert worker_events.index("start_load_kv") > worker_events.index(
333
"save_kv_layer"
334
)
335
assert storage2_scheduler_events[:7] == [
336
- "bind_gpu_block_pool",
336
+ "bind_kv_cache_manager",
337
"get_finished_count",
338
"set_xfer_handshake_metadata_pp_aware",
339
"on_new_request",
tests/v1/kv_connector/unit/test_nixl_connector.py
+140
-38
@@ -29,7 +29,10 @@ from vllm.distributed.kv_transfer.kv_connector.utils import (
29
get_current_attn_backend,
30
)
31
from vllm.distributed.kv_transfer.kv_connector.v1 import nixl
32
-from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
32
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
33
+ KVConnectorRole,
34
+ KVConnectorTransferResults,
35
+)
36
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
37
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
38
MultiKVConnectorStats,
@@ -518,6 +521,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
521
slot_size_bytes = 4096
522
self.slot_size_per_layer = [slot_size_bytes]
523
self.block_len_per_layer = [slot_size_bytes * self.block_size]
524
+ self.num_blocks = self.kv_cache_config.num_blocks
525
self.num_regions = 1
526
self.block_stride_per_layer = list(self.block_len_per_layer)
527
self.region_num_blocks = [self.num_blocks]
@@ -635,6 +639,12 @@ class TestNixlHandshake:
639
if pcp_rank > 0:
640
assert connector.get_finished(set()) == (set(), set())
641
642
+ worker.get_transfer_results = MagicMock(
643
+ return_value=KVConnectorTransferResults(finished_sending={"sent"})
644
+ )
645
+ results = connector.get_transfer_results(set())
646
+ assert results.finished_sending == ({"sent"} if pcp_rank == 0 else set())
647
+
648
@patch(
649
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
650
FakeNixlWrapper,
@@ -1195,6 +1205,7 @@ class TestNixlHandshake:
1205
ssm_sizes=(0, 0),
1206
attn_backend_name=worker.backend_name,
1207
physical_blocks_per_logical_kv_block=1,
1208
+ region_num_blocks=None,
1209
)
1210
1211
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):
1446
metadata = NixlConnectorMetadata()
1447
metadata.add_new_req_to_recv(
1448
request_id=request_id,
1438
- local_block_ids=([1, 2, 3],),
1449
+ local_block_ids=([0],),
1450
kv_transfer_params={
1440
- "remote_block_ids": ([4, 5, 6],),
1451
+ "remote_block_ids": ([0],),
1452
"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1453
"remote_request_id": f"prefill-{request_id}",
1454
"remote_host": "localhost",
@@ -1482,6 +1493,10 @@ def test_kv_connector_stats(default_vllm_config, dist_init):
1493
assert stats_after_reset is None
1494
1495
1496
+@patch(
1497
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1498
+ FakeNixlWrapper,
1499
+)
1500
def test_reqs_to_send_deadline_rebased_to_worker_clock(default_vllm_config, dist_init):
1501
"""reqs_to_send deadlines are stamped with the scheduler process's
1502
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():
1921
assert handle == 22
1922
assert worker._dram_src_handles_by_block_size[worker.block_size] == 11
1923
assert [block[2] for block in blocks] == [0, 0, 3, 3]
1909
- assert [
1924
+ memory_types = [
1925
call.args[1] for call in worker.nixl_wrapper.get_xfer_descs.call_args_list
1911
- ] == ["DRAM", "VRAM"]
1926
+ ]
1927
+ assert memory_types == ["DRAM", "VRAM"]
1928
1929
1930
def test_mixed_memory_read_notifies_after_both_transfers_finish():
@@ -1919,6 +1935,7 @@ def test_mixed_memory_read_notifies_after_both_transfers_finish():
1935
worker._recving_metadata = {"request": MagicMock()}
1936
worker._recving_transfers = defaultdict(list)
1937
worker._pending_recv_notifs = {}
1938
+ worker._failed_inflight_recvs = set()
1939
worker.xfer_stats = MagicMock()
1940
worker.nixl_wrapper = MagicMock()
1941
worker.nixl_wrapper.make_prepped_xfer.side_effect = [101, 102]
@@ -1959,6 +1976,7 @@ def test_mixed_memory_read_failure_does_not_notify_producer():
1976
worker._recving_metadata = {"request": MagicMock()}
1977
worker._recving_transfers = {"request": [101, 102]}
1978
worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]}
1979
+ worker._failed_inflight_recvs = set()
1980
worker._is_hma_required = True
1981
worker._failed_recv_reqs = MagicMock()
1982
worker._log_failure = MagicMock() # type: ignore[method-assign]
@@ -2871,9 +2889,11 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init):
2889
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2890
FailingNixlWrapper,
2891
)
2874
-def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init):
2875
- """Test that transfer setup failures mark blocks invalid
2876
- and return via get_finished."""
2892
+@pytest.mark.parametrize("is_hma", [False, True])
2893
+def test_transfer_setup_failure_returns_finished(
2894
+ default_vllm_config, dist_init, is_hma
2895
+):
2896
+ """Setup failures report the request; only non-HMA reports block IDs."""
2897
vllm_config = create_vllm_config()
2898
2899
connector = NixlConnector(
@@ -2883,6 +2903,7 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)
2903
vllm_config, connector.engine_id, hand_shake_latency=0
2904
)
2905
connector.connector_worker.nixl_wrapper.fail_transfer_setup = True
2906
+ connector.connector_worker._is_hma_required = is_hma
2907
2908
request_id = "test_transfer_fail"
2909
metadata = NixlConnectorMetadata()
@@ -2912,13 +2933,106 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)
2933
time.sleep(0.1)
2934
connector.start_load_kv(dummy_ctx)
2935
2915
- # check that blocks were marked invalid
2936
+ results = connector.get_transfer_results(finished_req_ids=set())
2937
+ assert request_id in results.finished_recving
2938
+ assert results.failed_recving == {request_id}
2939
invalid_blocks = connector.get_block_ids_with_load_errors()
2917
- assert invalid_blocks == {7, 8, 9}
2940
+ assert invalid_blocks == (set() if is_hma else {7, 8, 9})
2941
+
2942
+
2943
+class _ScriptedXferWrapper(FakeNixlWrapper):
2944
+ """Scripts per-handle xfer states; forbids releasing in-flight handles."""
2945
+
2946
+ def __init__(self, *args, **kwargs):
2947
+ super().__init__(*args, **kwargs)
2948
+ # handle -> successive states; the last one repeats.
2949
+ self.xfer_states: dict[int, list[str]] = {}
2950
+ self.released: list[int] = []
2951
+
2952
+ def check_xfer_state(self, handle: int) -> str:
2953
+ states = self.xfer_states[handle]
2954
+ return states.pop(0) if len(states) > 1 else states[0]
2955
+
2956
+ def release_xfer_handle(self, handle: int) -> None:
2957
+ # A posted-but-unfinished transfer cannot be aborted: releasing it
2958
+ # leaves the RDMA READ armed. Production code must never do this.
2959
+ assert self.xfer_states.get(handle, ["DONE"])[0] != "PROC", (
2960
+ f"released in-flight handle {handle}"
2961
+ )
2962
+ self.released.append(handle)
2963
+
2964
2919
- # ensure request appears in get_finished
2965
+def _make_split_read_connector(vllm_config, request_id, states):
2966
+ """Seed a request with scripted in-flight xfer handles + recv metadata."""
2967
+ connector = NixlConnector(
2968
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
2969
+ )
2970
+ connector.connector_worker = FakeNixlConnectorWorker(
2971
+ vllm_config, connector.engine_id, hand_shake_latency=0
2972
+ )
2973
+ worker = connector.connector_worker
2974
+ wrapper = _ScriptedXferWrapper("agent")
2975
+ worker.nixl_wrapper = wrapper
2976
+ metadata = NixlConnectorMetadata()
2977
+ metadata.add_new_req_to_recv(
2978
+ request_id=request_id,
2979
+ local_block_ids=([7, 8, 9],),
2980
+ kv_transfer_params={
2981
+ "remote_block_ids": ([10, 11, 12],),
2982
+ "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
2983
+ "remote_request_id": f"prefill-{request_id}",
2984
+ "remote_host": "localhost",
2985
+ "remote_port": 1234,
2986
+ "remote_tp_size": 1,
2987
+ },
2988
+ )
2989
+ worker._recving_metadata[request_id] = metadata.reqs_to_recv[request_id]
2990
+ wrapper.xfer_states = dict(states)
2991
+ worker._recving_transfers[request_id] = list(states)
2992
+ return connector, worker, wrapper
2993
+
2994
+
2995
+@patch(
2996
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2997
+ FakeNixlWrapper,
2998
+)
2999
+def test_split_read_failure_defers_report_until_last_handle(
3000
+ default_vllm_config, dist_init
3001
+):
3002
+ """One half of a split (mixed DRAM/VRAM) read failing must not report the
3003
+ request — nor invalidate its blocks — while the sibling xfer is still in
3004
+ flight: a posted READ cannot be aborted and would DMA into blocks the
3005
+ scheduler could free and reuse. The report happens exactly once, when the
3006
+ last handle is terminal."""
3007
+ request_id = "split_read_partial_failure"
3008
+ err_handle, live_handle = 11, 22
3009
+ connector, worker, wrapper = _make_split_read_connector(
3010
+ create_vllm_config(),
3011
+ request_id,
3012
+ {err_handle: ["ERR"], live_handle: ["PROC", "PROC", "DONE"]},
3013
+ )
3014
+
3015
+ # Poll 1: one half fails; the sibling is in flight -> nothing reported.
3016
_, done_recving = connector.get_finished(finished_req_ids=set())
2921
- assert request_id in done_recving
3017
+ assert done_recving == set()
3018
+ assert connector.get_block_ids_with_load_errors() == set()
3019
+ assert request_id in worker._recving_metadata
3020
+ assert wrapper.released == [err_handle]
3021
+
3022
+ # Poll 2: sibling still in flight.
3023
+ _, done_recving = connector.get_finished(finished_req_ids=set())
3024
+ assert done_recving == set()
3025
+
3026
+ # Poll 3: sibling terminal -> reported exactly once, blocks invalidated.
3027
+ _, done_recving = connector.get_finished(finished_req_ids=set())
3028
+ assert done_recving == {request_id}
3029
+ assert connector.get_block_ids_with_load_errors() == {7, 8, 9}
3030
+ assert request_id not in worker._recving_metadata
3031
+ assert wrapper.released == [err_handle, live_handle]
3032
+
3033
+ # Poll 4: nothing left; no double report.
3034
+ _, done_recving = connector.get_finished(finished_req_ids=set())
3035
+ assert done_recving == set()
3036
3037
3038
@patch(
@@ -2938,7 +3052,7 @@ def test_failed_request_skips_kv_postprocessing(
3052
default_vllm_config, dist_init, failure_mode
3053
):
3054
"""Test that failed requests skip KV sync and post-processing in
2941
- get_finished().
3055
+ get_transfer_results().
3056
3057
This is the core safety behavior: when a KV transfer fails at any stage,
3058
the request must still appear in done_recving (so the scheduler can apply
@@ -3023,11 +3137,10 @@ def test_failed_request_skips_kv_postprocessing(
3137
patch.object(worker, "sync_recved_kv_to_device") as mock_sync,
3138
patch.object(worker, "post_process_device_kv_on_receive") as mock_postprocess,
3139
):
3026
- _, done_recving = connector.get_finished(finished_req_ids=set())
3140
+ results = connector.get_transfer_results(finished_req_ids=set())
3141
3028
- # The failed request must appear in done_recving so the scheduler
3029
- # can handle it (e.g., trigger recompute via kv_load_failure_policy).
3030
- assert request_id in done_recving
3142
+ assert request_id in results.finished_recving
3143
+ assert results.failed_recving == {request_id}
3144
3145
# Critical: KV sync and post-processing must NOT have been called
3146
# since no valid KV data was received for the failed request.
@@ -3049,19 +3162,7 @@ def test_failed_request_skips_kv_postprocessing(
3162
def test_handles_failing_in_separate_polls_do_not_kill_the_engine(
3163
default_vllm_config, dist_init
3164
):
3052
- """A request whose handles fail in different polls must not crash the engine.
3053
-
3054
- One transfer handle is created per remote rank, and when a peer goes away
3055
- they do not all fail in the same poll: one errors while another is still
3056
- PROC. The first failure reports the request via _failed_recv_reqs and
3057
- get_finished() pops its metadata; when the remaining handles fail in a
3058
- later poll, the request must be cleaned up without being reported again.
3059
-
3060
- Reporting twice kills the EngineCore: the scheduler's assert in
3061
- _update_from_kv_xfer_finished only expects a finished recv for a request
3062
- still waiting for KVs, having moved this one out of
3063
- WAITING_FOR_REMOTE_KVS to recompute locally after the first report.
3064
- """
3165
+ """Failures across polls report once, after the last writer is terminal."""
3166
vllm_config = create_vllm_config()
3167
connector = NixlConnector(
3168
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(
3208
):
3209
_, done_recving = connector.get_finished(finished_req_ids=set())
3210
3110
- assert request_id in done_recving
3111
- assert request_id not in worker._recving_metadata
3211
+ assert request_id not in done_recving
3212
+ assert request_id in worker._recving_metadata
3213
assert worker._recving_transfers[request_id] == [second]
3214
+ assert connector.get_block_ids_with_load_errors() == set()
3215
3114
- # Poll 2: the second handle fails too. The request was already reported,
3115
- # so it must not be reported a second time: the scheduler has since moved
3116
- # it out of WAITING_FOR_REMOTE_KVS to recompute locally, and asserts on a
3117
- # finished recv for a request in that state. Its remaining handles are
3118
- # still released and the transfer entry removed.
3216
+ # Only the last terminal handle permits the scheduler to reuse the blocks.
3217
with patch.object(worker.nixl_wrapper, "check_xfer_state", return_value="ERR"):
3218
_, done_recving = connector.get_finished(finished_req_ids=set())
3219
3122
- assert request_id not in done_recving
3220
+ assert done_recving == {request_id}
3221
assert request_id not in worker._recving_transfers
3222
+ assert request_id not in worker._recving_metadata
3223
+ assert connector.get_block_ids_with_load_errors() == {1, 2, 3}
3224
+ _, done_recving = connector.get_finished(finished_req_ids=set())
3225
+ assert done_recving == set()
3226
3227
3228
def _set_test_speculative_config(
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+188
-28
@@ -5,16 +5,31 @@
5
import gc
6
from unittest.mock import MagicMock, patch
7
8
+import msgspec
9
import pytest
10
import torch
11
12
from tests.v1.attention.utils import MockMambaBuilder
13
from vllm import LLM, SamplingParams
13
-from vllm.config import KVTransferConfig
14
+from vllm.config import KVTransferConfig, set_current_vllm_config
15
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw
16
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
17
+ NixlAgentMetadata,
18
+)
19
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
20
+ NixlConnectorWorker,
21
+)
22
from vllm.v1.core.single_type_kv_cache_manager import (
23
FullAttentionManager,
24
SlidingWindowManager,
25
)
26
+from vllm.v1.kv_cache_interface import (
27
+ FullAttentionSpec,
28
+ KVCacheConfig,
29
+ KVCacheGroupSpec,
30
+ KVCacheTensor,
31
+ MLAAttentionSpec,
32
+)
33
34
from .utils import (
35
create_request,
@@ -272,26 +287,24 @@ def test_logical_to_kernel_block_ids_with_hma():
287
288
@pytest.mark.cpu_test
289
@pytest.mark.parametrize(
275
- "is_rocm,has_mamba,use_host_buffer,done_recving,failed_recving,expected_syncs",
290
+ "is_rocm,has_mamba,use_host_buffer,done_recving,expected_syncs",
291
[
277
- (True, True, False, {"req"}, set(), 1),
278
- (False, True, False, {"req"}, set(), 0),
279
- (True, False, False, {"req"}, set(), 0),
280
- (True, True, True, {"req"}, set(), 0),
281
- (True, True, False, set(), set(), 0),
282
- (True, True, False, {"req"}, {"req"}, 0),
292
+ (True, True, False, {"req"}, 1),
293
+ (False, False, False, {"req"}, 0),
294
+ (False, True, False, {"req"}, 0),
295
+ (True, True, True, {"req"}, 0),
296
+ (True, True, False, set(), 0),
297
],
298
)
285
-def test_sync_device_after_mamba_recv_gates(
299
+def test_sync_device_after_direct_recv_gates(
300
monkeypatch,
301
is_rocm,
302
has_mamba,
303
use_host_buffer,
304
done_recving,
291
- failed_recving,
305
expected_syncs,
306
):
294
- """Only direct-GPU Mamba receives on ROCm need a device fence."""
307
+ """Only backends requiring visibility fences synchronize direct receives."""
308
from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker
309
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
310
NixlConnectorWorker,
@@ -309,7 +322,7 @@ def test_sync_device_after_mamba_recv_gates(
322
lambda: sync_calls.append(True),
323
)
324
312
- worker._sync_device_after_mamba_recv(done_recving, failed_recving)
325
+ worker._sync_device_after_direct_recv(done_recving)
326
327
assert len(sync_calls) == expected_syncs
328
@@ -399,6 +412,8 @@ def test_read_blocks_for_req_expands_remote_ids(
412
)
413
414
remote_engine_id = "remote-engine"
415
+ worker.region_group_ids = list(range(len(resolved_types)))
416
+ worker.dst_region_group_ids = {remote_engine_id: list(worker.region_group_ids)}
417
418
worker.transfer_topo = MagicMock()
419
# tp_ratio not exercised (all_source_ranks is empty so no reads run),
@@ -888,6 +903,45 @@ def _make_mock_worker_for_desc_ids(
903
return worker
904
905
906
+@pytest.mark.cpu_test
907
+def test_shared_nixl_region_broadcasts_every_group_block_table():
908
+ """A shared region must receive every group's block table exactly once."""
909
+ worker = _make_mock_worker_for_desc_ids(
910
+ num_regions=3,
911
+ has_mamba=False,
912
+ group_spec_types=(FullAttentionSpec,) * 3,
913
+ block_len_per_layer=[100] * 3,
914
+ )
915
+ worker.region_group_ids = [0, bw._SHARED_REGION_GROUP_ID, 2]
916
+ worker._uses_region_group_mapping = True
917
+
918
+ result = worker._compute_desc_ids(
919
+ block_ids=([1], [2], [3]),
920
+ dst_num_blocks=10,
921
+ block_size_ratio=None,
922
+ physical_blocks_per_logical=1,
923
+ region_num_blocks=[10, 10, 10],
924
+ )
925
+
926
+ assert result.tolist() == [1, 11, 12, 13, 23]
927
+
928
+
929
+@pytest.mark.cpu_test
930
+def test_nixl_region_sort_preserves_pipeline_layer_order():
931
+ """Numeric layer order prevents PP peers from pairing different regions."""
932
+ names = [
933
+ "model.layers.10.self_attn",
934
+ "model.layers.2.self_attn",
935
+ "model.layers.1.self_attn",
936
+ ]
937
+
938
+ assert sorted(names, key=bw._region_sort_key) == [
939
+ "model.layers.1.self_attn",
940
+ "model.layers.2.self_attn",
941
+ "model.layers.10.self_attn",
942
+ ]
943
+
944
+
945
@pytest.mark.cpu_test
946
def test_get_block_descs_ids_hybrid_ssm():
947
"""Test _compute_desc_ids uses per-group strides for hybrid
@@ -1097,10 +1151,7 @@ def test_post_process_zeroes_untransferred_tail():
1151
worker.device_kv_caches = {"attn.0": attn_cache, "mamba.0": mamba_cache}
1152
fa_group = MagicMock(layer_names=["attn.0"])
1153
ssm_group = MagicMock(layer_names=["mamba.0"])
1100
- worker.kv_cache_config = MagicMock(
1101
- kv_cache_groups=[fa_group, ssm_group],
1102
- transfer_groups=[fa_group, ssm_group],
1103
- )
1154
+ worker.kv_cache_config = MagicMock(transfer_groups=[fa_group, ssm_group])
1155
# The cached property filters mamba layers out of the permuted caches.
1156
attn_caches = NixlConnectorWorker._attention_kv_caches.func(worker)
1157
assert len(attn_caches) == 1 and attn_caches[0] is attn_cache
@@ -1188,6 +1239,11 @@ def _make_fake_kv_cache_manager():
1239
1240
manager = object.__new__(KVCacheManager)
1241
manager.coordinator = MagicMock()
1242
+ manager.kv_cache_config = MagicMock()
1243
+ manager.kv_cache_config.kv_cache_groups = (
1244
+ MagicMock(host_resident=False),
1245
+ MagicMock(host_resident=False),
1246
+ )
1247
manager.coordinator.single_type_managers = (
1248
_FakeSingleTypeManager(True, 16, [10, 11, 12, 13, 14, 15]), # attention
1249
_FakeSingleTypeManager(False, 16, [20, 21, 22, 23, 24, 25]), # mamba
@@ -1723,16 +1779,6 @@ def test_exchange_clipped_blocks_ssm_positional_states():
1779
assert clipped == ([1, 2, 3], [0, 5, 6, 7])
1780
1781
1726
-@pytest.mark.cpu_test
1727
-def test_exchange_clipped_blocks_excludes_nontransfer_groups():
1728
- """Scheduler block tables must match the worker's registered regions."""
1729
- sched = make_nixl_scheduler()
1730
- sched.kv_cache_config = make_kv_cache_config(block_size=16, swa_enabled=True)
1731
- sched.kv_cache_config.kv_cache_groups[1].enable_kv_transfer = False
1732
-
1733
- assert sched.get_exchange_clipped_blocks(([1, 2], [9])) == ([1, 2],)
1734
-
1735
-
1782
# ── Hybrid MLA+SSM (KimiLinear-shaped KDA+MLA) tests ─────────────────────
1783
1784
@@ -1786,6 +1832,95 @@ def _make_hybrid_mla_kv_cache_config(num_blocks: int = 4):
1832
)
1833
1834
1835
+@pytest.mark.cpu_test
1836
+@pytest.mark.parametrize("kernel_block_size", [16, 8])
1837
+def test_nixl_keeps_device_block_count_with_hisparse_host_pool(kernel_block_size):
1838
+ host_num_blocks = 4
1839
+ gpu_num_blocks = 9
1840
+ spec = MLAAttentionSpec(
1841
+ block_size=16,
1842
+ num_kv_heads=1,
1843
+ head_size=8,
1844
+ dtype=torch.float16,
1845
+ )
1846
+ kv_cache_config = KVCacheConfig(
1847
+ num_blocks=gpu_num_blocks,
1848
+ hisparse_host_num_blocks=host_num_blocks,
1849
+ kv_cache_tensors=[
1850
+ KVCacheTensor(
1851
+ size=host_num_blocks * spec.page_size_bytes,
1852
+ layers=["mla.host"],
1853
+ layer_stride=host_num_blocks * spec.page_size_bytes,
1854
+ block_stride=spec.page_size_bytes,
1855
+ host_resident=True,
1856
+ ),
1857
+ KVCacheTensor(
1858
+ size=gpu_num_blocks * spec.page_size_bytes,
1859
+ layers=["mla.device"],
1860
+ layer_stride=gpu_num_blocks * spec.page_size_bytes,
1861
+ block_stride=spec.page_size_bytes,
1862
+ ),
1863
+ ],
1864
+ kv_cache_groups=[
1865
+ KVCacheGroupSpec(
1866
+ ["mla.host"],
1867
+ spec,
1868
+ host_resident=True,
1869
+ ),
1870
+ KVCacheGroupSpec(["mla.device"], spec),
1871
+ ],
1872
+ )
1873
+ vllm_config = create_vllm_config(block_size=16)
1874
+ vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
1875
+ fake_backend = MagicMock()
1876
+ fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size]
1877
+ fake_backend.get_name.return_value = "FLASHMLA"
1878
+ fake_backend.full_cls_name.return_value = "fake.FLASHMLA"
1879
+ fake_backend.get_kv_cache_shape.return_value = (1, kernel_block_size, 1, 1)
1880
+ fake_platform = MagicMock()
1881
+ fake_platform.device_type = "cuda"
1882
+ fake_platform.get_nixl_memory_type.return_value = "VRAM"
1883
+ physical_per_logical = spec.block_size // kernel_block_size
1884
+ physical_page_bytes = spec.page_size_bytes // physical_per_logical
1885
+ host_cache = torch.zeros(
1886
+ host_num_blocks * physical_per_logical, physical_page_bytes, dtype=torch.uint8
1887
+ )
1888
+ device_cache = torch.zeros(
1889
+ gpu_num_blocks * physical_per_logical, physical_page_bytes, dtype=torch.uint8
1890
+ )
1891
+
1892
+ with (
1893
+ patch.object(bw, "NixlWrapper"),
1894
+ patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),
1895
+ patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),
1896
+ patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]),
1897
+ patch.object(bw, "current_platform", fake_platform),
1898
+ set_current_vllm_config(vllm_config),
1899
+ ):
1900
+ worker = NixlConnectorWorker(vllm_config, "test-engine", kv_cache_config)
1901
+ worker.use_mla = True
1902
+ worker.nixl_wrapper.get_agent_metadata.return_value = b"metadata"
1903
+ worker.register_kv_caches({"mla.host": host_cache, "mla.device": device_cache})
1904
+
1905
+ assert worker.num_blocks == gpu_num_blocks * physical_per_logical
1906
+ assert dict(zip(worker.region_group_ids, worker.region_num_blocks)) == {
1907
+ 0: host_num_blocks * physical_per_logical,
1908
+ 1: gpu_num_blocks * physical_per_logical,
1909
+ }
1910
+ caches = {"mla.host": host_cache, "mla.device": device_cache}
1911
+ expected_addrs = [
1912
+ caches[name].data_ptr() + block * stride
1913
+ for name, stride, count in zip(
1914
+ worker.region_names,
1915
+ worker.block_stride_per_layer,
1916
+ worker.region_num_blocks,
1917
+ strict=True,
1918
+ )
1919
+ for block in range(count)
1920
+ ]
1921
+ assert worker.src_blocks_data[:, 0].tolist() == expected_addrs
1922
+
1923
+
1924
@pytest.mark.cpu_test
1925
def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
1926
"""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():
1938
kv_cache_config = _make_hybrid_mla_kv_cache_config()
1939
unified_page = kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes
1940
vllm_config = create_vllm_config(block_size=12)
1941
+ # The engine rewrites this global value to the smallest scheduler-group
1942
+ # block size. NIXL must retain the transferable groups' 12-token geometry.
1943
+ vllm_config.cache_config.block_size = 4
1944
# kv_buffer_device defaults to the *real* platform's device type, which on
1945
# a CPU-only test host would make this a host-buffer worker: host xfer
1946
# 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():
1971
worker.use_mla = True # opt-125m test config is not MLA; force the flag
1972
worker.nixl_wrapper.get_agent_metadata.return_value = b"fake-agent-metadata"
1973
1836
- tensors = [torch.zeros(4 * unified_page, dtype=torch.uint8) for _ in range(2)]
1974
+ kernel_page = unified_page // 3
1975
+ block_stride = 2 * kernel_page
1976
+ backing = torch.zeros(12, block_stride, dtype=torch.uint8)
1977
+ tensors = [
1978
+ backing[:, :kernel_page],
1979
+ backing[:, kernel_page:],
1980
+ ]
1981
# KDA layer first per tensor: exercises the dual-purpose flag merge.
1982
worker.register_kv_caches(
1983
{
@@ -1849,6 +1993,8 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
1993
# 12-token logical blocks over the 4-token MLA kernel block.
1994
assert worker._physical_blocks_per_logical_kv_block == 3
1995
assert worker.block_size == 4 and worker.num_blocks == 12
1996
+ assert worker.region_group_ids == [bw._SHARED_REGION_GROUP_ID] * 2
1997
+ assert worker.block_stride_per_layer == [block_stride, block_stride]
1998
# Both shared tensors are dual-purpose: their FA view is MLA even though
1999
# a KDA layer registered them first.
2000
assert worker._region_is_mla == [True, True]
@@ -1861,8 +2007,20 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
2007
# Mamba descs: 2 regions x (3 conv sub-projections + 1 ssm) x 4 blocks.
2008
assert worker.src_blocks_data.shape == (24 + 32, 3)
2009
fa_descs = worker.src_blocks_data[:24]
1864
- assert fa_descs[1][0] - fa_descs[0][0] == unified_page // 3
2010
+ expected_addrs = [
2011
+ tensor.data_ptr() + block * block_stride
2012
+ for tensor in tensors
2013
+ for block in range(12)
2014
+ ]
2015
+ assert fa_descs[:, 0].tolist() == expected_addrs
2016
assert all(size == unified_page // 3 for size in fa_descs[:, 1])
2017
+ worker.nixl_wrapper.register_memory.assert_called_once()
2018
+ metadata = msgspec.msgpack.decode(
2019
+ worker.xfer_handshake_metadata.agent_metadata_bytes,
2020
+ type=NixlAgentMetadata,
2021
+ )
2022
+ assert metadata.block_strides == [block_stride, block_stride]
2023
+ assert metadata.region_num_blocks == [12, 12]
2024
2025
2026
@pytest.mark.cpu_test
@@ -1909,6 +2067,8 @@ def test_push_write_hybrid_mla_replicates_attention():
2067
worker.dst_xfer_side_handles = {engine_id: {0: 100, 1: 101}}
2068
worker.src_xfer_handles_by_tp_ratio = {(-2, 4): [200, 201]}
2069
worker.src_xfer_handles_by_block_size = {4: 300}
2070
+ worker.region_group_ids = [0, 1]
2071
+ worker.dst_region_group_ids = {engine_id: [0, 1]}
2072
worker._sending_transfers = defaultdict(list)
2073
worker._sending_transfers_lock = threading.Lock()
2074
worker.kv_cache_config = _make_hybrid_mla_kv_cache_config()
tests/v1/kv_connector/unit/test_nixl_desc_geometry.py
+19
-8
@@ -186,11 +186,16 @@ def test_overlaid_transfer_groups_share_region_geometry():
186
worker.host_buffer_kv_cache_layout = "NHD"
187
worker._physical_blocks_per_logical_kv_block = 1
188
worker._logical_num_blocks = num_blocks
189
+ worker.region_mem_types = []
190
worker.region_group_ids = []
191
worker.region_mem_types = []
192
worker._mixed_mem_types = False
193
worker.region_names = []
194
worker.region_num_blocks = []
195
+ worker._mixed_mem_types = False
196
+ worker._desc_is_dram_by_block_size = {}
197
+ worker._desc_pos_by_block_size = {}
198
+ worker._dram_src_handles_by_block_size = {}
199
worker._region_is_mla = []
200
worker.block_len_per_layer = []
201
worker.block_stride_per_layer = []
@@ -304,8 +309,6 @@ def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blo
309
# would not be deduplicated. Pin it to the faked device type.
310
vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
311
307
- from unittest.mock import MagicMock
308
-
312
fake_backend = MagicMock()
313
fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size]
314
fake_backend.get_name.return_value = "FLASHMLA"
@@ -552,6 +555,16 @@ def _make_remote_meta(
555
)
556
557
558
+def _register_remote_agents(worker, metadata, tp_size):
559
+ """Mirror the async handshake callback that publishes prepared agents."""
560
+ worker._remote_agents[metadata.engine_id] = {
561
+ (0, rank): worker.add_remote_agent(
562
+ metadata, remote_tp_rank=rank, remote_tp_size=tp_size
563
+ )
564
+ for rank in range(tp_size)
565
+ }
566
+
567
+
568
def _owned_byte_ranges(worker, group_logical_ids):
569
"""Byte ranges owned by a request: for each HMA region tensor, every
570
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():
622
remote_num_logical=12,
623
remote_ssm_sizes=(24, 32),
624
)
612
- for rank in (0, 1):
613
- worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=2)
625
+ _register_remote_agents(worker, meta_r, 2)
626
627
# Request B: 17 matched tokens. Local: 2 logical blocks (24 tok
628
# capacity); remote: 16 prefilled tokens -> 2 remote logical blocks.
@@ -626,7 +638,7 @@ def test_hetero_ppl_multi_read_writes_stay_within_request_blocks():
638
"remote_block_ids": remote_ids,
639
"remote_engine_id": "remote-engine",
640
"remote_request_id": "prefill-req-b",
629
- "remote_host": "localhost",
641
+ "remote_host": "remote-host",
642
"remote_port": 1234,
643
"tp_size": 2,
644
},
@@ -709,8 +721,7 @@ def _run_hetero_case(
721
remote_num_logical=max(2 * n_remote + 4, 8),
722
remote_ssm_sizes=(48 // tp_size, 64 // tp_size),
723
)
712
- for rank in range(tp_size):
713
- worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=tp_size)
724
+ _register_remote_agents(worker, meta_r, tp_size)
725
726
# Sparse ids so neighbors exist between the request's blocks.
727
local_attn = [2 * i + 1 for i in range(n_local)]
@@ -726,7 +737,7 @@ def _run_hetero_case(
737
"remote_block_ids": remote_ids,
738
"remote_engine_id": "remote-engine",
739
"remote_request_id": "prefill-req-b",
729
- "remote_host": "localhost",
740
+ "remote_host": "remote-host",
741
"remote_port": 1234,
742
"tp_size": tp_size,
743
},
tests/v1/kv_connector/unit/test_nixl_push_connector.py
+54
-22
@@ -15,7 +15,7 @@ requiring a real NIXL agent or network:
15
* The worker matches D registrations against P finished blocks (both
16
scenario directions) and forwards non-PUSH_REG NIXL notifs to the main
17
thread's ``_get_new_notifs``.
18
-* ``get_finished`` enqueues evictions for the writer.
18
+* ``get_transfer_results`` enqueues evictions for the writer.
19
"""
20
21
from __future__ import annotations
@@ -32,6 +32,9 @@ from unittest.mock import MagicMock, patch
32
import msgspec
33
import pytest
34
35
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
36
+ KVConnectorTransferResults,
37
+)
38
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
39
PUSH_REG_NOTIF_PREFIX,
40
NixlAgentMetadata,
@@ -347,6 +350,8 @@ class _StubWriterWorker(NixlPushConnectorWorker):
350
351
# Base worker fields touched by start_load_kv / _get_new_notifs.
352
w._recving_metadata = {}
353
+ w._pending_recv_notifs = {}
354
+ w._failed_inflight_recvs = set()
355
w._recving_transfers = defaultdict(list)
356
w._reqs_to_process = set()
357
w._reqs_to_send = {}
@@ -742,21 +747,19 @@ class TestPushWriterNotifs:
747
assert notified == set()
748
assert request_id in w._recving_transfers
749
745
- def test_get_finished_evicts_completed_state(self):
746
- """``get_finished`` should enqueue evictions and wake the writer."""
750
+ def test_get_transfer_results_evicts_completed_state(self):
751
+ """Transfer completion should enqueue evictions and wake the writer."""
752
w = _StubWriterWorker.fresh()
753
749
- # Stub the base ``get_finished`` to return one done_sending entry.
750
- # Patch via the MRO's parent class.
754
with patch.object(
755
NixlPushConnectorWorker.__mro__[1],
753
- "get_finished",
754
- return_value=({"req-done"}, set()),
756
+ "get_transfer_results",
757
+ return_value=KVConnectorTransferResults(finished_sending={"req-done"}),
758
):
756
- done_sending, done_recving = w.get_finished()
759
+ results = w.get_transfer_results()
760
758
- assert "req-done" in done_sending
759
- assert done_recving == set()
761
+ assert results.finished_sending == {"req-done"}
762
+ assert results.finished_recving == set()
763
# Eviction enqueued for the writer.
764
evicted = []
765
while True:
@@ -928,19 +931,21 @@ class TestPushWriterNegative:
931
assert len(w._pending_d_registrations) == 1
932
assert w.start_push_calls == []
933
931
- def test_get_finished_enqueues_eviction_for_each_done_request(self):
932
- """``get_finished`` must enqueue an eviction for every request
934
+ def test_get_transfer_results_enqueues_each_completed_request(self):
935
+ """``get_transfer_results`` must enqueue every completed request
936
in ``done_sending`` so the writer can drop stale matching state.
937
Unlike the happy-path test, this verifies the *cardinality*: N
938
completed requests -> N evictions, in order."""
939
w = _StubWriterWorker.fresh()
940
with patch.object(
941
NixlPushConnectorWorker.__mro__[1],
939
- "get_finished",
940
- return_value=({"req-1", "req-2", "req-3"}, set()),
942
+ "get_transfer_results",
943
+ return_value=KVConnectorTransferResults(
944
+ finished_sending={"req-1", "req-2", "req-3"}
945
+ ),
946
):
942
- done_sending, _ = w.get_finished()
943
- assert done_sending == {"req-1", "req-2", "req-3"}
947
+ results = w.get_transfer_results()
948
+ assert results.finished_sending == {"req-1", "req-2", "req-3"}
949
950
evicted: list[str] = []
951
while True:
@@ -950,23 +955,47 @@ class TestPushWriterNegative:
955
break
956
assert sorted(evicted) == ["req-1", "req-2", "req-3"]
957
953
- def test_get_finished_with_no_completions_does_not_enqueue_eviction(self):
958
+ def test_empty_transfer_results_do_not_enqueue_eviction(self):
959
"""If there's nothing newly done, no eviction should be enqueued.
960
The wake event IS still set because ``get_finished`` always wakes
961
the writer to drain notifs."""
962
w = _StubWriterWorker.fresh()
963
with patch.object(
964
NixlPushConnectorWorker.__mro__[1],
960
- "get_finished",
961
- return_value=(set(), set()),
965
+ "get_transfer_results",
966
+ return_value=KVConnectorTransferResults(),
967
):
963
- done_sending, done_recving = w.get_finished()
964
- assert done_sending == set()
965
- assert done_recving == set()
968
+ results = w.get_transfer_results()
969
+ assert results.finished_sending == set()
970
+ assert results.finished_recving == set()
971
assert w._evict_finished_inbox.qsize() == 0
972
# Wake set so the writer drains NIXL notifs even when idle.
973
assert w._push_writer_wake.is_set()
974
975
+ def test_failed_push_transfer_finishes_sending(self):
976
+ w = _StubWriterWorker.fresh()
977
+ w.nixl_wrapper = MagicMock()
978
+ w.nixl_wrapper.check_xfer_state.return_value = "ERR"
979
+ w.xfer_stats = MagicMock()
980
+ w._log_failure = MagicMock()
981
+ w._sending_transfers["req-failed"] = [17]
982
+ w._reqs_to_send["req-failed"] = time.perf_counter() + 30
983
+ w._reqs_to_process.add("req-failed")
984
+
985
+ with patch.object(
986
+ NixlPushConnectorWorker.__mro__[1],
987
+ "get_transfer_results",
988
+ return_value=KVConnectorTransferResults(),
989
+ ):
990
+ results = w.get_transfer_results()
991
+
992
+ assert results.finished_sending == {"req-failed"}
993
+ assert results.finished_recving == set()
994
+ assert "req-failed" not in w._sending_transfers
995
+ assert "req-failed" not in w._reqs_to_send
996
+ assert "req-failed" not in w._reqs_to_process
997
+ w.nixl_wrapper.release_xfer_handle.assert_called_once_with(17)
998
+
999
def test_get_new_notifs_unknown_request_is_logged_and_skipped(self, caplog):
1000
"""A completion notif for a request the worker doesn't know
1001
about should be logged but not crash."""
@@ -1160,6 +1189,7 @@ class TestPushWriterMlaReplication:
1189
)
1190
}
1191
w._logical_to_kernel_block_ids = lambda block_ids, ratio: block_ids
1192
+ w.dst_region_group_ids[engine_id] = [0]
1193
w.dst_xfer_side_handles = {engine_id: {r: 1000 + r for r in d_ranks}}
1194
w.src_xfer_handles_by_block_size = {16: 2000}
1195
w._remote_agents = {engine_id: {(0, r): f"agent-{r}" for r in d_ranks}}
@@ -1240,6 +1270,8 @@ class TestPushWriterMlaReplication:
1270
)
1271
}
1272
w._logical_to_kernel_block_ids = lambda block_ids, ratio: block_ids
1273
+ w.region_group_ids = [0, 1, 2, 3]
1274
+ w.dst_region_group_ids[engine_id] = [0, 1, 2, 3]
1275
w.dst_xfer_side_handles = {engine_id: {0: 1000, 1: 1001}}
1276
w.src_xfer_handles_by_tp_ratio = {(-2, 16): [2000, 2001]}
1277
tests/v1/kv_connector/unit/test_output_aggregator.py
+11
-1
@@ -15,12 +15,14 @@ class DummyModelRunnerOutput(ModelRunnerOutput):
15
finished_sending: set[str] | None = None,
16
finished_recving: set[str] | None = None,
17
invalid_block_ids: set[int] | None = None,
18
+ failed_recving: set[str] | None = None,
19
expected_finished_count: int = 0,
20
):
21
self.kv_connector_output = KVConnectorOutput(
22
finished_sending=finished_sending,
23
finished_recving=finished_recving,
24
invalid_block_ids=invalid_block_ids or set(),
25
+ failed_recving=failed_recving or set(),
26
expected_finished_count=expected_finished_count,
27
)
28
@@ -73,7 +75,9 @@ def test_aggregate_workers_output():
75
76
output1 = DummyModelRunnerOutput(invalid_block_ids={3, 4})
77
output2 = DummyModelRunnerOutput(
76
- finished_recving={"req2"}, invalid_block_ids={4, 5}
78
+ finished_recving={"req2"},
79
+ invalid_block_ids={4, 5},
80
+ failed_recving={"req3"},
81
)
82
83
aggregated = aggregator.aggregate([output1, output2])
@@ -83,6 +87,12 @@ def test_aggregate_workers_output():
87
assert aggregated.finished_sending is None
88
assert aggregated.finished_recving == {"req2"}
89
assert aggregated.invalid_block_ids == {3, 4, 5}
90
+ assert not aggregated.failed_recving
91
+
92
+ output1 = DummyModelRunnerOutput(finished_recving={"req3"})
93
+ output2 = DummyModelRunnerOutput(finished_recving={"req3"})
94
+ aggregated = aggregator.aggregate([output1, output2])
95
+ assert aggregated.kv_connector_output.failed_recving == {"req3"}
96
97
98
def test_aggregate_workers_output_with_expected_finished_count():
tests/v1/kv_connector/unit/utils.py
+1
@@ -539,6 +539,7 @@ def make_nixl_scheduler(
539
sched._heartbeat_interval = kv_lease_duration // 6
540
# Fields touched by build_connector_meta / request_finished:
541
sched._reqs_need_recv = {}
542
+ sched._hisparse_host_blocks_to_recv = {}
543
sched._reqs_need_send = {}
544
sched._reqs_in_batch = set()
545
sched._reqs_not_processed = set()
tests/v1/kv_offload/test_factory.py
+3
-3
@@ -62,7 +62,7 @@ def _make_offloading_config(
62
normalized_extra_config["cpu_bytes_to_use"] = cpu_bytes_to_use
63
64
if groups is None:
65
- groups = (OffloadingGroupConfig(16, ("layer",)),)
65
+ groups = (OffloadingGroupConfig(16, ("layer",), 0),)
66
67
return OffloadingConfig(
68
groups=groups,
@@ -493,8 +493,8 @@ def test_offloading_spec_has_replicated_layout_default():
493
494
def test_offloading_spec_uses_normalized_chunk_geometry():
495
groups = (
496
- OffloadingGroupConfig(12, ("full_layer",)),
497
- OffloadingGroupConfig(16, ("mla_layer",)),
496
+ OffloadingGroupConfig(12, ("full_layer",), 0),
497
+ OffloadingGroupConfig(16, ("mla_layer",), 1),
498
)
499
spec = _create_spec(
500
groups=groups,
tests/v1/kv_offload/test_file_mapper.py
+4
-1
@@ -26,8 +26,11 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper:
26
OffloadingGroupConfig(
27
tokens_per_block=tokens_per_block,
28
layer_names=(layer_name,),
29
+ group_id=group_id,
30
+ )
31
+ for group_id, (tokens_per_block, layer_name) in enumerate(
32
+ kwargs.get("groups", ())
33
)
30
- for tokens_per_block, layer_name in kwargs.get("groups", ())
34
),
35
worker_kv_bytes_per_block=0,
36
enable_kv_cache_events=False,
tests/v1/simple_kv_offload/test_worker.py
+6
@@ -400,6 +400,8 @@ def test_register_mixed_page_sizes_in_one_cache_group(monkeypatch):
400
vllm_config.cache_config = CacheConfig()
401
vllm_config.cache_config.kv_cache_layout = layout.name
402
vllm_config.cache_config.num_gpu_blocks_override = None
403
+ vllm_config.attention_config.hisparse_config = None
404
+ vllm_config.kv_transfer_config = None
405
406
pages = [spec.page_size_bytes for spec in specs.values()]
407
num_blocks = 4
@@ -458,6 +460,8 @@ def test_register_mixed_page_sizes_odd_block_counts(monkeypatch, rank_blocks):
460
vllm_config.cache_config = CacheConfig()
461
vllm_config.cache_config.kv_cache_layout = layout.name
462
vllm_config.cache_config.num_gpu_blocks_override = None
463
+ vllm_config.attention_config.hisparse_config = None
464
+ vllm_config.kv_transfer_config = None
465
466
pages = [spec.page_size_bytes for spec in specs.values()]
467
kv_cache_config = get_kv_cache_config_from_groups(
@@ -507,6 +511,8 @@ def test_mixed_page_byte_placement_is_dcp_invariant():
511
vllm_config.cache_config.kv_cache_layout = KVCacheLayout.LBNHC.name
512
vllm_config.cache_config.num_gpu_blocks_override = None
513
vllm_config.parallel_config.decode_context_parallel_size = dcp_size
514
+ vllm_config.attention_config.hisparse_config = None
515
+ vllm_config.kv_transfer_config = None
516
517
page_bytes = sum(spec.page_size_bytes for spec in specs.values())
518
config = get_kv_cache_config_from_groups(vllm_config, [group], page_bytes * 4)
tests/v1/test_kv_cache_spec_registry.py
+15
@@ -16,6 +16,7 @@ from vllm.v1.core.single_type_kv_cache_manager import (
16
ChunkedLocalAttentionManager,
17
CrossAttentionManager,
18
FullAttentionManager,
19
+ HiSparseSourceManager,
20
MambaManager,
21
SingleTypeKVCacheManager,
22
SinkFullAttentionManager,
@@ -27,6 +28,7 @@ from vllm.v1.kv_cache_interface import (
28
CrossAttentionSpec,
29
FullAttentionSpec,
30
HiddenStateCacheSpec,
31
+ KVCacheGroupRole,
32
KVCacheSpec,
33
KVCacheSpecKind,
34
MambaSpec,
@@ -176,6 +178,19 @@ class TestKVCacheSpecRegistry:
178
is spec_uniform_base_map[spec_cls]
179
)
180
181
+ @pytest.mark.parametrize("role", list(KVCacheGroupRole))
182
+ def test_mla_manager_selection_by_role(self, role):
183
+ """Only the source role overrides ordinary MLA manager selection."""
184
+ expected = (
185
+ HiSparseSourceManager
186
+ if role == KVCacheGroupRole.HISPARSE_SOURCE
187
+ else FullAttentionManager
188
+ )
189
+ assert (
190
+ KVCacheSpecRegistry.get_manager_class(make_spec(MLAAttentionSpec), role)
191
+ is expected
192
+ )
193
+
194
@pytest.mark.parametrize("spec_cls", list(spec_manager_map))
195
def test_custom_spec_register(self, spec_cls):
196
"""A decorated custom spec resolves to the declared manager."""
tests/v1/worker/test_attn_utils.py
+126
-1
@@ -13,15 +13,18 @@ import pytest
13
import torch
14
15
from tests.v1.attention.utils import dense_kv_cache_views
16
-from vllm.v1.attention.backend import AttentionCGSupport
16
+from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf
17
from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
18
+from vllm.v1.hisparse.binding import allocate_hisparse_kv_caches
19
from vllm.v1.kv_cache_interface import (
20
FullAttentionSpec,
21
+ HiSparseResidentSpec,
22
KVCacheConfig,
23
KVCacheGroupSpec,
24
KVCacheLayout,
25
KVCacheTensor,
26
MLAAttentionSpec,
27
+ SparseCacheRole,
28
compute_layout_strides,
29
)
30
from vllm.v1.worker.gpu import attn_utils
@@ -36,6 +39,60 @@ from vllm.v1.worker.utils import (
39
)
40
41
42
+@pytest.mark.parametrize(
43
+ ("enabled", "block_size", "main_sizes", "indexer_sizes", "expected"),
44
+ [
45
+ (True, 256, [64], [64], 64),
46
+ (True, 64, [32, 64], [16, 32], 32),
47
+ (True, 64, [MultipleOf(16)], [32], 32),
48
+ (True, 64, [64], [32], None),
49
+ (False, 256, [64], [64], 256),
50
+ ],
51
+)
52
+def test_get_kv_cache_spec_resolves_hisparse_block_size(
53
+ monkeypatch, enabled, block_size, main_sizes, indexer_sizes, expected
54
+):
55
+ """Resolve shared MLA geometry before planning; leave other specs alone."""
56
+ specs = {
57
+ "main": MLAAttentionSpec(
58
+ block_size=block_size, num_kv_heads=1, head_size=576, dtype=torch.bfloat16
59
+ ),
60
+ "indexer": MLAAttentionSpec(
61
+ block_size=block_size,
62
+ num_kv_heads=1,
63
+ head_size=128,
64
+ dtype=torch.bfloat16,
65
+ cache_role=SparseCacheRole.INDEXER,
66
+ ),
67
+ "dense": FullAttentionSpec(
68
+ block_size=block_size, num_kv_heads=1, head_size=128, dtype=torch.bfloat16
69
+ ),
70
+ }
71
+ layers = {}
72
+ for name, sizes in zip(specs, [main_sizes, indexer_sizes, [block_size]]):
73
+ backend = SimpleNamespace(
74
+ customize_spec=AttentionBackend.customize_spec,
75
+ get_supported_kernel_block_sizes=lambda sizes=sizes: sizes,
76
+ )
77
+ layers[name] = SimpleNamespace(
78
+ get_kv_cache_spec=lambda _, spec=specs[name]: spec,
79
+ get_attn_backend=lambda backend=backend: backend,
80
+ )
81
+ monkeypatch.setattr(attn_utils, "get_layers_from_vllm_config", lambda *_: layers)
82
+ config = SimpleNamespace(
83
+ attention_config=SimpleNamespace(hisparse_config=object() if enabled else None)
84
+ )
85
+ if expected is None:
86
+ with pytest.raises(ValueError, match="supported by every sparse"):
87
+ attn_utils.get_kv_cache_spec(config)
88
+ return
89
+
90
+ resolved = attn_utils.get_kv_cache_spec(config)
91
+ assert resolved["main"].block_size == resolved["indexer"].block_size == expected
92
+ assert resolved["dense"] is specs["dense"]
93
+ assert all(spec.block_size == block_size for spec in specs.values())
94
+
95
+
96
class _FakeMetadataBuilder:
97
def __init__(self, support: AttentionCGSupport):
98
self.support = support
@@ -390,3 +447,71 @@ def test_copy_kv_cache_blocks_with_virtual_block_splitting(
447
torch.testing.assert_close(
448
cache[dst_start + physical_idx], expected[layer_idx][physical_idx]
449
)
450
+
451
+
452
+def test_allocate_hisparse_kv_caches_host_pool_and_view_less_specs():
453
+ """Host tensors get their own backing; view-less specs keep the raw one."""
454
+ spec = FullAttentionSpec(
455
+ block_size=2, num_kv_heads=1, head_size=4, dtype=torch.float32
456
+ )
457
+ page = spec.page_size_bytes
458
+ resident_spec = HiSparseResidentSpec(block_size=2, page_size=page)
459
+ device_size = 4 * page
460
+ config = KVCacheConfig(
461
+ num_blocks=4,
462
+ hisparse_host_num_blocks=3,
463
+ kv_cache_tensors=[
464
+ KVCacheTensor(
465
+ size=3 * page,
466
+ layers=["source"],
467
+ layer_stride=3 * page,
468
+ block_stride=page,
469
+ host_resident=True,
470
+ ),
471
+ KVCacheTensor(
472
+ size=device_size,
473
+ layers=["indexer"],
474
+ layer_stride=device_size,
475
+ block_stride=page,
476
+ ),
477
+ KVCacheTensor(
478
+ size=device_size,
479
+ layers=["resident"],
480
+ layer_stride=device_size,
481
+ block_stride=page,
482
+ ),
483
+ ],
484
+ kv_cache_groups=[
485
+ KVCacheGroupSpec(["source"], spec, host_resident=True),
486
+ KVCacheGroupSpec(["indexer"], spec),
487
+ KVCacheGroupSpec(["resident"], resident_spec),
488
+ ],
489
+ )
490
+ host_buffers: list[torch.Tensor] = []
491
+
492
+ def host_allocator(size: int) -> torch.Tensor:
493
+ host_buffers.append(torch.zeros(size, dtype=torch.int8))
494
+ return host_buffers[-1]
495
+
496
+ caches = allocate_hisparse_kv_caches(
497
+ config,
498
+ torch.device("cpu"),
499
+ KVCacheLayout.LBHNC,
500
+ [2, 2, 2],
501
+ SimpleNamespace(allocate=host_allocator),
502
+ )
503
+ assert len(config.kv_cache_tensors) == 3
504
+
505
+ assert [buf.numel() for buf in host_buffers] == [3 * page]
506
+ assert caches["source"].shape[0] == 3
507
+ assert (
508
+ caches["source"].untyped_storage().data_ptr()
509
+ == host_buffers[0].untyped_storage().data_ptr()
510
+ )
511
+ assert caches["indexer"].shape[0] == 4
512
+ backing = caches["resident"]
513
+ assert backing.dtype == torch.int8 and backing.numel() >= device_size
514
+ assert (
515
+ backing.untyped_storage().data_ptr()
516
+ == caches["indexer"].untyped_storage().data_ptr()
517
+ )
tests/v1/worker/test_dsv4_packed_zeroer_geometry.py
+1
@@ -134,6 +134,7 @@ def test_overlaid_zeroer_dedups_segments_with_max_span():
134
vllm_config.cache_config = CacheConfig()
135
vllm_config.cache_config.num_gpu_blocks_override = None
136
vllm_config.cache_config.kv_cache_layout = "BLHNC"
137
+ vllm_config.attention_config.hisparse_config = None
138
config = get_kv_cache_config_from_groups(vllm_config, groups, 8 * 1024 * 1024)
139
views = allocate_kv_cache(config, torch.device("cpu"), KVCacheLayout.BLHNC, None)
140
buf_ptr = views["g1.big"].data_ptr()
tests/v1/worker/test_gpu_autoregressive_speculator.py
+2
@@ -351,6 +351,8 @@ def test_multi_step_decode_replays_captured_graph_as_expected(
351
speculator = object.__new__(_TestSpeculator)
352
speculator.num_speculative_steps = 4
353
speculator.current_draft_step = torch.tensor(0)
354
+ speculator.slot_mapping_observer = None
355
+ speculator.host_mirror_forward_observer = None
356
speculator.input_buffers = SimpleNamespace(
357
positions=torch.arange(2),
358
query_start_loc=torch.arange(3),
tests/v1/worker/test_gpu_kv_connector.py
+23
-3
@@ -5,9 +5,11 @@ from types import SimpleNamespace
5
from unittest.mock import Mock
6
7
import pytest
8
+import torch
9
10
import vllm.v1.worker.gpu.kv_connector as kv_connector_module
11
from vllm.config import KVTransferConfig
12
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorTransferResults
13
from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector
14
15
@@ -18,9 +20,9 @@ def _make_connector(
20
backend = Mock()
21
backend.handle_preemptions.side_effect = lambda _: events.append("handle")
22
backend.bind_connector_metadata.side_effect = lambda _: events.append("bind")
21
- backend.start_load_kv.side_effect = lambda _: events.append("start")
23
+ backend.start_load_kv.side_effect = lambda *_args, **_kwargs: events.append("start")
24
backend.wait_for_save.side_effect = lambda: events.append("wait")
23
- backend.get_finished.side_effect = lambda _: (set(), set())
25
+ backend.get_transfer_results.return_value = KVConnectorTransferResults()
26
backend.get_block_ids_with_load_errors.return_value = set()
27
backend.get_kv_connector_stats.return_value = None
28
backend.get_kv_connector_kv_cache_events.return_value = None
@@ -61,7 +63,15 @@ def test_load_start_phase(
63
connector = _make_connector(monkeypatch, events)
64
output = _scheduler_output(has_sync_kv_loads)
65
64
- connector.pre_forward(output) # type: ignore[arg-type]
66
+ request_indices = torch.tensor([3, 1])
67
+ request_ids = ["first", "second"]
68
+ attn_metadata = {"layer": object()}
69
+ connector.pre_forward( # type: ignore[arg-type]
70
+ output,
71
+ request_state_indices=request_indices,
72
+ request_ids=request_ids,
73
+ attn_metadata=attn_metadata,
74
+ )
75
assert events == (
76
["handle", "bind", "start"] if has_sync_kv_loads else ["handle", "bind"]
77
)
@@ -69,6 +79,16 @@ def test_load_start_phase(
79
connector.post_forward(set())
80
assert events == ["handle", "bind", "start", "wait", "clear"]
81
82
+ kwargs = connector.kv_connector.start_load_kv.call_args.kwargs
83
+ assert kwargs["request_state_indices"] is request_indices
84
+ assert kwargs["request_ids"] is request_ids
85
+ assert kwargs["attn_metadata"] is attn_metadata
86
+
87
+ # A subsequent step without a forward must not reuse the prior batch.
88
+ connector.no_forward(_scheduler_output(False)) # type: ignore[arg-type]
89
+ assert connector.kv_connector.start_load_kv.call_count == 2
90
+ assert connector.kv_connector.start_load_kv.call_args.kwargs == {}
91
+
92
93
def test_no_forward_starts_deferred_load_once(monkeypatch: pytest.MonkeyPatch):
94
events: list[str] = []
tests/v1/worker/test_gpu_model_runner.py
+4
@@ -312,6 +312,10 @@ def test_select_common_block_size_uses_largest_shared_int():
312
assert selected_size == 64
313
314
315
+def test_select_common_block_size_without_active_backends_uses_manager_size():
316
+ assert select_common_block_size(256, []) == 256
317
+
318
+
319
def test_select_common_block_size_accepts_rocm_sparse_block_size_16(monkeypatch):
320
monkeypatch.setattr(current_platform, "is_rocm", lambda: True)
321
tests/v1/worker/test_gpu_model_runner_v2.py
+1
@@ -237,6 +237,7 @@ def _make_capture_runner(captured: bool) -> GPUModelRunner:
237
runner.attn_groups = None
238
runner.kv_cache_config = None
239
runner.use_aux_hidden_state_outputs = False
240
+ runner.kv_connector = model_runner_module.NO_OP_KV_CONNECTOR
241
return runner
242
243
tests/v1/worker/test_gpu_warmup_blocks.py
+1
@@ -89,6 +89,7 @@ def _make_runner(
89
return SimpleNamespace(
90
num_speculative_steps=num_spec_steps,
91
adaptive_verification=None,
92
+ rejection_sampler=None,
93
decode_query_len=num_spec_steps + 1,
94
is_pooling_model=False,
95
is_encoder_decoder=False,
tests/v1/worker/test_kv_cache_allocation_scope.py
+31
-3
@@ -5,8 +5,10 @@ from contextlib import AbstractContextManager
5
from types import SimpleNamespace
6
from typing import Any, cast
7
8
+import pytest
9
import torch
10
11
+import vllm.v1.hisparse.binding as hisparse_binding
12
import vllm.v1.worker.gpu.attn_utils as attn_utils
13
import vllm.v1.worker.gpu_model_runner as gpu_model_runner
14
from vllm.v1.worker.gpu_worker import Worker
@@ -26,7 +28,8 @@ class _AllocationScope(AbstractContextManager):
28
self.active = False
29
30
29
-def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None:
31
+@pytest.mark.parametrize("hisparse", [False, True])
32
+def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch, hisparse) -> None:
33
scope = _AllocationScope()
34
kv_caches = {"layer": torch.empty(0)}
35
@@ -38,23 +41,48 @@ def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None:
41
assert not scope.active
42
43
monkeypatch.setattr(attn_utils, "allocate_kv_cache", allocate)
41
- monkeypatch.setattr(attn_utils, "bind_kv_cache", bind)
44
+ monkeypatch.setattr(attn_utils, "bind_kv_cache_to_layers", bind)
45
monkeypatch.setattr(attn_utils, "get_shared_kv_cache_layers", lambda config: {})
46
47
+ hisparse_bindings = []
48
+ if hisparse:
49
+ host_pool = SimpleNamespace()
50
+
51
+ def bind_hisparse(**kwargs):
52
+ hisparse_bindings.append(kwargs)
53
+ assert scope.active
54
+ assert kwargs["kv_caches"] is kv_caches
55
+ assert kwargs["host_pool"] is host_pool
56
+ return [
57
+ SimpleNamespace(
58
+ view=SimpleNamespace(cache=torch.empty(1, 1, 1), block_size=1),
59
+ runtime=SimpleNamespace(),
60
+ )
61
+ ]
62
+
63
+ monkeypatch.setattr(hisparse_binding, "HiSparseHostPool", lambda: host_pool)
64
+ monkeypatch.setattr(hisparse_binding, "allocate_hisparse_kv_caches", allocate)
65
+ monkeypatch.setattr(hisparse_binding, "bind_hisparse_kv_caches", bind_hisparse)
66
+
67
config = SimpleNamespace(
68
+ attention_config=SimpleNamespace(
69
+ hisparse_config=object() if hisparse else None
70
+ ),
71
+ scheduler_config=SimpleNamespace(max_num_seqs=1, max_num_batched_tokens=1),
72
cache_config=SimpleNamespace(get_resolved_kv_cache_layout=lambda: None),
73
model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type="test")),
74
)
75
result = attn_utils.init_kv_cache(
49
- [],
76
{},
77
SimpleNamespace(kv_cache_groups=[]),
78
torch.device("cpu"),
79
[],
80
config,
81
kv_cache_allocation_context=scope,
82
+ block_tables=SimpleNamespace(),
83
)
84
85
+ assert len(hisparse_bindings) == int(hisparse)
86
assert result is kv_caches
87
assert not scope.active
88
tests/v1/worker/test_mixed_warmup_gate.py
+7
-1
@@ -35,10 +35,15 @@ def test_mixed_warmup_skipped_for_single_seq(max_num_reqs):
35
def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_warmup):
36
"""Startup must warm fixed drafts before calibration and retain its manager."""
37
manager = SimpleNamespace(cost_tables=None)
38
- runner = SimpleNamespace(adaptive_verification=manager)
38
+ rejection_sampler = SimpleNamespace(enable_adaptive_verification=True)
39
+ runner = SimpleNamespace(
40
+ adaptive_verification=manager,
41
+ rejection_sampler=rejection_sampler,
42
+ )
43
44
def run_steps(model_runner, execute, sample):
45
assert model_runner.adaptive_verification is None
46
+ assert not model_runner.rejection_sampler.enable_adaptive_verification
47
if fail_warmup:
48
raise RuntimeError("warmup failed")
49
@@ -50,3 +55,4 @@ def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_
55
warmup.warmup_kernels(runner, _fail, _fail)
56
assert runner.adaptive_verification is manager
57
assert manager.cost_tables is None
58
+ assert rejection_sampler.enable_adaptive_verification
tests/v1/worker/test_utils.py
+891
-3
@@ -1,13 +1,897 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from collections import deque
5
+from contextlib import nullcontext
6
from types import SimpleNamespace
7
+from unittest.mock import MagicMock, call
8
9
+import numpy as np
10
+import pytest
11
import torch
12
13
+import vllm.v1.attention.backends.mla.index_group as index_group_module
14
+import vllm.v1.hisparse.runtime as hisparse_runtime_module
15
+from vllm.config import CUDAGraphMode
16
from vllm.config.mamba import MambaBackendEnum, MambaConfig
17
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import (
18
+ worker as hisparse_worker_module,
19
+)
20
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
21
+ HiSparseConnectorWorker,
22
+ _flatten_row_mirrors,
23
+ _select_written_row_mirrors,
24
+ _SlotMappingStaging,
25
+)
26
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
10
-from vllm.v1.worker.utils import bind_kv_cache
27
+from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
28
+from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror
29
+from vllm.v1.worker.utils import (
30
+ bind_kv_cache,
31
+ bind_kv_cache_to_layers,
32
+ copy_kv_cache_blocks_inplace,
33
+)
34
+
35
+
36
+def _make_hisparse_worker() -> HiSparseConnectorWorker:
37
+ worker = object.__new__(HiSparseConnectorWorker)
38
+ worker.refines_row_mirrors = False
39
+ worker._slot_mapping_staging = None
40
+ worker._row_mirror_num_rows = 0
41
+ worker._per_layer_mirrored = set()
42
+ worker._submitted_mirror_layers = set()
43
+ worker._pending_dma_descriptors = deque()
44
+ worker._dma_free_descriptors = []
45
+ worker.host_write_events = (MagicMock(), MagicMock())
46
+ worker.host_write_event = worker.host_write_events[1]
47
+ worker._next_host_write_event = 0
48
+ worker.dma_stream = None
49
+ worker.shared_host_region = None
50
+ worker._completed_host_copy_dst_ids = []
51
+ worker._slot_mapping_staging = None
52
+ return worker
53
+
54
+
55
+def test_hisparse_row_mirrors_follow_runner_request_order():
56
+ first = SparseKVRowMirror((1,), 10, 1)
57
+ second = SparseKVRowMirror((2,), 20, 1)
58
+ row_mirrors = {"first": (first,), "second": (second,)}
59
+
60
+ assert _flatten_row_mirrors(row_mirrors, ["second", "first"]) == (
61
+ second,
62
+ first,
63
+ )
64
+
65
+
66
+def test_hisparse_gpu_slots_select_written_rows_from_async_envelope():
67
+ """Rejected speculative rows must never become durable host KV."""
68
+ candidates = (
69
+ SparseKVRowMirror((100, 500), 1000, 4),
70
+ SparseKVRowMirror((20, 300), 2000, 4),
71
+ )
72
+ source_slots = np.array([102, 103, -1, 21, 22, 23], dtype=np.int64)
73
+
74
+ assert _select_written_row_mirrors(candidates, source_slots, 0) == (
75
+ SparseKVRowMirror((102, 502), 1002, 2),
76
+ SparseKVRowMirror((21, 301), 2001, 3),
77
+ )
78
+
79
+
80
+def test_hisparse_written_rows_preserve_scheduler_page_boundaries():
81
+ candidates = (
82
+ SparseKVRowMirror((0,), 100, 4),
83
+ SparseKVRowMirror((4,), 104, 4),
84
+ )
85
+
86
+ assert _select_written_row_mirrors(
87
+ candidates, np.array([2, 3, 4, 5], dtype=np.int64), 0
88
+ ) == (
89
+ SparseKVRowMirror((2,), 102, 2),
90
+ SparseKVRowMirror((4,), 104, 2),
91
+ )
92
+
93
+
94
+def test_hisparse_written_rows_without_host_destinations_are_skipped():
95
+ assert (
96
+ _select_written_row_mirrors((), np.array([100, 101], dtype=np.int64), 0) == ()
97
+ )
98
+
99
+
100
+def test_hisparse_appends_reference_slots_within_a_mirror_phase(monkeypatch):
101
+ """Separate context and query writes must share one mirror phase."""
102
+ worker = _make_hisparse_worker()
103
+ worker.is_host_writer = True
104
+ worker.refines_row_mirrors = True
105
+ state = _SlotMappingStaging(
106
+ MagicMock(), MagicMock(), torch.empty(4, dtype=torch.int64)
107
+ )
108
+ worker._slot_mapping_staging = state
109
+ handle = SimpleNamespace(runtime=SimpleNamespace(resident_source_index=1))
110
+ worker.cache_layer_names = ["layer"]
111
+ worker.cache_handles = [handle]
112
+ worker._layer_mirror_callbacks = (MagicMock(),)
113
+ worker._row_mirror_num_rows = 1
114
+ worker._submitted_mirror_layers = set()
115
+ main_stream = MagicMock()
116
+ monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: main_stream)
117
+ monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext())
118
+
119
+ handle.slot_mapping = torch.tensor([7, 8], dtype=torch.int64)
120
+ worker._stage_row_mirror_mapping(2)
121
+ handle.slot_mapping = torch.tensor([9, 10], dtype=torch.int64)
122
+ worker._stage_row_mirror_mapping(2)
123
+
124
+ torch.testing.assert_close(state.slots, torch.tensor([7, 8, 9, 10]))
125
+ assert state.stream.wait_stream.call_args_list == [call(main_stream)] * 2
126
+ assert state.event.record.call_args_list == [call(state.stream)] * 2
127
+ assert state.num_tokens == 4
128
+ assert state.source_index == 1
129
+
130
+
131
+@pytest.mark.parametrize(
132
+ ("mode", "expected_calls"),
133
+ [(CUDAGraphMode.PIECEWISE, 1), (CUDAGraphMode.FULL, 0)],
134
+)
135
+def test_hisparse_submits_layer_mirror_at_replayed_attention_boundary(
136
+ monkeypatch, mode, expected_calls
137
+):
138
+ """Piecewise replay must pipeline DMA; FULL replay uses the batch fallback."""
139
+ handle = object.__new__(hisparse_runtime_module.HiSparseCacheHandle)
140
+ handle.dummy_batch = False
141
+ handle.decode_batch = False
142
+ handle.submit_layer_mirror = MagicMock()
143
+ monkeypatch.setattr(
144
+ hisparse_runtime_module,
145
+ "get_forward_context",
146
+ lambda: SimpleNamespace(cudagraph_runtime_mode=mode),
147
+ )
148
+
149
+ handle.finish_kv_update()
150
+
151
+ assert handle.submit_layer_mirror.call_count == expected_calls
152
+
153
+
154
+def test_copy_cpu_kv_cache_logical_blocks_ignores_storage_padding():
155
+ backing = torch.full((10, 2, 3), -1, dtype=torch.float32)
156
+ cache = backing[1:9]
157
+ cache[2:4] = 7
158
+ cache[6:8] = 11
159
+
160
+ copy_kv_cache_blocks_inplace(
161
+ [cache],
162
+ num_blocks=4,
163
+ kv_cache_block_copies=[
164
+ KVCacheBlockCopy(1, 0),
165
+ KVCacheBlockCopy(3, 2),
166
+ ],
167
+ )
168
+
169
+ torch.testing.assert_close(cache[0:2], torch.full_like(cache[0:2], 7))
170
+ torch.testing.assert_close(cache[4:6], torch.full_like(cache[4:6], 11))
171
+ assert (backing[0] == -1).all()
172
+ assert (backing[9] == -1).all()
173
+
174
+
175
+@pytest.mark.parametrize("shared,rank", [(False, 0), (True, 0), (True, 1)])
176
+@pytest.mark.parametrize("has_copies", [False, True])
177
+def test_hisparse_host_copy_waits_for_writes_on_writer_only(
178
+ monkeypatch, shared, rank, has_copies
179
+):
180
+ """Host CoW must read completed writes; shared non-writers only barrier."""
181
+ worker = _make_hisparse_worker()
182
+ cache = torch.zeros((2, 1, 1))
183
+ worker.host_caches = [cache]
184
+ worker.host_num_blocks = 2
185
+ worker.shared_host_region = object() if shared else None
186
+ event = MagicMock()
187
+ event.synchronize.side_effect = lambda: cache[0].fill_(7)
188
+ barrier = MagicMock()
189
+ monkeypatch.setattr(
190
+ hisparse_worker_module, "get_tensor_model_parallel_rank", lambda: rank
191
+ )
192
+ monkeypatch.setattr(
193
+ hisparse_worker_module, "get_tp_group", lambda: SimpleNamespace(barrier=barrier)
194
+ )
195
+ copies = [KVCacheBlockCopy(0, 1)] if has_copies else []
196
+
197
+ worker._copy_host_blocks(copies, event)
198
+
199
+ writes = has_copies and (not shared or rank == 0)
200
+ assert cache[1].item() == (7 if writes else 0)
201
+ assert event.synchronize.call_count == int(writes)
202
+ assert barrier.call_count == int(shared and has_copies)
203
+ assert worker._completed_host_copy_dst_ids == ([1] if has_copies else [])
204
+
205
+
206
+def test_hisparse_worker_updates_request_state_mapping_in_place(monkeypatch):
207
+ worker = _make_hisparse_worker()
208
+ worker.request_state_indices = torch.arange(4, dtype=torch.int32)
209
+ worker._pending_invalid_block_ids = [5]
210
+ invalidations = []
211
+ worker.invalidate_blocks = lambda blocks, states: invalidations.append(
212
+ (blocks.copy(), states.clone())
213
+ )
214
+ original_ptr = worker.request_state_indices.data_ptr()
215
+ monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
216
+
217
+ worker.set_request_state_indices(torch.tensor([3, 1], dtype=torch.int32))
218
+
219
+ assert worker.request_state_indices.data_ptr() == original_ptr
220
+ assert worker.request_state_indices.tolist() == [3, 1, -1, -1]
221
+ assert len(invalidations) == 1
222
+ assert invalidations[0][0] == [5]
223
+ torch.testing.assert_close(
224
+ invalidations[0][1], torch.tensor([3, 1], dtype=torch.int32)
225
+ )
226
+ assert worker._pending_invalid_block_ids == []
227
+
228
+
229
+def test_hisparse_pre_forward_transfer_builds_page_descriptors():
230
+ worker = _make_hisparse_worker()
231
+ worker.is_host_writer = True
232
+ worker.kernel_block_size = 2
233
+ source = torch.empty((3, 2, 4), dtype=torch.uint8)
234
+ destination = torch.empty((6, 4), dtype=torch.uint8)
235
+ worker.resident_caches = (source,)
236
+ worker.host_caches = (destination,)
237
+ worker.cache_handles = [
238
+ SimpleNamespace(runtime=SimpleNamespace(resident_source_index=0))
239
+ ]
240
+ worker._dma_free_descriptors = []
241
+ worker._submit_dma_descriptors = MagicMock()
242
+
243
+ worker._enqueue_transfers([SparseKVPageTransfer(7, 2, (1,), after_forward=False)])
244
+
245
+ descriptors, count = worker._submit_dma_descriptors.call_args.args
246
+ assert count == 1
247
+ assert descriptors.src[:count].tolist() == [
248
+ source.data_ptr() + source.stride(0) * source.element_size()
249
+ ]
250
+ assert descriptors.dst[:count].tolist() == [destination.data_ptr() + 4 * 4]
251
+ assert descriptors.sizes[:count].tolist() == [2 * 4]
252
+ assert worker._submit_dma_descriptors.call_args.kwargs["transfer_ids"] == (7,)
253
+
254
+
255
+def test_hisparse_eager_mirror_records_transfer_without_page_copy():
256
+ worker = _make_hisparse_worker()
257
+ worker.cache_handles = [
258
+ SimpleNamespace(runtime=SimpleNamespace(eager_host_mirror=True))
259
+ ]
260
+ worker._record_transfer_completion = MagicMock()
261
+ worker._enqueue_transfers = MagicMock()
262
+ transfers = [SparseKVPageTransfer(7, 2, (1,), after_forward=False)]
263
+
264
+ worker._submit_transfers(transfers)
265
+
266
+ worker._record_transfer_completion.assert_called_once_with(transfers)
267
+ worker._enqueue_transfers.assert_not_called()
268
+
269
+
270
+def test_hisparse_lazy_mirror_copies_transfer_pages():
271
+ worker = _make_hisparse_worker()
272
+ worker.cache_handles = [
273
+ SimpleNamespace(runtime=SimpleNamespace(eager_host_mirror=False))
274
+ ]
275
+ worker._record_transfer_completion = MagicMock()
276
+ worker._enqueue_transfers = MagicMock()
277
+ transfers = [SparseKVPageTransfer(7, 2, (1,), after_forward=False)]
278
+
279
+ worker._submit_transfers(transfers)
280
+
281
+ worker._enqueue_transfers.assert_called_once_with(transfers)
282
+ worker._record_transfer_completion.assert_not_called()
283
+
284
+
285
+def test_hisparse_dma_row_mirror_builds_descriptors(monkeypatch):
286
+ worker = _make_hisparse_worker()
287
+ worker.is_host_writer = True
288
+ worker.kernel_block_size = 2
289
+ source = torch.empty((3, 2, 4), dtype=torch.uint8)
290
+ destination = torch.empty((6, 4), dtype=torch.uint8)
291
+ worker.resident_caches = (source,)
292
+ worker.host_caches = (destination,)
293
+ worker.cache_handles = [
294
+ SimpleNamespace(
295
+ runtime=SimpleNamespace(resident_source_index=0), decode_batch=True
296
+ )
297
+ ]
298
+ worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0"))
299
+ worker.dma_stream = MagicMock()
300
+ worker.host_write_event = MagicMock()
301
+ worker._dma_free_descriptors = []
302
+ worker._pending_dma_descriptors = deque()
303
+ worker._pending_transfer_events = deque()
304
+ worker._enqueued_transfer_ids = []
305
+ worker._dma_submitted = False
306
+ current_stream = MagicMock()
307
+ event = MagicMock()
308
+ swap_blocks_batch = MagicMock()
309
+ monkeypatch.setattr(
310
+ hisparse_worker_module, "current_stream", lambda: current_stream
311
+ )
312
+ monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext())
313
+ monkeypatch.setattr(torch, "Event", lambda: event)
314
+ monkeypatch.setattr(
315
+ hisparse_worker_module.ops, "swap_blocks_batch", swap_blocks_batch
316
+ )
317
+
318
+ worker._set_row_mirrors((SparseKVRowMirror((3,), 4, 1),))
319
+ worker._enqueue_row_dma(range(1))
320
+
321
+ src_ptrs, dst_ptrs, sizes = swap_blocks_batch.call_args.args
322
+ assert src_ptrs.tolist() == [source.data_ptr() + 3 * 4]
323
+ assert dst_ptrs.tolist() == [destination.data_ptr() + 4 * 4]
324
+ assert sizes.tolist() == [4]
325
+ worker.dma_stream.wait_stream.assert_called_once_with(current_stream)
326
+ worker.host_write_event.record.assert_called_once_with(worker.dma_stream)
327
+ event.record.assert_called_once_with(worker.dma_stream)
328
+ assert worker._dma_submitted
329
+
330
+
331
+def test_hisparse_row_dma_uses_resident_spans():
332
+ worker = _make_hisparse_worker()
333
+ source = torch.empty((2, 2, 4), dtype=torch.uint8)
334
+ destination = torch.empty((12, 4), dtype=torch.uint8)
335
+ worker.is_host_writer = True
336
+ worker.kernel_block_size = 2
337
+ worker.resident_caches = (source,)
338
+ worker.host_caches = (destination,)
339
+ worker.cache_handles = [
340
+ SimpleNamespace(
341
+ runtime=SimpleNamespace(resident_source_index=0), decode_batch=True
342
+ )
343
+ ]
344
+ worker._set_row_mirrors(
345
+ (
346
+ SparseKVRowMirror((0,), 4, 2),
347
+ SparseKVRowMirror((2,), 9, 1),
348
+ )
349
+ )
350
+ worker._dma_free_descriptors = []
351
+ worker._submit_dma_descriptors = MagicMock()
352
+
353
+ worker._enqueue_row_dma(range(1))
354
+
355
+ descriptors, count = worker._submit_dma_descriptors.call_args.args
356
+ assert count == 2
357
+ assert descriptors.src[:count].tolist() == [
358
+ source.data_ptr(),
359
+ source.data_ptr() + 2 * 4,
360
+ ]
361
+ assert descriptors.dst[:count].tolist() == [
362
+ destination.data_ptr() + 4 * 4,
363
+ destination.data_ptr() + 9 * 4,
364
+ ]
365
+ assert descriptors.sizes[:count].tolist() == [2 * 4, 4]
366
+
367
+
368
+def test_hisparse_finish_forward_mirrors_all_layers_once(monkeypatch):
369
+ dst_slots = torch.tensor([7, 8, 9], dtype=torch.int64)
370
+ req_ids = torch.tensor([0, 1, 2], dtype=torch.int32)
371
+ leader = SimpleNamespace(
372
+ eager_host_mirror=True,
373
+ is_group_leader=True,
374
+ invalidate_written_slots=MagicMock(),
375
+ )
376
+ follower = SimpleNamespace(
377
+ eager_host_mirror=True,
378
+ is_group_leader=False,
379
+ invalidate_written_slots=MagicMock(),
380
+ )
381
+ handles = [
382
+ SimpleNamespace(
383
+ runtime=runtime,
384
+ decode_batch=True,
385
+ host_mirror_required=True,
386
+ num_actual_tokens=3,
387
+ num_decode_tokens=2,
388
+ req_id_per_token=req_ids,
389
+ mirror_slot_mapping=dst_slots,
390
+ )
391
+ for runtime in (leader, follower)
392
+ ]
393
+ worker = _make_hisparse_worker()
394
+ worker.is_host_writer = True
395
+ worker.cache_handles = handles
396
+ worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0"))
397
+ worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 3),))
398
+ worker._dma_submitted = False
399
+ worker._per_layer_mirrored = set()
400
+ worker._submitted_mirror_layers = set()
401
+ worker._post_forward_transfers = []
402
+ worker._enqueue_row_dma = MagicMock()
403
+ worker.host_write_event = MagicMock()
404
+ worker._forward_ready_event = MagicMock()
405
+ current_stream = MagicMock()
406
+ monkeypatch.setattr(
407
+ hisparse_worker_module, "current_stream", lambda: current_stream
408
+ )
409
+
410
+ worker.finish_forward()
411
+
412
+ worker._forward_ready_event.record.assert_called_once_with()
413
+ worker._enqueue_row_dma.assert_called_once_with(
414
+ (0, 1), ready_event=worker._forward_ready_event
415
+ )
416
+ leader.invalidate_written_slots.assert_called_once()
417
+ torch.testing.assert_close(
418
+ leader.invalidate_written_slots.call_args.args[0], dst_slots
419
+ )
420
+ follower.invalidate_written_slots.assert_not_called()
421
+ worker.host_write_event.record.assert_called_once_with(current_stream)
422
+
423
+
424
+def test_hisparse_finish_forward_does_not_repeat_per_layer_mirrors():
425
+ slots = torch.tensor([7, 8], dtype=torch.int64)
426
+ runtime = SimpleNamespace(
427
+ eager_host_mirror=False,
428
+ is_group_leader=False,
429
+ invalidate_written_slots=MagicMock(),
430
+ )
431
+ handles = [
432
+ SimpleNamespace(
433
+ runtime=runtime,
434
+ decode_batch=False,
435
+ host_mirror_required=True,
436
+ num_actual_tokens=2,
437
+ num_decode_tokens=0,
438
+ req_id_per_token=torch.empty(0, dtype=torch.int32),
439
+ mirror_slot_mapping=slots,
440
+ )
441
+ for _ in range(2)
442
+ ]
443
+ worker = _make_hisparse_worker()
444
+ worker.is_host_writer = True
445
+ worker.cache_handles = handles
446
+ worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 2),))
447
+ worker._per_layer_mirrored = {0, 1}
448
+ worker._submitted_mirror_layers = {0, 1}
449
+ worker._enqueue_row_dma = MagicMock()
450
+
451
+ worker._enqueue_host_mirror()
452
+
453
+ worker._enqueue_row_dma.assert_not_called()
454
+
455
+
456
+def test_hisparse_finish_forward_submits_lazy_post_forward_transfer(monkeypatch):
457
+ runtime = SimpleNamespace(eager_host_mirror=False)
458
+ worker = _make_hisparse_worker()
459
+ worker.is_host_writer = True
460
+ worker.cache_handles = [SimpleNamespace(runtime=runtime, num_actual_tokens=0)]
461
+ transfer = SparseKVPageTransfer(7, 2, (1,), after_forward=True)
462
+ worker._post_forward_transfers = [transfer]
463
+ worker._forward_ready_event = MagicMock()
464
+ worker._enqueue_host_mirror = MagicMock()
465
+ worker._submit_transfers = MagicMock()
466
+ worker._dma_submitted = False
467
+ worker._submitted_mirror_layers = set()
468
+ worker._finish_mirror_phase = MagicMock()
469
+ worker._release_completed_dma_descriptors = MagicMock()
470
+ stream = MagicMock()
471
+ monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream)
472
+
473
+ worker.finish_forward()
474
+
475
+ worker._finish_mirror_phase.assert_called_once_with(worker._forward_ready_event)
476
+ worker._submit_transfers.assert_called_once_with([transfer])
477
+ worker.host_write_event.record.assert_called_once_with(stream)
478
+
479
+
480
+def test_hisparse_prefill_mirrors_source_groups_and_flushes_partial_group():
481
+ slots = torch.tensor([7, 8], dtype=torch.int64)
482
+ source_indices = [0, 0, 1, 1, 1, 1, 2]
483
+ handles = [
484
+ SimpleNamespace(
485
+ runtime=SimpleNamespace(
486
+ eager_host_mirror=False,
487
+ is_group_leader=False,
488
+ resident_source_index=source_index,
489
+ ),
490
+ decode_batch=False,
491
+ host_mirror_required=layer_index < 5,
492
+ num_actual_tokens=2 if layer_index < 5 else 0,
493
+ num_decode_tokens=0,
494
+ req_id_per_token=torch.empty(0, dtype=torch.int32),
495
+ mirror_slot_mapping=slots,
496
+ )
497
+ for layer_index, source_index in enumerate(source_indices)
498
+ ]
499
+ worker = _make_hisparse_worker()
500
+ worker.is_host_writer = True
501
+ worker.cache_handles = handles
502
+ worker._set_row_mirrors((SparseKVRowMirror((0, 0, 0), 7, 2),))
503
+ worker._per_layer_mirrored = set()
504
+ worker._submitted_mirror_layers = set()
505
+ worker._layer_ready_events = tuple(MagicMock() for _ in handles)
506
+ worker._enqueue_row_dma = MagicMock()
507
+
508
+ for layer_index in range(5):
509
+ worker._enqueue_layer_mirror(layer_index)
510
+
511
+ worker._enqueue_row_dma.assert_called_once_with(
512
+ (0, 1), ready_event=worker._layer_ready_events[1]
513
+ )
514
+ worker._enqueue_host_mirror(ready_event=worker._layer_ready_events[-1])
515
+ assert worker._enqueue_row_dma.call_args_list[1].args == ((2, 3, 4),)
516
+ assert worker._enqueue_row_dma.call_args_list[1].kwargs == {
517
+ "ready_event": worker._layer_ready_events[-1]
518
+ }
519
+
520
+
521
+def test_hisparse_finish_forward_rejects_partial_per_layer_mirror():
522
+ slots = torch.tensor([7, 8], dtype=torch.int64)
523
+ runtime = SimpleNamespace(
524
+ eager_host_mirror=False,
525
+ is_group_leader=False,
526
+ invalidate_written_slots=MagicMock(),
527
+ )
528
+ handles = [
529
+ SimpleNamespace(
530
+ runtime=runtime,
531
+ decode_batch=False,
532
+ host_mirror_required=True,
533
+ num_actual_tokens=2,
534
+ num_decode_tokens=0,
535
+ req_id_per_token=torch.empty(0, dtype=torch.int32),
536
+ mirror_slot_mapping=slots,
537
+ )
538
+ for _ in range(2)
539
+ ]
540
+ worker = _make_hisparse_worker()
541
+ worker.is_host_writer = True
542
+ worker.cache_handles = handles
543
+ worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 2),))
544
+ worker._per_layer_mirrored = {0}
545
+ worker._submitted_mirror_layers = set()
546
+ worker._enqueue_row_dma = MagicMock()
547
+
548
+ with pytest.raises(RuntimeError, match="did not mirror every active layer"):
549
+ worker._enqueue_host_mirror()
550
+
551
+
552
+def test_hisparse_finish_forward_orders_next_forward_after_dma(monkeypatch):
553
+ current_stream = MagicMock()
554
+ worker = _make_hisparse_worker()
555
+ worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0"))
556
+ worker.is_host_writer = True
557
+ worker._dma_submitted = True
558
+ worker.host_write_event = MagicMock()
559
+ worker.cache_handles = [
560
+ SimpleNamespace(
561
+ runtime=SimpleNamespace(eager_host_mirror=True), num_actual_tokens=0
562
+ )
563
+ ]
564
+ worker._post_forward_transfers = []
565
+ worker._pending_dma_descriptors = deque()
566
+ worker._dma_free_descriptors = []
567
+ worker._forward_ready_event = MagicMock()
568
+ worker._finish_mirror_phase = MagicMock()
569
+ monkeypatch.setattr(
570
+ hisparse_worker_module, "current_stream", lambda: current_stream
571
+ )
572
+
573
+ worker.finish_forward()
574
+
575
+ current_stream.wait_event.assert_called_once_with(worker.host_write_event)
576
+ assert not worker._dma_submitted
577
+
578
+
579
+def test_hisparse_finish_forward_mirrors_standalone_mtp_cache(monkeypatch):
580
+ dst_slots = torch.tensor([7, 8], dtype=torch.int64)
581
+ runtime = SimpleNamespace(
582
+ eager_host_mirror=True,
583
+ is_group_leader=False,
584
+ invalidate_written_slots=MagicMock(),
585
+ )
586
+ inactive = SimpleNamespace(
587
+ runtime=runtime,
588
+ decode_batch=False,
589
+ host_mirror_required=False,
590
+ num_actual_tokens=0,
591
+ num_decode_tokens=0,
592
+ req_id_per_token=None,
593
+ mirror_slot_mapping=dst_slots,
594
+ )
595
+ mtp = SimpleNamespace(
596
+ runtime=runtime,
597
+ decode_batch=True,
598
+ host_mirror_required=True,
599
+ num_actual_tokens=2,
600
+ num_decode_tokens=2,
601
+ req_id_per_token=torch.tensor([0, 1], dtype=torch.int32),
602
+ mirror_slot_mapping=dst_slots,
603
+ )
604
+ worker = _make_hisparse_worker()
605
+ worker.is_host_writer = True
606
+ worker.cache_handles = [inactive, mtp]
607
+ worker._set_row_mirrors((SparseKVRowMirror((0,), 7, 2),))
608
+ worker._per_layer_mirrored = set()
609
+ worker._submitted_mirror_layers = set()
610
+ worker._enqueue_row_dma = MagicMock()
611
+
612
+ worker._enqueue_host_mirror()
613
+
614
+ worker._enqueue_row_dma.assert_called_once_with((1,), ready_event=None)
615
+
616
+
617
+def test_hisparse_shared_host_reader_skips_mirror(monkeypatch):
618
+ """A non-writer TP rank must not mirror rows into the shared host pool."""
619
+ dst_slots = torch.tensor([7, 8], dtype=torch.int64)
620
+ leader = SimpleNamespace(
621
+ eager_host_mirror=True,
622
+ is_group_leader=True,
623
+ invalidate_written_slots=MagicMock(),
624
+ )
625
+ handle = SimpleNamespace(
626
+ runtime=leader,
627
+ decode_batch=True,
628
+ host_mirror_required=True,
629
+ num_actual_tokens=2,
630
+ num_decode_tokens=1,
631
+ req_id_per_token=torch.tensor([0], dtype=torch.int32),
632
+ mirror_slot_mapping=dst_slots,
633
+ )
634
+ worker = _make_hisparse_worker()
635
+ worker.is_host_writer = False
636
+ worker.cache_handles = [handle]
637
+ worker._set_row_mirrors((SparseKVRowMirror((0,), 7, 2),))
638
+ worker._per_layer_mirrored = set()
639
+ worker._enqueue_row_dma = MagicMock()
640
+
641
+ worker._enqueue_host_mirror()
642
+
643
+ worker._enqueue_row_dma.assert_not_called()
644
+ leader.invalidate_written_slots.assert_called_once()
645
+
646
+
647
+def test_hisparse_shared_host_reader_skips_transfer_completion():
648
+ """A non-writer TP rank must not acknowledge host writes."""
649
+ worker = _make_hisparse_worker()
650
+ worker.is_host_writer = False
651
+ worker.kernel_block_size = 1
652
+ worker._enqueued_transfer_ids = []
653
+ worker._pending_transfer_events = []
654
+
655
+ worker._record_transfer_completion([SparseKVPageTransfer(1, 2, (3,), True)])
656
+
657
+ assert worker._enqueued_transfer_ids == []
658
+ assert worker._pending_transfer_events == []
659
+
660
+
661
+def test_hisparse_writer_records_transfer_completion_after_dma(monkeypatch):
662
+ worker = _make_hisparse_worker()
663
+ worker.is_host_writer = True
664
+ worker.dma_stream = MagicMock()
665
+ worker._enqueued_transfer_ids = []
666
+ worker._pending_transfer_events = deque()
667
+ event = MagicMock()
668
+ monkeypatch.setattr(torch, "Event", lambda: event)
669
+
670
+ worker._record_transfer_completion(
671
+ [
672
+ SparseKVPageTransfer(3, 2, (1,), False),
673
+ SparseKVPageTransfer(7, 4, (5,), True),
674
+ ]
675
+ )
676
+
677
+ event.record.assert_called_once_with(worker.dma_stream)
678
+ assert worker._enqueued_transfer_ids == [3, 7]
679
+ assert worker._pending_transfer_events == deque([(event, (3, 7))])
680
+
681
+
682
+@pytest.mark.parametrize("is_host_writer", [False, True])
683
+def test_hisparse_step_waits_for_previous_host_write(monkeypatch, is_host_writer):
684
+ worker = _make_hisparse_worker()
685
+ worker.is_host_writer = is_host_writer
686
+ worker.hot_backing = SimpleNamespace(device=torch.device("cuda:1"))
687
+ host_write_events = (MagicMock(), MagicMock())
688
+ worker.host_write_events = host_write_events
689
+ worker.host_write_event = host_write_events[1]
690
+ worker._next_host_write_event = 0
691
+ worker.host_caches = ()
692
+ worker.host_num_blocks = 1
693
+ worker._post_forward_transfers = []
694
+ worker._pending_invalid_block_ids = []
695
+ worker.cache_handles = [
696
+ SimpleNamespace(
697
+ runtime=SimpleNamespace(eager_host_mirror=True),
698
+ decode_batch=False,
699
+ host_mirror_required=False,
700
+ num_actual_tokens=0,
701
+ num_decode_tokens=0,
702
+ req_id_per_token=None,
703
+ )
704
+ ]
705
+ worker._per_layer_mirrored = set()
706
+ worker._submitted_mirror_layers = set()
707
+ worker._layer_mirror_callbacks = ()
708
+ stream = MagicMock()
709
+ monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream)
710
+
711
+ worker.start_step(
712
+ SimpleNamespace(
713
+ host_block_copies=[],
714
+ command=None,
715
+ source_block_ids=[],
716
+ row_mirrors={},
717
+ all_context_pages_resident=True,
718
+ row_mirrors_from_resident=False,
719
+ ),
720
+ None,
721
+ )
722
+ worker.start_step(
723
+ SimpleNamespace(
724
+ host_block_copies=[],
725
+ command=None,
726
+ source_block_ids=[],
727
+ row_mirrors={},
728
+ all_context_pages_resident=True,
729
+ row_mirrors_from_resident=False,
730
+ ),
731
+ None,
732
+ )
733
+
734
+ assert stream.wait_event.call_args_list == [
735
+ call(host_write_events[1]),
736
+ call(host_write_events[0]),
737
+ ]
738
+ assert worker.host_write_event is host_write_events[1]
739
+
740
+
741
+@pytest.mark.parametrize("tp_rank, expected_copies", [(0, 1), (1, 0)])
742
+def test_hisparse_shared_host_block_copy_has_one_writer(
743
+ monkeypatch, tp_rank, expected_copies
744
+):
745
+ worker = _make_hisparse_worker()
746
+ worker.shared_host_region = object()
747
+ worker.host_caches = (torch.empty(4, 1),)
748
+ worker.host_num_blocks = 4
749
+ previous_event = MagicMock()
750
+ copies = (KVCacheBlockCopy(0, 1),)
751
+ copy_blocks = MagicMock()
752
+ tp_group = MagicMock()
753
+ monkeypatch.setattr(
754
+ hisparse_worker_module, "get_tensor_model_parallel_rank", lambda: tp_rank
755
+ )
756
+ monkeypatch.setattr(hisparse_worker_module, "get_tp_group", lambda: tp_group)
757
+ monkeypatch.setattr(
758
+ hisparse_worker_module, "copy_kv_cache_blocks_inplace", copy_blocks
759
+ )
760
+
761
+ worker._copy_host_blocks(copies, previous_event)
762
+
763
+ assert copy_blocks.call_count == expected_copies
764
+ tp_group.barrier.assert_called_once_with()
765
+ assert worker.take_completed_host_copies() == [1]
766
+ assert worker.take_completed_host_copies() == []
767
+
768
+
769
+def test_hisparse_empty_step_does_not_replay_stale_host_mirror(monkeypatch):
770
+ handle = SimpleNamespace(
771
+ runtime=SimpleNamespace(eager_host_mirror=True),
772
+ decode_batch=True,
773
+ num_actual_tokens=2,
774
+ num_decode_tokens=2,
775
+ req_id_per_token=torch.tensor([0, 1]),
776
+ mirror_slot_mapping=torch.tensor([4, 5]),
777
+ )
778
+ worker = _make_hisparse_worker()
779
+ worker.is_host_writer = True
780
+ worker.hot_backing = SimpleNamespace(device=torch.device("cpu"))
781
+ worker.host_write_event = MagicMock()
782
+ worker._forward_ready_event = MagicMock()
783
+ worker.host_caches = ()
784
+ worker.host_num_blocks = 1
785
+ worker.cache_handles = [handle]
786
+ worker._per_layer_mirrored = set()
787
+ worker._submitted_mirror_layers = set()
788
+ worker._layer_mirror_callbacks = (MagicMock(),)
789
+ worker._post_forward_transfers = []
790
+ worker._pending_invalid_block_ids = []
791
+ worker._enqueue_host_mirror = MagicMock(wraps=worker._enqueue_host_mirror)
792
+ stream = MagicMock()
793
+ monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream)
794
+
795
+ worker.start_step(
796
+ SimpleNamespace(
797
+ host_block_copies=[],
798
+ command=None,
799
+ source_block_ids=[],
800
+ row_mirrors={},
801
+ all_context_pages_resident=True,
802
+ row_mirrors_from_resident=False,
803
+ ),
804
+ None,
805
+ )
806
+ worker.finish_forward()
807
+
808
+ worker._enqueue_host_mirror.assert_called_once_with(worker._forward_ready_event)
809
+ assert handle.num_actual_tokens == 0
810
+ torch.testing.assert_close(handle.mirror_slot_mapping, torch.tensor([4, 5]))
811
+
812
+
813
+def test_hisparse_cache_handles_join_index_groups_during_construction(monkeypatch):
814
+ """Followers must not allocate duplicate runtime state before profiling."""
815
+ config = SimpleNamespace(
816
+ scheduler_config=SimpleNamespace(
817
+ max_num_seqs=2,
818
+ max_num_batched_tokens=2,
819
+ async_scheduling=False,
820
+ ),
821
+ speculative_config=None,
822
+ kv_transfer_config=None,
823
+ )
824
+ resolved = hisparse_runtime_module.ResolvedHiSparseConfig(
825
+ top_k=4,
826
+ device_buffer_size=8,
827
+ )
828
+ monkeypatch.setattr(hisparse_runtime_module, "_has_hisparse_ops", lambda: True)
829
+ monkeypatch.setattr(
830
+ hisparse_runtime_module.ResolvedHiSparseConfig,
831
+ "from_vllm_config",
832
+ classmethod(lambda cls, vllm_config, model_top_k: resolved),
833
+ )
834
+ shared_states: list[object] = []
835
+ streams: list[object] = []
836
+
837
+ def create_shared_state(_device, _max_rows, _top_k):
838
+ shared_states.append(object())
839
+ return shared_states[-1]
840
+
841
+ def create_stream(_device):
842
+ streams.append(object())
843
+ return streams[-1]
844
+
845
+ monkeypatch.setattr(
846
+ hisparse_runtime_module, "_create_shared_topk_state", create_shared_state
847
+ )
848
+ monkeypatch.setattr(hisparse_runtime_module, "_create_copy_stream", create_stream)
849
+ monkeypatch.setattr(index_group_module, "_create_side_stream", lambda _: object())
850
+ monkeypatch.setattr(index_group_module, "_create_event", lambda: object())
851
+ index_group_builder = index_group_module.SparseMLAIndexGroupBuilder(
852
+ torch.empty((2, 4), dtype=torch.int32)
853
+ )
854
+
855
+ def make_cache_handle(is_leader: bool):
856
+ index_group, _ = index_group_builder.register_layer(is_leader)
857
+ cache_handle = hisparse_runtime_module.create_hisparse_cache_handle(
858
+ config,
859
+ model_top_k=4,
860
+ is_index_group_leader=is_leader,
861
+ row_width=8,
862
+ kv_dtype=torch.float32,
863
+ index_group=index_group,
864
+ device="cpu",
865
+ )
866
+ assert cache_handle is not None
867
+ return cache_handle
868
+
869
+ first_leader = make_cache_handle(True)
870
+ first_index_group = index_group_builder.current_group
871
+ first_follower = make_cache_handle(False)
872
+ second_leader = make_cache_handle(True)
873
+ second_index_group = index_group_builder.current_group
874
+ second_follower = make_cache_handle(False)
875
+
876
+ assert first_index_group is not None
877
+ assert second_index_group is not None
878
+ assert first_follower.runtime.index_group is first_leader.runtime.index_group
879
+ assert second_follower.runtime.index_group is second_leader.runtime.index_group
880
+ assert first_leader.runtime.index_group is not second_leader.runtime.index_group
881
+ assert first_leader.runtime.index_group.copy_stream is first_index_group.side_stream
882
+ assert (
883
+ first_leader.runtime.index_group.logical_topk_ready
884
+ is first_index_group.logical_topk_ready
885
+ )
886
+ assert (
887
+ second_leader.runtime.index_group.copy_stream is second_index_group.side_stream
888
+ )
889
+ assert (
890
+ second_leader.runtime.index_group.logical_topk_ready
891
+ is second_index_group.logical_topk_ready
892
+ )
893
+ assert len(shared_states) == 2
894
+ assert streams == []
895
896
897
class _TestReplaySSMMixer(MambaMixer2):
@@ -30,7 +914,8 @@ def _packed_replayssm_cache(num_blocks: int) -> torch.Tensor:
914
return torch.full((num_blocks, 1, 1, 80), 0, dtype=torch.int8)
915
916
33
-def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group():
917
+@pytest.mark.parametrize("layers_only", [False, True])
918
+def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(layers_only):
919
mixers = [_TestReplaySSMMixer() for _ in range(3)]
920
layer_names = [f"layers.{i}.mixer" for i in range(3)]
921
ctx = dict(zip(layer_names, mixers))
@@ -45,7 +930,10 @@ def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group():
930
SimpleNamespace(layer_names=[layer_names[1]]),
931
]
932
48
- bind_kv_cache(kv_cache, ctx, [], kv_cache_groups=kv_cache_groups)
933
+ if layers_only:
934
+ bind_kv_cache_to_layers(kv_cache, ctx, kv_cache_groups=kv_cache_groups)
935
+ else:
936
+ bind_kv_cache(kv_cache, ctx, [], kv_cache_groups=kv_cache_groups)
937
938
assert (
939
mixers[0]._replayssm_ring_start.data_ptr()
vllm/_custom_ops.py
+15
-1
@@ -3008,6 +3008,9 @@ def cp_gather_and_upconvert_fp8_kv_cache(
3008
workspace_starts: torch.Tensor,
3009
batch_size: int,
3010
seq_starts: torch.Tensor | None = None,
3011
+ host_cache: torch.Tensor | None = None,
3012
+ host_row_ids: torch.Tensor | None = None,
3013
+ device_row_ids: torch.Tensor | None = None,
3014
) -> None:
3015
"""Gather and upconvert FP8 KV cache to BF16 workspace.
3016
@@ -3018,9 +3021,20 @@ def cp_gather_and_upconvert_fp8_kv_cache(
3021
workspace_starts: Workspace start offsets [num_reqs]
3022
batch_size: Number of requests
3023
seq_starts: Optional source sequence offsets [num_reqs]
3024
+ host_cache: Optional pinned host rows used for non-resident entries
3025
+ host_row_ids: Host source row for each remapped cache row
3026
+ device_row_ids: Device source row, or -1, for each remapped cache row
3027
"""
3028
torch.ops._C_cache_ops.cp_gather_and_upconvert_fp8_kv_cache(
3023
- src_cache, dst, block_table, workspace_starts, batch_size, seq_starts
3029
+ src_cache,
3030
+ dst,
3031
+ block_table,
3032
+ workspace_starts,
3033
+ batch_size,
3034
+ seq_starts,
3035
+ host_cache,
3036
+ host_row_ids,
3037
+ device_row_ids,
3038
)
3039
3040
vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py
+4
-1
@@ -35,7 +35,10 @@ def fused_rope_unified_mla_kv_cache_update_impl(
35
layer_name: LayerNameType,
36
) -> torch.Tensor:
37
layer_name = _resolve_layer_name(layer_name)
38
- attn_metadata, _, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
38
+ attn_metadata, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(
39
+ layer_name
40
+ )
41
+ attn_layer.prepare_kv_cache_update(attn_metadata)
42
if layer_slot_mapping is not None:
43
ops.concat_and_cache_mla_rope_fused(
44
positions,
vllm/config/__init__.py
+2
-1
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from vllm.config.attention import AttentionConfig
4
+from vllm.config.attention import AttentionConfig, HiSparseConfig
5
from vllm.config.cache import CacheConfig
6
from vllm.config.compilation import (
7
CompilationConfig,
@@ -68,6 +68,7 @@ from vllm.config.weight_transfer import WeightTransferConfig
68
__all__ = [
69
# From vllm.config.attention
70
"AttentionConfig",
71
+ "HiSparseConfig",
72
# From vllm.config.cache
73
"CacheConfig",
74
# From vllm.config.compilation
vllm/config/attention.py
+24
-1
@@ -4,7 +4,7 @@
4
from dataclasses import field
5
from typing import Any, Literal
6
7
-from pydantic import field_validator
7
+from pydantic import Field, field_validator
8
9
from vllm.config.utils import config
10
from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
@@ -14,6 +14,25 @@ IndexerKVDType = Literal["auto", "bf16", "fp8", "mxfp4", "nvfp4"]
14
MiniMaxM3MSADecodeBackend = Literal["triton", "cutlass"]
15
16
17
+@config
18
+class HiSparseConfig:
19
+ """Configuration for HiSparse sparse-MLA KV offloading."""
20
+
21
+ device_buffer_size: int | None = Field(default=None, gt=0)
22
+ """Total per-request GPU hot-buffer rows, including the newest-token slot.
23
+
24
+ Defaults to one top-k per decode query plus one top-k of LRU slack. The
25
+ physical allocation is rounded up to the GPU cache block size selected
26
+ from the active backends.
27
+ """
28
+
29
+ eager_host_mirror: bool = True
30
+ """Mirror decode-written KV rows to the host pool during the forward so
31
+ page spills complete without moving data. When disabled, decode rows stay
32
+ resident-only and evicted pages are copied to host at spill time. Prefill
33
+ rows are always mirrored during the forward."""
34
+
35
+
36
@config
37
class AttentionConfig:
38
"""Configuration for attention mechanisms in vLLM."""
@@ -67,6 +86,10 @@ class AttentionConfig:
86
indexer). Quantized formats (fp8, mxfp4, nvfp4) require indexer kernel
87
support in the backend."""
88
89
+ hisparse_config: HiSparseConfig | None = None
90
+ """HiSparse host-resident KV configuration. Setting this enables experimental
91
+ Model Runner V2-only HiSparse sparse-MLA decode hot-buffering."""
92
+
93
use_non_causal: bool = False
94
"""Whether to use non-causal (bidirectional) attention."""
95
vllm/config/kv_transfer.py
+34
@@ -13,6 +13,40 @@ KVConsumer = Literal["kv_consumer", "kv_both"]
13
KVRole = Literal[KVProducer, KVConsumer]
14
15
16
+def hisparse_host_pool_gib(
17
+ kv_transfer_config: "KVTransferConfig | None",
18
+) -> float | None:
19
+ if kv_transfer_config is None:
20
+ return None
21
+ if kv_transfer_config.kv_connector == "MultiConnector":
22
+ connectors = kv_transfer_config.kv_connector_extra_config.get("connectors", [])
23
+ else:
24
+ connectors = [
25
+ {
26
+ "kv_connector": kv_transfer_config.kv_connector,
27
+ "kv_connector_extra_config": (
28
+ kv_transfer_config.kv_connector_extra_config
29
+ ),
30
+ }
31
+ ]
32
+ entries = [
33
+ connector.get("kv_connector_extra_config", {})
34
+ for connector in connectors
35
+ if connector.get("kv_connector") == "HiSparseConnector"
36
+ ]
37
+ if len(entries) > 1:
38
+ raise ValueError("Only one HiSparseConnector may be configured")
39
+ if not entries:
40
+ return None
41
+ host_pool_gib = entries[0].get("host_pool_gib")
42
+ if host_pool_gib is None:
43
+ raise ValueError("HiSparseConnector requires host_pool_gib")
44
+ host_pool_gib = float(host_pool_gib)
45
+ if host_pool_gib <= 0:
46
+ raise ValueError("HiSparseConnector host_pool_gib must be positive")
47
+ return host_pool_gib
48
+
49
+
50
def kv_buffer_device_default_factory() -> str:
51
from vllm.platforms import current_platform
52
vllm/config/speculative.py
+7
@@ -1919,6 +1919,13 @@ class SpeculativeConfig:
1919
def uses_draft_model(self) -> bool:
1920
return self.method == "draft_model"
1921
1922
+ def uses_draft_kv_cache(self) -> bool:
1923
+ return (
1924
+ self.use_eagle()
1925
+ or self.uses_draft_model()
1926
+ or self.uses_extract_hidden_states()
1927
+ )
1928
+
1929
def uses_extract_hidden_states(self) -> bool:
1930
return self.method == "extract_hidden_states"
1931
vllm/config/vllm.py
+51
@@ -666,6 +666,14 @@ class VllmConfig:
666
667
@property
668
def use_v2_model_runner(self) -> bool:
669
+ if self.attention_config.hisparse_config is not None:
670
+ if envs.VLLM_USE_V2_MODEL_RUNNER is False:
671
+ raise ValueError(
672
+ "HiSparse requires Model Runner V2; remove "
673
+ "VLLM_USE_V2_MODEL_RUNNER=0."
674
+ )
675
+ return True
676
+
677
if getattr(self, "watermark_config", None) is not None:
678
if envs.VLLM_USE_V2_MODEL_RUNNER is False:
679
logger.info_once(
@@ -1559,6 +1567,39 @@ class VllmConfig:
1567
self._maybe_disable_dynamic_sd_for_data_parallel()
1568
self._maybe_override_dynamic_sd_cudagraph_mode()
1569
1570
+ if self.attention_config.hisparse_config is not None:
1571
+ if not current_platform.is_cuda():
1572
+ raise ValueError("HiSparse currently requires NVIDIA CUDA.")
1573
+ if self.parallel_config.pipeline_parallel_size > 1:
1574
+ raise ValueError("HiSparse does not support pipeline parallelism.")
1575
+ if self.parallel_config.decode_context_parallel_size > 1:
1576
+ raise ValueError(
1577
+ "HiSparse does not support decode context parallelism."
1578
+ )
1579
+ if self.model_config is not None and not hasattr(
1580
+ self.model_config.hf_config, "index_topk"
1581
+ ):
1582
+ raise ValueError(
1583
+ "HiSparse is only supported for DSA models with index_topk."
1584
+ )
1585
+ if self.kv_transfer_config is not None and (
1586
+ self.kv_transfer_config.kv_connector
1587
+ not in (
1588
+ None,
1589
+ "NixlConnector",
1590
+ "MooncakeStoreConnector",
1591
+ "MultiConnector",
1592
+ )
1593
+ ):
1594
+ logger.warning(
1595
+ "HiSparse host-resident KV is configured with connector "
1596
+ "%s. NixlConnector (GPU-staged host imports) and "
1597
+ "MooncakeStoreConnector (shared-store offload) are the "
1598
+ "validated paths; other connectors are treated as "
1599
+ "debug/fallback paths.",
1600
+ self.kv_transfer_config.kv_connector,
1601
+ )
1602
+
1603
self._normalize_piecewise_cudagraph_mode(
1604
breakable_cudagraph_enabled=breakable_cudagraph_enabled
1605
)
@@ -1937,6 +1978,16 @@ class VllmConfig:
1978
# Default to enable HMA if not explicitly disabled by user or logic above.
1979
self.scheduler_config.disable_hybrid_kv_cache_manager = False
1980
1981
+ if (
1982
+ self.attention_config.hisparse_config is not None
1983
+ and self.scheduler_config.disable_hybrid_kv_cache_manager
1984
+ ):
1985
+ raise ValueError(
1986
+ "HiSparse requires the hybrid KV cache manager; remove "
1987
+ "--disable-hybrid-kv-cache-manager or use connectors that "
1988
+ "support HMA."
1989
+ )
1990
+
1991
if self.compilation_config.debug_dump_path:
1992
self.compilation_config.debug_dump_path = (
1993
self.compilation_config.debug_dump_path.absolute().expanduser()
vllm/distributed/kv_transfer/kv_connector/factory.py
+6
@@ -197,6 +197,12 @@ KVConnectorFactory.register_connector(
197
"MultiConnector",
198
)
199
200
+KVConnectorFactory.register_connector(
201
+ "HiSparseConnector",
202
+ "vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector",
203
+ "HiSparseConnector",
204
+)
205
+
206
KVConnectorFactory.register_connector(
207
"MoRIIOConnector",
208
"vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector",
vllm/distributed/kv_transfer/kv_connector/utils.py
+6
@@ -60,6 +60,7 @@ class KVOutputAggregator:
60
# [req_id -> n_remaining_workers]
61
self._recv_remaining_count = dict[str, int]()
62
self._send_remaining_count = dict[str, int]()
63
+ self._failed_recving_pending = set[str]()
64
self._expected_finished_count = expected_finished_count
65
66
@classmethod
@@ -156,6 +157,10 @@ class KVOutputAggregator:
157
combined_kv_cache_events.increment_workers(1)
158
159
invalid_block_ids |= kv_output.invalid_block_ids
160
+ self._failed_recving_pending |= kv_output.failed_recving
161
+
162
+ failed_recving = self._failed_recving_pending & finished_recving
163
+ self._failed_recving_pending -= failed_recving
164
165
# select output of the worker specified by output_rank
166
output = outputs[output_rank]
@@ -168,6 +173,7 @@ class KVOutputAggregator:
173
kv_cache_events=combined_kv_cache_events or None,
174
kv_connector_worker_meta=aggregated_kv_connector_worker_meta or None,
175
invalid_block_ids=invalid_block_ids,
176
+ failed_recving=failed_recving,
177
expected_finished_count=self._expected_finished_count,
178
)
179
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+49
-6
@@ -34,8 +34,8 @@ The class provides the following primitives:
34
save_kv_layer() - starts saving KV for layer i (maybe async)
35
wait_for_save() - blocks until all saves are done
36
37
- get_finished() - called with ids of finished requests, returns
38
- ids of requests that have completed async sending/recving.
37
+ get_transfer_results() - returns async send/receive completions and
38
+ receive failures in one snapshot.
39
build_connector_worker_meta() - builds metadata to be sent
40
back to the scheduler-side connector
41
"""
@@ -43,6 +43,7 @@ The class provides the following primitives:
43
import enum
44
from abc import ABC, abstractmethod
45
from collections.abc import Callable, Iterable
46
+from dataclasses import dataclass, field
47
from typing import TYPE_CHECKING, Any, Literal
48
49
import torch
@@ -63,7 +64,7 @@ if TYPE_CHECKING:
64
)
65
from vllm.forward_context import ForwardContext
66
from vllm.v1.core.block_pool import BlockPool
66
- from vllm.v1.core.kv_cache_manager import KVCacheBlocks
67
+ from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
68
from vllm.v1.kv_cache_interface import KVCacheConfig
69
from vllm.v1.request import Request
70
@@ -82,6 +83,19 @@ CopyBlocksOp = Callable[
83
logger = init_logger(__name__)
84
85
86
+@dataclass
87
+class KVConnectorTransferResults:
88
+ """Asynchronous transfer completions from one worker snapshot.
89
+
90
+ Failed receives also appear in ``finished_recving`` so the scheduler can
91
+ release the request from its transfer wait state.
92
+ """
93
+
94
+ finished_sending: set[str] = field(default_factory=set)
95
+ finished_recving: set[str] = field(default_factory=set)
96
+ failed_recving: set[str] = field(default_factory=set)
97
+
98
+
99
class SupportsHMA(ABC):
100
"""
101
The class that indicates the corresponding connector supports hybrid memory
@@ -211,6 +225,7 @@ class KVConnectorBase_V1(ABC):
225
else:
226
raise ValueError("kv_transfer_config must be set for KVConnectorBase_V1")
227
self._kv_cache_config = kv_cache_config
228
+ self._kv_cache_manager: KVCacheManager | None = None
229
self._role = role
230
231
@property
@@ -278,6 +293,15 @@ class KVConnectorBase_V1(ABC):
293
"""
294
return
295
296
+ # TODO(NickLucche): group model-runner lifecycle hooks in the interface.
297
+ def finish_forward(self) -> None:
298
+ """Notify the connector that the model no longer reads this step's KV."""
299
+ return
300
+
301
+ def reset_capture_state(self) -> None:
302
+ """Reset worker state mutated while capturing CUDA graphs."""
303
+ return
304
+
305
def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata):
306
"""
307
Handle preempted requests or evicted blocks BEFORE they are overwritten.
@@ -368,6 +392,16 @@ class KVConnectorBase_V1(ABC):
392
"""
393
return None, None
394
395
+ def get_transfer_results(
396
+ self, finished_req_ids: set[str]
397
+ ) -> KVConnectorTransferResults:
398
+ """Return completed sends, receives, and receive failures together."""
399
+ finished_sending, finished_recving = self.get_finished(finished_req_ids)
400
+ return KVConnectorTransferResults(
401
+ finished_sending=set(finished_sending or ()),
402
+ finished_recving=set(finished_recving or ()),
403
+ )
404
+
405
def get_block_ids_with_load_errors(self) -> set[int]:
406
"""
407
Get the set of block IDs that failed to load.
@@ -380,9 +414,9 @@ class KVConnectorBase_V1(ABC):
414
- Applies to both sync- and async-loading requests.
415
- Async loading: failed blocks may be reported in any forward pass
416
up to and including the pass where the request ID is returned by
383
- `get_finished()`. Even if failures occur, the request must still
384
- be reported via `get_finished()`, and the failed block IDs must
385
- appear here no later than that same pass.
417
+ `get_transfer_results()`. Even if failures occur, the request
418
+ must still be reported as finished receiving, and the failed
419
+ block IDs must appear here no later than that same pass.
420
- Sync loading: failed blocks should be reported in the forward
421
pass in which they are detected.
422
"""
@@ -436,6 +470,11 @@ class KVConnectorBase_V1(ABC):
470
# Scheduler-side methods
471
# ==============================
472
473
+ def bind_kv_cache_manager(self, kv_cache_manager: "KVCacheManager") -> None:
474
+ """Bind the scheduler's cache manager after it has been constructed."""
475
+ self._kv_cache_manager = kv_cache_manager
476
+ self.bind_gpu_block_pool(kv_cache_manager.block_pool)
477
+
478
def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None:
479
"""
480
Bind the GPU block pool to the connector for per-GPU block status tracking.
@@ -583,6 +622,10 @@ class KVConnectorBase_V1(ABC):
622
"""
623
return ()
624
625
+ def has_pending_block_frees(self) -> bool:
626
+ """Whether pending transfers can release blocks instead of preemption."""
627
+ return False
628
+
629
def has_pending_push_work(self) -> bool:
630
"""Return True if the connector has push-mode work that requires
631
the engine main loop to keep stepping (e.g. a P-side request whose
vllm/distributed/kv_transfer/kv_connector/v1/hisparse/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py
new
+329
@@ -0,0 +1,329 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Connector boundary for HiSparse scheduler and worker state."""
4
+
5
+from __future__ import annotations
6
+
7
+from dataclasses import dataclass
8
+from typing import TYPE_CHECKING, Any
9
+
10
+import torch
11
+
12
+from vllm.config import VllmConfig
13
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
14
+ KVConnectorBase_V1,
15
+ KVConnectorMetadata,
16
+ KVConnectorRole,
17
+ KVConnectorWorkerMetadata,
18
+ SupportsHMA,
19
+)
20
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
21
+ HiSparseConnectorWorker,
22
+)
23
+from vllm.v1.attention.backend import AttentionMetadata
24
+from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
25
+from vllm.v1.core.sched.output import SchedulerOutput
26
+from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
27
+from vllm.v1.hisparse.types import SparseKVOffloadCommand, SparseKVRowMirror
28
+from vllm.v1.outputs import KVConnectorOutput
29
+
30
+if TYPE_CHECKING:
31
+ from vllm.forward_context import ForwardContext
32
+ from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
33
+ from vllm.v1.hisparse.coordinator import HiSparseCoordinator
34
+ from vllm.v1.kv_cache_interface import KVCacheConfig
35
+ from vllm.v1.request import Request
36
+
37
+
38
+@dataclass
39
+class HiSparseConnectorMetadata(KVConnectorMetadata):
40
+ command: SparseKVOffloadCommand | None
41
+ host_block_copies: tuple[KVCacheBlockCopy, ...]
42
+ source_block_ids: tuple[int, ...]
43
+ row_mirrors: dict[str, tuple[SparseKVRowMirror, ...]]
44
+ all_context_pages_resident: bool
45
+
46
+
47
+@dataclass
48
+class HiSparseConnectorWorkerMetadata(KVConnectorWorkerMetadata):
49
+ enqueued_transfer_counts: dict[int, int]
50
+ completed_transfer_counts: dict[int, int]
51
+ completed_host_copy_dst_ids: tuple[int, ...] = ()
52
+
53
+ def aggregate(self, other: KVConnectorWorkerMetadata) -> KVConnectorWorkerMetadata:
54
+ assert isinstance(other, HiSparseConnectorWorkerMetadata)
55
+
56
+ def add_counts(first: dict[int, int], second: dict[int, int]) -> dict[int, int]:
57
+ combined = first.copy()
58
+ for transfer_id, count in second.items():
59
+ combined[transfer_id] = combined.get(transfer_id, 0) + count
60
+ return combined
61
+
62
+ return HiSparseConnectorWorkerMetadata(
63
+ enqueued_transfer_counts=add_counts(
64
+ self.enqueued_transfer_counts, other.enqueued_transfer_counts
65
+ ),
66
+ completed_transfer_counts=add_counts(
67
+ self.completed_transfer_counts, other.completed_transfer_counts
68
+ ),
69
+ # Every worker runs the same copies, so a union is the ack.
70
+ completed_host_copy_dst_ids=tuple(
71
+ {*self.completed_host_copy_dst_ids, *other.completed_host_copy_dst_ids}
72
+ ),
73
+ )
74
+
75
+
76
+class HiSparseConnectorScheduler:
77
+ def __init__(
78
+ self,
79
+ *,
80
+ async_speculative: bool,
81
+ draft_kv_lookahead: int = 0,
82
+ ) -> None:
83
+ self.coordinator: HiSparseCoordinator | None = None
84
+ self.async_speculative = async_speculative
85
+ self.draft_kv_lookahead = draft_kv_lookahead
86
+ self.requests: dict[str, Request] = {}
87
+
88
+ def bind_coordinator(self, coordinator: HiSparseCoordinator) -> None:
89
+ assert self.coordinator is None
90
+ self.coordinator = coordinator
91
+
92
+ def build_connector_meta(
93
+ self, scheduler_output: SchedulerOutput
94
+ ) -> HiSparseConnectorMetadata:
95
+ assert self.coordinator is not None
96
+ # HiSparse rebinds worker state every step, so the load must start
97
+ # before the forward rather than being deferred to post-forward.
98
+ scheduler_output.has_sync_kv_loads = True
99
+ scheduler_output.block_table_updates = (
100
+ self.coordinator.take_block_table_updates() or None
101
+ )
102
+ command = self.coordinator.build_offload_command()
103
+ host_block_copies = self.coordinator.take_host_block_copies()
104
+ source_group_id = self.coordinator.host_group_id
105
+ assert source_group_id is not None
106
+ source_block_ids = [
107
+ block_id
108
+ for request in scheduler_output.scheduled_new_reqs
109
+ for block_id in request.block_ids[source_group_id]
110
+ ]
111
+ for new_block_ids in scheduler_output.scheduled_cached_reqs.new_block_ids:
112
+ if new_block_ids is not None:
113
+ source_block_ids.extend(new_block_ids[source_group_id])
114
+ num_computed_tokens = {
115
+ request.req_id: request.num_computed_tokens
116
+ for request in scheduler_output.scheduled_new_reqs
117
+ }
118
+ num_computed_tokens.update(
119
+ zip(
120
+ scheduler_output.scheduled_cached_reqs.req_ids,
121
+ scheduler_output.scheduled_cached_reqs.num_computed_tokens,
122
+ )
123
+ )
124
+ scheduled_requests = tuple(
125
+ (
126
+ request_id,
127
+ num_computed_tokens[request_id],
128
+ scheduled_count,
129
+ )
130
+ for request_id, scheduled_count in (
131
+ scheduler_output.num_scheduled_tokens.items()
132
+ )
133
+ )
134
+ row_mirrors = {}
135
+ for request_id, scheduled_start, scheduled_count in scheduled_requests:
136
+ mirror_start = scheduled_start
137
+ if self.async_speculative:
138
+ mirror_start = max(
139
+ 0,
140
+ scheduled_start - self.requests[request_id].num_output_placeholders,
141
+ )
142
+ row_mirrors[request_id] = self.coordinator.build_row_mirrors(
143
+ (
144
+ (
145
+ request_id,
146
+ mirror_start,
147
+ scheduled_count
148
+ + scheduled_start
149
+ - mirror_start
150
+ + self.draft_kv_lookahead,
151
+ ),
152
+ ),
153
+ )
154
+ return HiSparseConnectorMetadata(
155
+ command,
156
+ host_block_copies,
157
+ tuple(source_block_ids),
158
+ row_mirrors,
159
+ self.coordinator.all_context_pages_resident(scheduled_requests),
160
+ )
161
+
162
+ def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
163
+ assert self.coordinator is not None
164
+ for request_id in connector_output.finished_recving or ():
165
+ self.coordinator.finish_host_import(
166
+ request_id, failed=request_id in connector_output.failed_recving
167
+ )
168
+ metadata = connector_output.kv_connector_worker_meta
169
+ if metadata is None:
170
+ return
171
+ assert isinstance(metadata, HiSparseConnectorWorkerMetadata)
172
+ self.coordinator.release_completed_host_copies(
173
+ metadata.completed_host_copy_dst_ids
174
+ )
175
+ self.coordinator.update_spills(
176
+ metadata.enqueued_transfer_counts,
177
+ metadata.completed_transfer_counts,
178
+ )
179
+
180
+
181
+class HiSparseConnector(KVConnectorBase_V1, SupportsHMA):
182
+ """Join the scheduler coordinator to the worker's transfer engine."""
183
+
184
+ @classmethod
185
+ def get_required_kvcache_layout(cls, vllm_config: VllmConfig) -> str:
186
+ return "BLHNC"
187
+
188
+ def __init__(
189
+ self,
190
+ vllm_config: VllmConfig,
191
+ role: KVConnectorRole,
192
+ kv_cache_config: KVCacheConfig,
193
+ ) -> None:
194
+ super().__init__(vllm_config, role, kv_cache_config)
195
+ if kv_cache_config.hisparse_host_num_blocks is None:
196
+ raise ValueError("HiSparseConnector requires a HiSparse host pool")
197
+ self.connector_scheduler: HiSparseConnectorScheduler | None = None
198
+ self.connector_worker: HiSparseConnectorWorker | None = None
199
+ if role == KVConnectorRole.SCHEDULER:
200
+ speculative_config = vllm_config.speculative_config
201
+ self.connector_scheduler = HiSparseConnectorScheduler(
202
+ async_speculative=bool(
203
+ vllm_config.scheduler_config.async_scheduling
204
+ and speculative_config is not None
205
+ ),
206
+ draft_kv_lookahead=vllm_config.num_lookahead_tokens,
207
+ )
208
+ elif role == KVConnectorRole.WORKER:
209
+ self.connector_worker = HiSparseConnectorWorker(
210
+ vllm_config, kv_cache_config
211
+ )
212
+ else:
213
+ raise ValueError(f"Unsupported KV connector role: {role}")
214
+
215
+ def bind_kv_cache_manager(self, kv_cache_manager: KVCacheManager) -> None:
216
+ assert self.connector_scheduler is not None
217
+ self.connector_scheduler.bind_coordinator(
218
+ get_hisparse_coordinator(kv_cache_manager)
219
+ )
220
+
221
+ @property
222
+ def requires_kv_delivery(self) -> bool:
223
+ return False
224
+
225
+ def has_pending_push_work(self) -> bool:
226
+ assert self.connector_scheduler is not None
227
+ assert self.connector_scheduler.coordinator is not None
228
+ return self.connector_scheduler.coordinator.has_pending_work()
229
+
230
+ def has_pending_block_frees(self) -> bool:
231
+ assert self.connector_scheduler is not None
232
+ assert self.connector_scheduler.coordinator is not None
233
+ return self.connector_scheduler.coordinator.has_pending_reclamation()
234
+
235
+ def finish_forward(self) -> None:
236
+ assert self.connector_worker is not None
237
+ self.connector_worker.finish_forward()
238
+
239
+ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None:
240
+ assert self.connector_worker is not None
241
+ self.connector_worker.register_kv_caches(kv_caches)
242
+
243
+ def reset_capture_state(self) -> None:
244
+ assert self.connector_worker is not None
245
+ self.connector_worker.reset_hot_state()
246
+
247
+ def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None:
248
+ assert self.connector_worker is not None
249
+ metadata = self._get_connector_metadata()
250
+ assert isinstance(metadata, HiSparseConnectorMetadata)
251
+ request_state_indices = kwargs.get("request_state_indices")
252
+ assert request_state_indices is None or isinstance(
253
+ request_state_indices, torch.Tensor
254
+ )
255
+ request_ids = kwargs.get("request_ids")
256
+ assert request_ids is None or isinstance(request_ids, list)
257
+ attn_metadata = kwargs.get("attn_metadata")
258
+ assert attn_metadata is None or isinstance(attn_metadata, dict)
259
+ self.connector_worker.start_step(
260
+ metadata,
261
+ request_state_indices,
262
+ request_ids,
263
+ num_tokens=int(kwargs.get("num_tokens") or 0),
264
+ )
265
+ self.connector_worker.prepare_forward(attn_metadata)
266
+
267
+ def wait_for_layer_load(self, layer_name: str) -> None:
268
+ return
269
+
270
+ def save_kv_layer(
271
+ self,
272
+ layer_name: str,
273
+ kv_layer: torch.Tensor,
274
+ attn_metadata: AttentionMetadata,
275
+ **kwargs: Any,
276
+ ) -> None:
277
+ return
278
+
279
+ def wait_for_save(self) -> None:
280
+ return
281
+
282
+ def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None:
283
+ assert self.connector_worker is not None
284
+ enqueued, completed = self.connector_worker.take_transfer_updates()
285
+ host_copies = self.connector_worker.take_completed_host_copies()
286
+ if not enqueued and not completed and not host_copies:
287
+ return None
288
+ return HiSparseConnectorWorkerMetadata(
289
+ enqueued_transfer_counts={transfer_id: 1 for transfer_id in enqueued},
290
+ completed_transfer_counts={transfer_id: 1 for transfer_id in completed},
291
+ completed_host_copy_dst_ids=tuple(host_copies),
292
+ )
293
+
294
+ def shutdown(self) -> None:
295
+ if self.connector_worker is not None:
296
+ self.connector_worker.shutdown()
297
+
298
+ def get_num_new_matched_tokens(
299
+ self, request: Request, num_computed_tokens: int
300
+ ) -> tuple[int | None, bool]:
301
+ return 0, False
302
+
303
+ def update_state_after_alloc(
304
+ self,
305
+ request: Request,
306
+ blocks: KVCacheBlocks,
307
+ num_external_tokens: int,
308
+ ) -> None:
309
+ assert self.connector_scheduler is not None
310
+ self.connector_scheduler.requests[request.request_id] = request
311
+
312
+ def build_connector_meta(
313
+ self, scheduler_output: SchedulerOutput
314
+ ) -> KVConnectorMetadata:
315
+ assert self.connector_scheduler is not None
316
+ return self.connector_scheduler.build_connector_meta(scheduler_output)
317
+
318
+ def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
319
+ assert self.connector_scheduler is not None
320
+ self.connector_scheduler.update_connector_output(connector_output)
321
+
322
+ def request_finished_all_groups(
323
+ self,
324
+ request: Request,
325
+ block_ids: tuple[list[int], ...],
326
+ ) -> tuple[bool, dict[str, Any] | None]:
327
+ assert self.connector_scheduler is not None
328
+ self.connector_scheduler.requests.pop(request.request_id, None)
329
+ return False, None
vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py
new
+832
@@ -0,0 +1,832 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""HiSparse worker-side host/hot state and data movement."""
4
+
5
+from __future__ import annotations
6
+
7
+from collections import deque
8
+from collections.abc import Mapping, Sequence
9
+from dataclasses import dataclass
10
+from functools import partial
11
+from typing import TYPE_CHECKING, Any, NamedTuple
12
+
13
+import numpy as np
14
+import torch
15
+
16
+from vllm import _custom_ops as ops
17
+from vllm.config import VllmConfig
18
+from vllm.distributed.parallel_state import (
19
+ get_tensor_model_parallel_rank,
20
+ get_tp_group,
21
+)
22
+from vllm.utils.torch_utils import current_stream
23
+from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
24
+from vllm.v1.hisparse.layout import HISPARSE_HOT_SUFFIX
25
+from vllm.v1.hisparse.runtime import HiSparseCacheHandle, release_pinned_state
26
+from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror
27
+from vllm.v1.kv_cache_interface import (
28
+ HiSparseHotSpec,
29
+ KVCacheConfig,
30
+)
31
+from vllm.v1.worker.utils import copy_kv_cache_blocks_inplace
32
+
33
+if TYPE_CHECKING:
34
+ from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
35
+ HiSparseConnectorMetadata,
36
+ )
37
+ from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
38
+
39
+
40
+class _DMADescriptors(NamedTuple):
41
+ src: torch.Tensor
42
+ dst: torch.Tensor
43
+ sizes: torch.Tensor
44
+ src_np: np.ndarray
45
+ dst_np: np.ndarray
46
+ sizes_np: np.ndarray
47
+
48
+
49
+@dataclass
50
+class _SlotMappingStaging:
51
+ stream: torch.Stream
52
+ event: torch.Event
53
+ slots: torch.Tensor
54
+ candidates: tuple[SparseKVRowMirror, ...] = ()
55
+ num_tokens: int = 0
56
+ source_index: int = 0
57
+
58
+
59
+def _allocate_dma_descriptors(size: int) -> _DMADescriptors:
60
+ src, dst, sizes = (torch.empty(size, dtype=torch.int64) for _ in range(3))
61
+ return _DMADescriptors(
62
+ src,
63
+ dst,
64
+ sizes,
65
+ src.numpy(),
66
+ dst.numpy(),
67
+ sizes.numpy(),
68
+ )
69
+
70
+
71
+def _get_hisparse_cache(
72
+ forward_context: dict[str, Any], layer_name: str
73
+) -> HiSparseCacheHandle:
74
+ attention_layer = forward_context[layer_name]
75
+ hisparse_cache = attention_layer.hisparse_cache
76
+ assert hisparse_cache is not None
77
+ return hisparse_cache
78
+
79
+
80
+def _flatten_row_mirrors(
81
+ row_mirrors: Mapping[str, tuple[SparseKVRowMirror, ...]],
82
+ request_ids: Sequence[str] | None,
83
+) -> tuple[SparseKVRowMirror, ...]:
84
+ ordered_ids = row_mirrors if request_ids is None else request_ids
85
+ return tuple(
86
+ mirror
87
+ for request_id in ordered_ids
88
+ for mirror in row_mirrors.get(request_id, ())
89
+ )
90
+
91
+
92
+def _select_written_row_mirrors(
93
+ candidates: tuple[SparseKVRowMirror, ...],
94
+ source_slots: np.ndarray,
95
+ source_index: int,
96
+) -> tuple[SparseKVRowMirror, ...]:
97
+ """Select and coalesce GPU-written rows from a scheduler-owned envelope."""
98
+ source_slots = source_slots[source_slots >= 0]
99
+ if source_slots.size == 0 or not candidates:
100
+ return ()
101
+ starts = np.asarray([mirror.source_starts for mirror in candidates], dtype=np.int64)
102
+ counts = np.fromiter((mirror.num_rows for mirror in candidates), dtype=np.int64)
103
+ destinations = np.fromiter(
104
+ [mirror.destination_start for mirror in candidates], dtype=np.int64
105
+ )
106
+ if source_index >= starts.shape[1]:
107
+ raise RuntimeError("HiSparse row-mirror source index is out of range.")
108
+ source_starts = starts[:, source_index]
109
+ order = np.argsort(source_starts)
110
+ matches = np.searchsorted(source_starts[order], source_slots, side="right") - 1
111
+ candidate_ids = order[np.maximum(matches, 0)]
112
+ offsets = source_slots - source_starts[candidate_ids]
113
+ valid = (matches >= 0) & (offsets < counts[candidate_ids])
114
+ if not np.all(valid):
115
+ # A page the scheduler could not map has no host destination, so its
116
+ # rows are simply not mirrorable this step.
117
+ source_slots = source_slots[valid]
118
+ candidate_ids = candidate_ids[valid]
119
+ offsets = offsets[valid]
120
+ if source_slots.size == 0:
121
+ return ()
122
+ sources = starts[candidate_ids] + offsets[:, None]
123
+ destinations = destinations[candidate_ids] + offsets
124
+ boundaries = (
125
+ np.flatnonzero(
126
+ (candidate_ids[1:] != candidate_ids[:-1])
127
+ | (np.diff(destinations) != 1)
128
+ | np.any(np.diff(sources, axis=0) != 1, axis=1)
129
+ )
130
+ + 1
131
+ )
132
+ run_starts = np.concatenate(([0], boundaries))
133
+ run_ends = np.concatenate((boundaries, [destinations.size]))
134
+ return tuple(
135
+ SparseKVRowMirror(
136
+ source_starts=tuple(int(value) for value in sources[start]),
137
+ destination_start=int(destinations[start]),
138
+ num_rows=int(end - start),
139
+ )
140
+ for start, end in zip(run_starts, run_ends, strict=True)
141
+ )
142
+
143
+
144
+def _create_hisparse_host_events(
145
+ shared_host_region: SharedOffloadRegion | None,
146
+ is_host_writer: bool,
147
+ device: torch.device,
148
+) -> tuple[torch.Event, torch.Event]:
149
+ if shared_host_region is None:
150
+ return torch.Event(), torch.Event()
151
+
152
+ events: tuple[torch.Event, torch.Event] | None = None
153
+ ipc_handles = None
154
+ if is_host_writer:
155
+ events = (
156
+ torch.cuda.Event(interprocess=True),
157
+ torch.cuda.Event(interprocess=True),
158
+ )
159
+ stream = current_stream()
160
+ for event in events:
161
+ event.record(stream)
162
+ ipc_handles = tuple(event.ipc_handle() for event in events)
163
+ ipc_handles = get_tp_group().broadcast_object(ipc_handles, src=0)
164
+ if events is None:
165
+ events = (
166
+ torch.cuda.Event.from_ipc_handle(device, ipc_handles[0]),
167
+ torch.cuda.Event.from_ipc_handle(device, ipc_handles[1]),
168
+ )
169
+ return events
170
+
171
+
172
+class HiSparseConnectorWorker:
173
+ """Own HiSparse host/hot state and execute its worker-side transfers."""
174
+
175
+ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> None:
176
+ self.vllm_config = vllm_config
177
+ self.kv_cache_config = kv_cache_config
178
+ self._initialized = False
179
+
180
+ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None:
181
+ forward_context = self.vllm_config.compilation_config.static_forward_context
182
+ cache_handles: list[HiSparseCacheHandle] = []
183
+ cache_layer_names: list[str] = []
184
+ for group in self.kv_cache_config.kv_cache_groups:
185
+ if not isinstance(group.kv_cache_spec, HiSparseHotSpec):
186
+ continue
187
+ for cache_name in group.layer_names:
188
+ assert cache_name.endswith(HISPARSE_HOT_SUFFIX)
189
+ layer_name = cache_name[: -len(HISPARSE_HOT_SUFFIX)]
190
+ cache_handles.append(_get_hisparse_cache(forward_context, layer_name))
191
+ cache_layer_names.append(layer_name)
192
+
193
+ if not cache_handles:
194
+ raise RuntimeError("HiSparse connector found no hot-cache handles.")
195
+ hot_backings: dict[int, torch.Tensor] = {}
196
+ registered_host_pools: dict[int, torch.Tensor] = {}
197
+ for cache in cache_handles:
198
+ hot_backing = cache.runtime.hot_backing
199
+ hot_backings[hot_backing.untyped_storage().data_ptr()] = hot_backing
200
+ registered_pool = cache.runtime.registered_host_pool
201
+ registered_host_pools[registered_pool.data_ptr()] = registered_pool
202
+ if len(hot_backings) != 1:
203
+ raise RuntimeError("HiSparse hot tensors must share one GPU backing.")
204
+ hot_backing = next(iter(hot_backings.values()))
205
+ pinned_host_pools = list(registered_host_pools.values())
206
+
207
+ resident = cache_handles[0].view
208
+ assert resident is not None
209
+ host_num_blocks = self.kv_cache_config.hisparse_host_num_blocks
210
+ assert host_num_blocks is not None
211
+ try:
212
+ self.initialize(
213
+ cache_handles,
214
+ cache_layer_names,
215
+ hot_backing,
216
+ self.vllm_config.scheduler_config.max_num_seqs,
217
+ host_num_blocks,
218
+ hot_backing.device,
219
+ pinned_host_pools,
220
+ )
221
+ except Exception:
222
+ release_pinned_state(
223
+ [cache.runtime for cache in cache_handles],
224
+ pinned_host_pools,
225
+ )
226
+ raise
227
+
228
+ def initialize(
229
+ self,
230
+ cache_handles: list[HiSparseCacheHandle],
231
+ cache_layer_names: list[str],
232
+ hot_backing: torch.Tensor,
233
+ max_num_reqs: int,
234
+ host_num_blocks: int,
235
+ device: torch.device,
236
+ pinned_host_pools: list[torch.Tensor],
237
+ *,
238
+ shared_host_region: SharedOffloadRegion | None = None,
239
+ is_host_writer: bool = True,
240
+ ) -> None:
241
+ if self._initialized:
242
+ raise RuntimeError("HiSparse connector worker is already initialized.")
243
+ resident = cache_handles[0].view
244
+ assert resident is not None
245
+ self.kernel_block_size = resident.block_size
246
+ self.host_num_blocks = host_num_blocks
247
+ self.pinned_host_pools = pinned_host_pools
248
+ self.shared_host_region = shared_host_region
249
+ self.is_host_writer = is_host_writer
250
+ self.dma_stream = (
251
+ torch.cuda.Stream(device=device) if self.is_host_writer else None
252
+ )
253
+ self._slot_mapping_staging = None
254
+ if self.is_host_writer:
255
+ max_mirror_rows = (
256
+ self.vllm_config.scheduler_config.max_num_batched_tokens
257
+ + max_num_reqs * (self.vllm_config.num_lookahead_tokens + 1)
258
+ )
259
+ self._slot_mapping_staging = _SlotMappingStaging(
260
+ stream=torch.Stream(device=device),
261
+ event=torch.Event(),
262
+ slots=torch.empty(
263
+ max_mirror_rows,
264
+ dtype=torch.int64,
265
+ pin_memory=True,
266
+ ),
267
+ )
268
+ self._dma_free_descriptors: list[_DMADescriptors] = []
269
+ self._pending_dma_descriptors: deque[tuple[torch.Event, _DMADescriptors]] = (
270
+ deque()
271
+ )
272
+ self._dma_submitted = False
273
+ self._per_layer_mirrored: set[int] = set()
274
+ self._submitted_mirror_layers: set[int] = set()
275
+ self._layer_ready_events = tuple(torch.Event() for _ in cache_handles)
276
+ self._forward_ready_event = torch.Event()
277
+ self._set_row_mirrors(())
278
+ self.cache_layer_names = cache_layer_names
279
+ self._group_leaders = tuple(
280
+ (layer_name, cache)
281
+ for layer_name, cache in zip(cache_layer_names, cache_handles, strict=True)
282
+ if cache.runtime.is_group_leader
283
+ )
284
+ self.host_write_events = _create_hisparse_host_events(
285
+ shared_host_region, is_host_writer, device
286
+ )
287
+ self.host_write_event = self.host_write_events[1]
288
+ self._next_host_write_event = 0
289
+ self.cache_handles = cache_handles
290
+ self.leader_runtimes = [cache.runtime for _, cache in self._group_leaders]
291
+ request_state_indices = {
292
+ indices.data_ptr(): indices
293
+ for cache in cache_handles
294
+ if (indices := cache.runtime.request_state_indices) is not None
295
+ }
296
+ if len(request_state_indices) != 1:
297
+ raise RuntimeError(
298
+ "HiSparse runtimes must share one request-state mapping."
299
+ )
300
+ self.request_state_indices = next(iter(request_state_indices.values()))
301
+ if self.request_state_indices.numel() != max_num_reqs:
302
+ raise RuntimeError(
303
+ "HiSparse request-state mapping does not match max_num_seqs."
304
+ )
305
+ self.hot_backing = hot_backing
306
+ self._pending_invalid_block_ids: list[int] = []
307
+ # Destination block ids of host copies this worker has run.
308
+ self._completed_host_copy_dst_ids: list[int] = []
309
+ self._post_forward_transfers: list[SparseKVPageTransfer] = []
310
+ self._enqueued_transfer_ids: list[int] = []
311
+ self._pending_transfer_events: deque[tuple[torch.Event, tuple[int, ...]]] = (
312
+ deque()
313
+ )
314
+ self._init_dma()
315
+ if self.is_host_writer:
316
+ for layer_index, handle in enumerate(cache_handles):
317
+ handle.submit_layer_mirror = partial(
318
+ self._enqueue_layer_mirror, layer_index
319
+ )
320
+ self._initialized = True
321
+
322
+ def set_request_state_indices(self, indices: torch.Tensor) -> None:
323
+ if indices.numel() > self.request_state_indices.numel():
324
+ raise ValueError(
325
+ "HiSparse request-state mapping exceeds max_num_seqs: "
326
+ f"{indices.numel()} > {self.request_state_indices.numel()}."
327
+ )
328
+ if torch.cuda.is_current_stream_capturing():
329
+ return
330
+ # Attention indexes persistent request state by input-batch row. Refresh
331
+ # that indirection after every batch compaction or reorder.
332
+ self.request_state_indices.fill_(-1)
333
+ self.request_state_indices[: indices.numel()].copy_(indices)
334
+ if self._pending_invalid_block_ids:
335
+ self.invalidate_blocks(self._pending_invalid_block_ids, indices)
336
+ self._pending_invalid_block_ids.clear()
337
+
338
+ def _init_dma(self) -> None:
339
+ host_caches = tuple(cache.runtime.host_cache for cache in self.cache_handles)
340
+ resident_caches = []
341
+ for cache in self.cache_handles:
342
+ assert cache.view is not None and cache.slot_mapping is not None
343
+ resident_caches.append(cache.view.cache)
344
+ self.host_caches = host_caches
345
+ self.resident_caches = tuple(resident_caches)
346
+
347
+ for resident_cache, host_cache in zip(resident_caches, host_caches):
348
+ row_bytes = resident_cache.shape[-1] * resident_cache.element_size()
349
+ if (
350
+ resident_cache.ndim != 3
351
+ or resident_cache.shape[1] != self.kernel_block_size
352
+ or resident_cache.stride(1) * resident_cache.element_size() != row_bytes
353
+ ):
354
+ raise RuntimeError("HiSparse DMA requires contiguous resident rows.")
355
+ if (
356
+ host_cache.ndim != 2
357
+ or not host_cache.is_contiguous()
358
+ or host_cache.shape[1] * host_cache.element_size() != row_bytes
359
+ ):
360
+ raise RuntimeError("HiSparse DMA requires contiguous host rows.")
361
+
362
+ def start_step(
363
+ self,
364
+ metadata: HiSparseConnectorMetadata,
365
+ request_state_indices: torch.Tensor | None,
366
+ request_ids: list[str] | None = None,
367
+ num_tokens: int = 0,
368
+ ) -> None:
369
+ self._stage_row_mirror_mapping(num_tokens)
370
+ previous_host_write_event = self.host_write_event
371
+ self.host_write_event = self.host_write_events[self._next_host_write_event]
372
+ self._next_host_write_event ^= 1
373
+ current_stream().wait_event(previous_host_write_event)
374
+ self._release_completed_dma_descriptors()
375
+ mirrors = _flatten_row_mirrors(metadata.row_mirrors, request_ids)
376
+ if self._slot_mapping_staging is not None:
377
+ self._slot_mapping_staging.candidates = mirrors
378
+ self._set_row_mirrors(mirrors)
379
+ self._dma_submitted = False
380
+ self._clear_forward_mirror_state()
381
+ for handle in self.cache_handles:
382
+ handle.all_context_pages_resident = metadata.all_context_pages_resident
383
+ handle.mirror_from_resident = True
384
+ self._copy_host_blocks(metadata.host_block_copies, previous_host_write_event)
385
+ transfers = (
386
+ metadata.command.page_transfers if metadata.command is not None else []
387
+ )
388
+ self._post_forward_transfers = [
389
+ transfer for transfer in transfers if transfer.after_forward
390
+ ]
391
+ self._submit_transfers(
392
+ [transfer for transfer in transfers if not transfer.after_forward]
393
+ )
394
+ self._pending_invalid_block_ids.extend(metadata.source_block_ids)
395
+ if request_state_indices is not None:
396
+ self.set_request_state_indices(request_state_indices)
397
+
398
+ def _clear_forward_mirror_state(self) -> None:
399
+ self._per_layer_mirrored.clear()
400
+ self._submitted_mirror_layers.clear()
401
+ for handle in self.cache_handles:
402
+ handle.decode_batch = False
403
+ handle.host_mirror_required = False
404
+ handle.num_actual_tokens = 0
405
+ handle.num_decode_tokens = 0
406
+ handle.req_id_per_token = None
407
+
408
+ def prepare_forward(self, attn_metadata: Mapping[str, Any] | None) -> None:
409
+ if attn_metadata is None:
410
+ return
411
+ for layer_name, handle in self._group_leaders:
412
+ metadata = attn_metadata.get(layer_name)
413
+ if metadata is not None:
414
+ handle.prepare_group_for_batch(metadata)
415
+
416
+ def _stage_row_mirror_mapping(self, num_tokens: int) -> None:
417
+ """Snapshot the rows this forward will write, off the compute stream.
418
+
419
+ The resident slot mapping is a persistent view bound at registration
420
+ time, so the rows are already staged by the time the forward launches.
421
+ """
422
+ state = self._slot_mapping_staging
423
+ if state is None or not num_tokens:
424
+ return
425
+ handle = self.cache_handles[0]
426
+ slots = handle.slot_mapping
427
+ assert slots is not None
428
+ source_index = handle.runtime.resident_source_index
429
+ start = state.num_tokens
430
+ end = start + num_tokens
431
+ if end > state.slots.shape[0]:
432
+ raise ValueError(
433
+ "HiSparse row mapping exceeds staging capacity: "
434
+ f"{end} > {state.slots.shape[0]}."
435
+ )
436
+ main_stream = current_stream()
437
+ state.stream.wait_stream(main_stream)
438
+ with torch.cuda.stream(state.stream):
439
+ state.slots[start:end].copy_(
440
+ slots[:num_tokens],
441
+ non_blocking=True,
442
+ )
443
+ state.event.record(state.stream)
444
+ state.num_tokens = end
445
+ state.source_index = source_index
446
+
447
+ def _resolve_row_mirrors(self, state: _SlotMappingStaging) -> None:
448
+ """Compute the final row mirrors from the staged GPU slot mapping."""
449
+ self._set_row_mirrors(
450
+ _select_written_row_mirrors(
451
+ state.candidates,
452
+ state.slots[: state.num_tokens].numpy(),
453
+ state.source_index,
454
+ )
455
+ )
456
+ state.num_tokens = 0
457
+
458
+ def _copy_host_blocks(
459
+ self,
460
+ host_block_copies: Sequence[KVCacheBlockCopy],
461
+ previous_host_write_event: torch.Event,
462
+ ) -> None:
463
+ if not host_block_copies:
464
+ return
465
+ self._completed_host_copy_dst_ids.extend(
466
+ copy.dst_block_id for copy in host_block_copies
467
+ )
468
+ if self.shared_host_region is None or get_tensor_model_parallel_rank() == 0:
469
+ if self.host_caches:
470
+ previous_host_write_event.synchronize()
471
+ copy_kv_cache_blocks_inplace(
472
+ self.host_caches,
473
+ self.host_num_blocks,
474
+ host_block_copies,
475
+ )
476
+ if self.shared_host_region is not None:
477
+ get_tp_group().barrier()
478
+
479
+ def invalidate_blocks(
480
+ self, block_ids: list[int], request_state_indices: torch.Tensor
481
+ ) -> None:
482
+ """Invalidate recycled host slots in this worker's leader runtimes."""
483
+ if not block_ids:
484
+ return
485
+ device = self.cache_handles[0].runtime.device
486
+ staging = torch.tensor(block_ids, dtype=torch.int32, pin_memory=True)
487
+ blocks = staging.to(device, dtype=torch.int32, non_blocking=True)
488
+ offsets = torch.arange(self.kernel_block_size, dtype=torch.int32, device=device)
489
+ slots = (blocks[:, None] * self.kernel_block_size + offsets[None, :]).flatten()
490
+ sorted_slots = torch.sort(slots).values
491
+ state_indices = request_state_indices.to(device=device, dtype=torch.long)
492
+ for runtime in self.leader_runtimes:
493
+ runtime.invalidate_sorted_slots(sorted_slots, state_indices)
494
+
495
+ def reset_hot_state(self) -> None:
496
+ for runtime in self.leader_runtimes:
497
+ runtime.reset_hot_state()
498
+
499
+ def _release_completed_dma_descriptors(self) -> None:
500
+ pending = self._pending_dma_descriptors
501
+ while pending and pending[0][0].query():
502
+ _, descriptors = pending.popleft()
503
+ self._dma_free_descriptors.append(descriptors)
504
+
505
+ def _acquire_dma_descriptors(self, size: int) -> _DMADescriptors:
506
+ for index, descriptors in enumerate(self._dma_free_descriptors):
507
+ if descriptors.src.numel() >= size:
508
+ return self._dma_free_descriptors.pop(index)
509
+ return _allocate_dma_descriptors(size)
510
+
511
+ def _submit_dma_descriptors(
512
+ self,
513
+ descriptors: _DMADescriptors,
514
+ descriptor_count: int,
515
+ transfer_ids: tuple[int, ...] = (),
516
+ ready_event: torch.Event | None = None,
517
+ ) -> None:
518
+ stream = self.dma_stream
519
+ assert stream is not None
520
+ if ready_event is None:
521
+ stream.wait_stream(current_stream())
522
+ else:
523
+ stream.wait_event(ready_event)
524
+ completion_event = torch.Event()
525
+ with torch.cuda.stream(stream):
526
+ ops.swap_blocks_batch(
527
+ descriptors.src[:descriptor_count],
528
+ descriptors.dst[:descriptor_count],
529
+ descriptors.sizes[:descriptor_count],
530
+ )
531
+ self.host_write_event.record(stream)
532
+ completion_event.record(stream)
533
+ self._pending_dma_descriptors.append((completion_event, descriptors))
534
+ self._dma_submitted = True
535
+ if transfer_ids:
536
+ self._pending_transfer_events.append((completion_event, transfer_ids))
537
+ self._enqueued_transfer_ids.extend(transfer_ids)
538
+
539
+ def _set_row_mirrors(self, mirrors: tuple[SparseKVRowMirror, ...]) -> None:
540
+ self._row_mirrors = mirrors
541
+ self._row_mirror_destination_starts = np.fromiter(
542
+ (mirror.destination_start for mirror in mirrors),
543
+ dtype=np.int64,
544
+ count=len(mirrors),
545
+ )
546
+ self._row_mirror_counts = np.fromiter(
547
+ (mirror.num_rows for mirror in mirrors),
548
+ dtype=np.int64,
549
+ count=len(mirrors),
550
+ )
551
+ if mirrors:
552
+ self._row_mirror_source_starts = np.asarray(
553
+ [mirror.source_starts for mirror in mirrors], dtype=np.int64
554
+ )
555
+ if self._row_mirror_source_starts.ndim != 2:
556
+ raise RuntimeError("HiSparse DMA source mappings must be rectangular.")
557
+ else:
558
+ self._row_mirror_source_starts = np.empty((0, 0), dtype=np.int64)
559
+ self._row_mirror_num_rows = int(self._row_mirror_counts.sum())
560
+
561
+ def _enqueue_row_dma(
562
+ self, layer_indices: Sequence[int], ready_event: torch.Event | None = None
563
+ ) -> None:
564
+ if (
565
+ not layer_indices
566
+ or not self._row_mirror_num_rows
567
+ or not self.is_host_writer
568
+ ):
569
+ return
570
+ mirrors = self._row_mirrors
571
+ num_layers = len(layer_indices)
572
+ descriptor_count = len(mirrors) * num_layers
573
+ descriptors = self._acquire_dma_descriptors(descriptor_count)
574
+ destination_starts = self._row_mirror_destination_starts
575
+ row_counts = self._row_mirror_counts
576
+ for descriptor_offset, layer_index in enumerate(layer_indices):
577
+ cache = self.cache_handles[layer_index]
578
+ source_index = cache.runtime.resident_source_index
579
+ if source_index >= self._row_mirror_source_starts.shape[1]:
580
+ raise RuntimeError("HiSparse row DMA source index is out of range.")
581
+ source_rows = self._row_mirror_source_starts[:, source_index]
582
+ source = self.resident_caches[layer_index]
583
+ destination = self.host_caches[layer_index]
584
+ row_bytes = source.shape[-1] * source.element_size()
585
+ if (
586
+ source.stride(1) * source.element_size() != row_bytes
587
+ or destination.shape[1] * destination.element_size() != row_bytes
588
+ ):
589
+ raise RuntimeError("HiSparse row DMA requires contiguous rows.")
590
+ source_blocks, source_row_offsets = np.divmod(
591
+ source_rows, self.kernel_block_size
592
+ )
593
+ source_out_of_range = (
594
+ np.any(source_blocks < 0)
595
+ or np.any(source_blocks >= source.shape[0])
596
+ or np.any(source_row_offsets + row_counts > self.kernel_block_size)
597
+ )
598
+ if source_out_of_range:
599
+ raise RuntimeError("HiSparse row DMA source is out of range.")
600
+ if np.any(destination_starts < 0) or np.any(
601
+ destination_starts + row_counts > destination.shape[0]
602
+ ):
603
+ raise RuntimeError("HiSparse row DMA index is out of range.")
604
+ descriptor_slice = slice(descriptor_offset, descriptor_count, num_layers)
605
+ descriptors.src_np[descriptor_slice] = (
606
+ source.data_ptr()
607
+ + source_blocks * source.stride(0) * source.element_size()
608
+ + source_row_offsets * row_bytes
609
+ )
610
+ descriptors.dst_np[descriptor_slice] = (
611
+ destination.data_ptr() + destination_starts * row_bytes
612
+ )
613
+ descriptors.sizes_np[descriptor_slice] = row_counts * row_bytes
614
+
615
+ self._submit_dma_descriptors(
616
+ descriptors, descriptor_count, ready_event=ready_event
617
+ )
618
+
619
+ def _enqueue_layer_mirror(self, layer_index: int) -> None:
620
+ handle = self.cache_handles[layer_index]
621
+ if not handle.host_mirror_required:
622
+ return
623
+ if layer_index in self._per_layer_mirrored:
624
+ raise RuntimeError(f"HiSparse layer {layer_index} mirrored twice.")
625
+ state = self._slot_mapping_staging
626
+ if state is not None and state.num_tokens:
627
+ # The staging stream waited on the compute stream before the
628
+ # forward launched, so this only drains pre-forward work that is
629
+ # already queued ahead of the layer kernels: no GPU bubble.
630
+ state.event.synchronize()
631
+ self._resolve_row_mirrors(state)
632
+ self._per_layer_mirrored.add(layer_index)
633
+ next_layer = layer_index + 1
634
+ if (
635
+ next_layer < len(self.cache_handles)
636
+ and self.cache_handles[next_layer].runtime.resident_source_index
637
+ == handle.runtime.resident_source_index
638
+ ):
639
+ return
640
+ ready_event = self._layer_ready_events[layer_index]
641
+ ready_event.record()
642
+ pending_layers = tuple(
643
+ sorted(self._per_layer_mirrored - self._submitted_mirror_layers)
644
+ )
645
+ self._enqueue_row_dma(
646
+ pending_layers,
647
+ ready_event=ready_event,
648
+ )
649
+ self._submitted_mirror_layers.update(pending_layers)
650
+
651
+ def _record_transfer_completion(
652
+ self, transfers: list[SparseKVPageTransfer]
653
+ ) -> None:
654
+ if not transfers or not self.is_host_writer:
655
+ return
656
+ stream = self.dma_stream
657
+ assert stream is not None
658
+ completion_event = torch.Event()
659
+ completion_event.record(stream)
660
+ transfer_ids = tuple(transfer.transfer_id for transfer in transfers)
661
+ self._pending_transfer_events.append((completion_event, transfer_ids))
662
+ self._enqueued_transfer_ids.extend(transfer_ids)
663
+
664
+ def _submit_transfers(self, transfers: list[SparseKVPageTransfer]) -> None:
665
+ if self.cache_handles[0].runtime.eager_host_mirror:
666
+ self._record_transfer_completion(transfers)
667
+ else:
668
+ self._enqueue_transfers(transfers)
669
+
670
+ def _enqueue_transfers(self, transfers: list[SparseKVPageTransfer]) -> None:
671
+ if not transfers or not self.is_host_writer:
672
+ return
673
+ num_layers = len(self.cache_handles)
674
+ descriptor_count = len(transfers) * num_layers
675
+ descriptors = self._acquire_dma_descriptors(descriptor_count)
676
+ destination_rows = np.fromiter(
677
+ (
678
+ transfer.destination_block_id * self.kernel_block_size
679
+ for transfer in transfers
680
+ ),
681
+ dtype=np.int64,
682
+ count=len(transfers),
683
+ )
684
+ source_blocks_by_transfer = np.asarray(
685
+ [transfer.source_block_ids for transfer in transfers], dtype=np.int64
686
+ )
687
+ if source_blocks_by_transfer.ndim != 2:
688
+ raise RuntimeError(
689
+ "HiSparse spill DMA source mappings must be rectangular."
690
+ )
691
+ for layer_index, cache in enumerate(self.cache_handles):
692
+ source_index = cache.runtime.resident_source_index
693
+ if source_index >= source_blocks_by_transfer.shape[1]:
694
+ raise RuntimeError("HiSparse spill DMA source index is out of range.")
695
+ source_blocks = source_blocks_by_transfer[:, source_index]
696
+ source = self.resident_caches[layer_index]
697
+ destination = self.host_caches[layer_index]
698
+ if np.any(source_blocks < 0) or np.any(source_blocks >= source.shape[0]):
699
+ raise RuntimeError("HiSparse spill DMA source is out of range.")
700
+ if np.any(destination_rows < 0) or np.any(
701
+ destination_rows + self.kernel_block_size > destination.shape[0]
702
+ ):
703
+ raise RuntimeError("HiSparse spill DMA destination is out of range.")
704
+ row_bytes = source.shape[-1] * source.element_size()
705
+ descriptor_slice = slice(layer_index, descriptor_count, num_layers)
706
+ descriptors.src_np[descriptor_slice] = (
707
+ source.data_ptr()
708
+ + source_blocks * source.stride(0) * source.element_size()
709
+ )
710
+ descriptors.dst_np[descriptor_slice] = (
711
+ destination.data_ptr() + destination_rows * row_bytes
712
+ )
713
+ descriptors.sizes_np[descriptor_slice] = self.kernel_block_size * row_bytes
714
+ self._submit_dma_descriptors(
715
+ descriptors,
716
+ descriptor_count,
717
+ transfer_ids=tuple(transfer.transfer_id for transfer in transfers),
718
+ )
719
+
720
+ def _enqueue_host_mirror(
721
+ self,
722
+ ready_event: torch.Event | None = None,
723
+ ) -> None:
724
+ active = [
725
+ (index, handle)
726
+ for index, handle in enumerate(self.cache_handles)
727
+ if handle.num_actual_tokens != 0
728
+ ]
729
+ if not active:
730
+ return
731
+
732
+ def mirror_key(handle: HiSparseCacheHandle) -> tuple:
733
+ slots = handle.mirror_slot_mapping
734
+ return (
735
+ handle.num_actual_tokens,
736
+ handle.num_decode_tokens,
737
+ handle.decode_batch,
738
+ handle.host_mirror_required,
739
+ handle.runtime.eager_host_mirror,
740
+ None if slots is None else slots.data_ptr(),
741
+ )
742
+
743
+ keys = {mirror_key(handle) for _, handle in active}
744
+ if len(keys) > 1:
745
+ raise RuntimeError(
746
+ f"HiSparse cache layers disagree on mirror metadata: {list(keys)}."
747
+ )
748
+ cache = active[0][1]
749
+ if not cache.host_mirror_required:
750
+ return
751
+ dst_slots = cache.mirror_slot_mapping
752
+ if dst_slots is None:
753
+ raise RuntimeError("HiSparse host mirror has no source slot mapping.")
754
+ num_rows = min(cache.num_actual_tokens, dst_slots.numel())
755
+ if num_rows == 0:
756
+ return
757
+ if self.is_host_writer:
758
+ expected_layers = {index for index, _ in active}
759
+ if self._per_layer_mirrored and self._per_layer_mirrored != expected_layers:
760
+ raise RuntimeError(
761
+ "HiSparse per-layer DMA did not mirror every active layer: "
762
+ f"expected {sorted(expected_layers)}, got "
763
+ f"{sorted(self._per_layer_mirrored)}."
764
+ )
765
+ pending_layers = tuple(
766
+ sorted(expected_layers - self._submitted_mirror_layers)
767
+ )
768
+ if pending_layers:
769
+ self._enqueue_row_dma(pending_layers, ready_event=ready_event)
770
+ self._submitted_mirror_layers.update(pending_layers)
771
+ assert cache.req_id_per_token is not None
772
+ for _, handle in active:
773
+ if handle.runtime.is_group_leader:
774
+ handle.runtime.invalidate_written_slots(
775
+ dst_slots[:num_rows],
776
+ cache.req_id_per_token[:num_rows],
777
+ )
778
+
779
+ def _finish_mirror_phase(self, ready_event: torch.Event | None = None) -> None:
780
+ state = self._slot_mapping_staging
781
+ if state is not None and state.num_tokens:
782
+ # Measured free: blocking here costs no throughput (26.7 vs 27.6
783
+ # gen tok/s) and lets every step mirror exactly the written rows.
784
+ state.event.synchronize()
785
+ self._resolve_row_mirrors(state)
786
+ self._enqueue_host_mirror(ready_event)
787
+ self._clear_forward_mirror_state()
788
+
789
+ def finish_forward(self) -> None:
790
+ compute_stream = current_stream()
791
+ self._forward_ready_event.record()
792
+ self._finish_mirror_phase(self._forward_ready_event)
793
+ transfers = self._post_forward_transfers
794
+ self._post_forward_transfers = []
795
+ self._submit_transfers(transfers)
796
+ if self.is_host_writer:
797
+ if self._dma_submitted:
798
+ compute_stream.wait_event(self.host_write_event)
799
+ self._dma_submitted = False
800
+ else:
801
+ self.host_write_event.record(compute_stream)
802
+ self._release_completed_dma_descriptors()
803
+
804
+ def take_completed_host_copies(self) -> list[int]:
805
+ """Drain host copies this worker has enqueued for this step."""
806
+ completed = self._completed_host_copy_dst_ids
807
+ self._completed_host_copy_dst_ids = []
808
+ return completed
809
+
810
+ def take_transfer_updates(self) -> tuple[list[int], list[int]]:
811
+ enqueued = self._enqueued_transfer_ids
812
+ self._enqueued_transfer_ids = []
813
+ completed: list[int] = []
814
+ while self._pending_transfer_events:
815
+ event, transfer_ids = self._pending_transfer_events[0]
816
+ if not event.query():
817
+ break
818
+ self._pending_transfer_events.popleft()
819
+ completed.extend(transfer_ids)
820
+ return enqueued, completed
821
+
822
+ def shutdown(self) -> None:
823
+ if not self._initialized:
824
+ return
825
+ if self._slot_mapping_staging is not None:
826
+ self._slot_mapping_staging.stream.synchronize()
827
+ if self.dma_stream is not None:
828
+ self.dma_stream.synchronize()
829
+ release_pinned_state(
830
+ [cache.runtime for cache in self.cache_handles], self.pinned_host_pools
831
+ )
832
+ self._initialized = False
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py
+3
-3
@@ -702,9 +702,9 @@ class LMCacheMPConnectorUpstream(KVConnectorBase_V1):
702
- Applies to both sync- and async-loading requests.
703
- Async loading: failed blocks may be reported in any forward pass
704
up to and including the pass where the request ID is returned by
705
- `get_finished()`. Even if failures occur, the request must still
706
- be reported via `get_finished()`, and the failed block IDs must
707
- appear here no later than that same pass.
705
+ `get_transfer_results()`. Even if failures occur, the request
706
+ must still be reported as finished receiving, and the failed
707
+ block IDs must appear here no later than that same pass.
708
- Sync loading: failed blocks should be reported in the forward
709
pass in which they are detected.
710
"""
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+17
-15
@@ -635,6 +635,7 @@ class MooncakeConnectorScheduler:
635
):
636
self.vllm_config = vllm_config
637
self.block_size = vllm_config.cache_config.block_size
638
+ self.kv_cache_config = kv_cache_config
639
640
assert vllm_config.kv_transfer_config
641
self.is_kv_producer: bool = (
@@ -649,7 +650,7 @@ class MooncakeConnectorScheduler:
650
not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
651
and any(
652
not isinstance(g.kv_cache_spec, FullAttentionSpec)
652
- for g in kv_cache_config.kv_cache_groups
653
+ for g in kv_cache_config.transfer_groups
654
)
655
)
656
# GDN is represented as a MambaSpec in vLLM. This Mooncake MambaSpec
@@ -670,7 +671,7 @@ class MooncakeConnectorScheduler:
671
(g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size)
672
if isinstance(g.kv_cache_spec, SlidingWindowSpec)
673
else (0, self.block_size)
673
- for g in kv_cache_config.kv_cache_groups
674
+ for g in kv_cache_config.transfer_groups
675
]
676
# cdiv(n_tokens, block_size) gives blocks/window; add 1 to
677
# conservatively account for boundary overlap.
@@ -684,11 +685,17 @@ class MooncakeConnectorScheduler:
685
block_ids: tuple[list[int], ...] | list[list[int]],
686
) -> list[list[int]]:
687
"""Clip per-group block IDs to sliding window size."""
687
- if len(block_ids) == 0 or not self._is_hma_required:
688
- return list(block_ids)
688
+ if len(block_ids) == 0:
689
+ return []
690
+ if len(block_ids) == len(self.kv_cache_config.transfer_group_ids):
691
+ selected = list(block_ids)
692
+ else:
693
+ selected = list(self.kv_cache_config.select_transfer_block_ids(block_ids))
694
+ if len(selected) == 0 or not self._is_hma_required:
695
+ return selected
696
return [
697
blocks[-self.blocks_per_sw[i] :] if self.blocks_per_sw[i] > 0 else blocks
691
- for i, blocks in enumerate(block_ids)
698
+ for i, blocks in enumerate(selected)
699
]
700
701
def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int:
@@ -1046,18 +1053,13 @@ class MooncakeConnectorWorker:
1053
1054
self._tp_size: dict[EngineId, int] = {self.engine_id: self.tp_size}
1055
self._layer_specs: dict[str, KVCacheSpec] = {}
1049
- for group in kv_cache_config.kv_cache_groups:
1056
+ for group in kv_cache_config.transfer_groups:
1057
group_spec = group.kv_cache_spec
1058
specs_by_layer = getattr(group_spec, "kv_cache_specs", {})
1059
for layer_name in group.layer_names:
1060
self._layer_specs[layer_name] = specs_by_layer.get(
1061
layer_name, group_spec
1062
)
1056
- self._layer_group_indices: dict[str, int] = {
1057
- layer: group_index
1058
- for group_index, group in enumerate(kv_cache_config.kv_cache_groups)
1059
- for layer in group.layer_names
1060
- }
1063
self.transfer_topo = TransferTopology(
1064
tp_rank=self.tp_rank,
1065
tp_size=self.tp_size,
@@ -1441,7 +1443,7 @@ class MooncakeConnectorWorker:
1443
block_arange = np.arange(self._physical_blocks_per_logical_kv_block).reshape(
1444
1, -1
1445
)
1444
- group_specs = self.kv_cache_config.kv_cache_groups
1446
+ group_specs = self.kv_cache_config.transfer_groups
1447
return [
1448
BlockTable.map_to_kernel_blocks(
1449
np.array(group),
@@ -1494,7 +1496,7 @@ class MooncakeConnectorWorker:
1496
local_block_ids_by_group: list[list[int]] = []
1497
remote_block_ids_by_group: list[list[int]] = []
1498
has_block_error = False
1497
- group_specs = self.kv_cache_config.kv_cache_groups
1499
+ group_specs = self.kv_cache_config.transfer_groups
1500
for group_index, (local_group, remote_group) in enumerate(
1501
zip(send_meta.local_block_ids, remote_block_ids_per_group)
1502
):
@@ -1750,7 +1752,7 @@ class MooncakeConnectorWorker:
1752
self.registered_layer_names.append(layer_name)
1753
self.registered_layer_indices.append(layer_index)
1754
self.registered_group_indices.append(
1753
- self._layer_group_indices[layer_name]
1755
+ self.kv_cache_config.transfer_group_index_by_layer[layer_name]
1756
)
1757
storage = cache.untyped_storage()
1758
storage_addr = storage.data_ptr()
@@ -2121,7 +2123,7 @@ class MooncakeConnectorWorker:
2123
) -> list[TransferRegion]:
2124
if not group_indices:
2125
group_indices = [
2124
- self._layer_group_indices.get(layer_name, 0)
2126
+ self.kv_cache_config.transfer_group_index_by_layer.get(layer_name, 0)
2127
for layer_name in layer_names
2128
]
2129
return _expand_transfer_regions(
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
+13
-4
@@ -3,18 +3,20 @@
3
"""External-store cache-hit coordinator for MooncakeStoreConnector."""
4
5
from collections.abc import Sequence
6
-from typing import cast
6
+from typing import NamedTuple, cast
7
8
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
9
chunk_hashes_for_block_size,
10
)
11
from vllm.utils.math_utils import cdiv
12
from vllm.v1.core.block_pool import BlockPool
13
-from vllm.v1.core.kv_cache_coordinator import SpecGroup
13
from vllm.v1.core.kv_cache_utils import (
14
BlockHash,
15
KVCacheBlock,
16
)
17
+from vllm.v1.core.single_type_kv_cache_manager import (
18
+ SingleTypeKVCacheManager,
19
+)
20
from vllm.v1.kv_cache_interface import (
21
FullAttentionSpec,
22
KVCacheGroupSpec,
@@ -25,6 +27,13 @@ from vllm.v1.kv_cache_interface import (
27
from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
28
29
30
+class StoreSpecGroup(NamedTuple):
31
+ spec: KVCacheSpec
32
+ group_ids: list[int]
33
+ manager_cls: type[SingleTypeKVCacheManager]
34
+ use_eagle: bool
35
+
36
+
37
class ExternalCachedBlockPool:
38
"""Duck-typed BlockPool backed by a ``(group_id, hash)`` exists set."""
39
@@ -121,7 +130,7 @@ class MooncakeStoreCoordinator:
130
"""Mirrors KVCacheCoordinator.verify_and_split_kv_cache_groups but
131
dispatches via spec_manager_map (we don't allocate managers).
132
"""
124
- attention_groups: list[SpecGroup] = []
133
+ attention_groups: list[StoreSpecGroup] = []
134
for i, g in enumerate(self.kv_cache_groups):
135
# Skip groups that opt out of prefix caching (e.g. GLM-5.3-Flash
136
# kpool tail): per-request scratch, never shareable, so they must
@@ -143,7 +152,7 @@ class MooncakeStoreCoordinator:
152
break
153
else:
154
attention_groups.append(
146
- SpecGroup(spec, [i], manager_cls, g.is_eagle_group)
155
+ StoreSpecGroup(spec, [i], manager_cls, g.is_eagle_group)
156
)
157
# Full attention first (matches upstream convergence ordering).
158
attention_groups.sort(key=lambda g: not isinstance(g.spec, FullAttentionSpec))
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+134
-17
@@ -67,6 +67,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import
67
from vllm.logger import init_logger
68
from vllm.utils.math_utils import cdiv
69
from vllm.utils.network_utils import get_ip, make_zmq_socket
70
+from vllm.utils.torch_utils import is_non_overlapping_and_dense
71
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
72
from vllm.v1.core.kv_cache_utils import (
73
BlockHash,
@@ -1909,21 +1910,31 @@ class MooncakeStoreWorker:
1910
logger.warning("No KV caches to offload.")
1911
return
1912
1913
+ # Resolve each entry to a representative tensor for storage
1914
+ # deduplication. For attention layers the value is already a tensor;
1915
+ # for Mamba layers it is a list of tensors that all share the same
1916
+ # underlying raw storage, so we take the first one.
1917
+ def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
1918
+ assert isinstance(v, torch.Tensor | list)
1919
+ return v if isinstance(v, torch.Tensor) else v[0]
1920
+
1921
assert self.cache_config.num_gpu_blocks is not None
1922
self.num_blocks = self.cache_config.num_gpu_blocks
1923
+ use_group_regions = self._kv_cache_config.hisparse_host_num_blocks is not None
1924
1915
- seen_storage_ptrs: set[int] = set()
1916
- cache_tensors: list[torch.Tensor] = []
1917
-
1918
- for cache in kv_caches.values():
1919
- cache = group_kernel_blocks(cache, self.num_blocks)
1920
- cache_tensors.append(cache)
1921
- cache_storage = cache.untyped_storage()
1922
- base_addr = cache_storage.data_ptr()
1923
- region_len = cache_storage.nbytes()
1924
-
1925
- if base_addr not in seen_storage_ptrs:
1925
+ if not use_group_regions:
1926
+ seen_storage_ptrs: set[int] = set()
1927
+ cache_tensors = [
1928
+ group_kernel_blocks(_repr_tensor(cache), self.num_blocks)
1929
+ for cache in kv_caches.values()
1930
+ ]
1931
+ for cache in cache_tensors:
1932
+ cache_storage = cache.untyped_storage()
1933
+ base_addr = cache_storage.data_ptr()
1934
+ if base_addr in seen_storage_ptrs:
1935
+ continue
1936
seen_storage_ptrs.add(base_addr)
1937
+ region_len = cache_storage.nbytes()
1938
ret = self.store.register_buffer(base_addr, region_len)
1939
if ret != 0:
1940
logger.error(
@@ -1932,17 +1943,110 @@ class MooncakeStoreWorker:
1943
region_len,
1944
ret,
1945
)
1946
+ for db in self.token_dbs:
1947
+ db.store_layout.register_kv_caches(cache_tensors, self.num_blocks)
1948
+
1949
+ registered_buffers: dict[int, int] = {}
1950
+ num_host_segments = 0
1951
+ num_device_segments = 0
1952
+
1953
+ for group, db in zip(self._kv_cache_groups, self.token_dbs, strict=True):
1954
+ if not use_group_regions:
1955
+ continue
1956
+ missing_layers = [
1957
+ layer_name
1958
+ for layer_name in group.layer_names
1959
+ if layer_name not in kv_caches
1960
+ ]
1961
+ if missing_layers:
1962
+ raise ValueError(
1963
+ f"Missing transferable KV caches for layers: {missing_layers}"
1964
+ )
1965
+ group_caches = [
1966
+ (layer_name, kv_caches[layer_name]) for layer_name in group.layer_names
1967
+ ]
1968
+
1969
+ seen_group_ptrs: set[int] = set()
1970
+ addrs: list[int] = []
1971
+ block_lens: list[int] = []
1972
+ for layer_name, value in group_caches:
1973
+ cache = _repr_tensor(value)
1974
+ is_host_resident = group.host_resident
1975
+ if is_host_resident:
1976
+ num_blocks = self._kv_cache_config.hisparse_host_num_blocks
1977
+ assert num_blocks is not None
1978
+ else:
1979
+ num_blocks = self._kv_cache_config.num_blocks
1980
+ cache = group_kernel_blocks(cache, num_blocks)
1981
+ cache_storage = cache.untyped_storage()
1982
+ if is_host_resident:
1983
+ base_addr = cache.data_ptr()
1984
+ region_len = cache.nbytes
1985
+ else:
1986
+ base_addr = cache_storage.data_ptr()
1987
+ region_len = cache_storage.nbytes()
1988
+ if region_len % num_blocks != 0:
1989
+ raise ValueError(
1990
+ f"KV cache region for {layer_name} has size {region_len}, "
1991
+ f"which is not divisible by its {num_blocks} blocks."
1992
+ )
1993
+
1994
+ registered_len = registered_buffers.get(base_addr)
1995
+ if registered_len is None:
1996
+ ret = self.store.register_buffer(base_addr, region_len)
1997
+ if ret != 0:
1998
+ mem_kind = "host" if is_host_resident else "device"
1999
+ raise RuntimeError(
2000
+ "Mooncake register_buffer failed for addr "
2001
+ f"{base_addr:#x} len {region_len} ({mem_kind}): {ret}"
2002
+ )
2003
+ registered_buffers[base_addr] = region_len
2004
+ elif registered_len != region_len:
2005
+ raise ValueError(
2006
+ f"KV cache views at {base_addr:#x} expose inconsistent "
2007
+ f"region sizes: {registered_len} and {region_len}."
2008
+ )
2009
+
2010
+ regions: list[tuple[int, int]] = []
2011
+ if not is_non_overlapping_and_dense(cache[0]):
2012
+ for head_idx in range(cache.shape[1]):
2013
+ head_cache = cache[:, head_idx]
2014
+ assert is_non_overlapping_and_dense(head_cache[0])
2015
+ regions.append(
2016
+ (
2017
+ head_cache.data_ptr(),
2018
+ head_cache.stride(0) * head_cache.element_size(),
2019
+ )
2020
+ )
2021
+ elif cache.stride(0) * cache.element_size() * num_blocks == region_len:
2022
+ regions.append((base_addr, region_len // num_blocks))
2023
+ else:
2024
+ regions.append(
2025
+ (cache.data_ptr(), cache.stride(0) * cache.element_size())
2026
+ )
2027
+
2028
+ for region_addr, block_len in regions:
2029
+ if region_addr in seen_group_ptrs:
2030
+ continue
2031
+ seen_group_ptrs.add(region_addr)
2032
+ addrs.append(region_addr)
2033
+ block_lens.append(block_len)
2034
+ if is_host_resident:
2035
+ num_host_segments += 1
2036
+ else:
2037
+ num_device_segments += 1
2038
+
2039
+ db.set_kv_caches_base_addr(addrs)
2040
+ db.set_block_len(block_lens)
2041
2042
logger.info(
1937
- "Registered KV caches: num_groups=%d, num_tensors=%d, num_blocks=%d",
2043
+ "Registered KV caches: num_groups=%d, num_segments=%d (host=%d, device=%d)",
2044
len(self.token_dbs),
1939
- len(cache_tensors),
1940
- self.num_blocks,
2045
+ num_host_segments + num_device_segments,
2046
+ num_host_segments,
2047
+ num_device_segments,
2048
)
2049
1943
- for db in self.token_dbs:
1944
- db.store_layout.register_kv_caches(cache_tensors, self.num_blocks)
1945
-
2050
# Start transfer threads
2051
if self.can_put:
2052
ready_event_sending = threading.Event()
@@ -2316,6 +2420,19 @@ class MooncakeStoreWorker:
2420
if store is None:
2421
return
2422
self.store = None
2423
+ deadline = time.monotonic() + 5.0
2424
+ for thread in (self.kv_send_thread, *self.kv_recv_threads):
2425
+ if thread is None:
2426
+ continue
2427
+ while thread.request_queue.unfinished_tasks and time.monotonic() < deadline:
2428
+ time.sleep(0.05)
2429
+ if thread.request_queue.unfinished_tasks:
2430
+ logger.warning(
2431
+ "Mooncake store %s still has %d in-flight requests at "
2432
+ "close; closing anyway.",
2433
+ thread.name,
2434
+ thread.request_queue.unfinished_tasks,
2435
+ )
2436
try:
2437
store.close()
2438
except Exception as e:
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+44
-22
@@ -17,6 +17,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
17
KVConnectorHandshakeMetadata,
18
KVConnectorMetadata,
19
KVConnectorRole,
20
+ KVConnectorTransferResults,
21
KVConnectorWorkerMetadata,
22
SupportsHMA,
23
)
@@ -35,7 +36,7 @@ if TYPE_CHECKING:
36
from vllm.distributed.kv_events import KVCacheEvent
37
from vllm.forward_context import ForwardContext
38
from vllm.v1.core.block_pool import BlockPool
38
- from vllm.v1.core.kv_cache_manager import KVCacheBlocks
39
+ from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
40
from vllm.v1.kv_cache_interface import KVCacheConfig
41
from vllm.v1.request import Request
42
@@ -208,6 +209,10 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
209
# Propagated from scheduler to worker side via the connector metadata.
210
self._extra_async_saves: dict[str, int] = {}
211
212
+ @property
213
+ def sub_connectors(self) -> list[KVConnectorBase_V1]:
214
+ return list(self._connectors)
215
+
216
@property
217
def supports_divergent_local_hybrid_hits(self) -> bool:
218
return bool(self._connectors) and all(
@@ -248,6 +253,10 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
253
for c in self._connectors:
254
c.register_kv_caches(kv_caches)
255
256
+ def bind_kv_cache_manager(self, kv_cache_manager: "KVCacheManager") -> None:
257
+ for connector in self._connectors:
258
+ connector.bind_kv_cache_manager(kv_cache_manager)
259
+
260
def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None:
261
for c in self._connectors:
262
c.bind_gpu_block_pool(gpu_block_pool)
@@ -292,6 +301,14 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
301
for c in self._connectors:
302
c.start_load_kv(forward_context, **kwargs)
303
304
+ def finish_forward(self) -> None:
305
+ for c in self._connectors:
306
+ c.finish_forward()
307
+
308
+ def reset_capture_state(self) -> None:
309
+ for c in self._connectors:
310
+ c.reset_capture_state()
311
+
312
def wait_for_layer_load(self, layer_name: str) -> None:
313
for c in self._connectors:
314
c.wait_for_layer_load(layer_name)
@@ -313,29 +330,31 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
330
def get_finished(
331
self, finished_req_ids: set[str]
332
) -> tuple[set[str] | None, set[str] | None]:
316
- finished_sending: set[str] = set()
317
- finished_recving: set[str] = set()
318
- for c in self._connectors:
319
- sending, recving = c.get_finished(finished_req_ids)
320
- if not recving and not sending:
321
- continue
322
- # Aggregate finished recving request ids.
323
- finished_recving.update(recving or ())
324
- # Aggregate finished sending request ids - only include
325
- # once we've drained the "extra" count (for cases where
326
- # more than one connector is async-saving the same request).
327
- for req_id in sending or ():
333
+ results = self.get_transfer_results(finished_req_ids)
334
+ return (
335
+ results.finished_sending or None,
336
+ results.finished_recving or None,
337
+ )
338
+
339
+ def get_transfer_results(
340
+ self, finished_req_ids: set[str]
341
+ ) -> KVConnectorTransferResults:
342
+ results = KVConnectorTransferResults()
343
+ for connector in self._connectors:
344
+ child_results = connector.get_transfer_results(finished_req_ids)
345
+ results.finished_recving.update(child_results.finished_recving)
346
+ results.failed_recving.update(child_results.failed_recving)
347
+ for req_id in child_results.finished_sending:
348
extra_pending = self._extra_async_saves.get(req_id)
349
if extra_pending is None:
330
- finished_sending.add(req_id)
331
- continue
332
- assert extra_pending > 0
333
- if extra_pending == 1:
334
- del self._extra_async_saves[req_id]
350
+ results.finished_sending.add(req_id)
351
else:
336
- self._extra_async_saves[req_id] = extra_pending - 1
337
-
338
- return finished_sending or None, finished_recving or None
352
+ assert extra_pending > 0
353
+ if extra_pending == 1:
354
+ del self._extra_async_saves[req_id]
355
+ else:
356
+ self._extra_async_saves[req_id] = extra_pending - 1
357
+ return results
358
359
def get_block_ids_with_load_errors(self) -> set[int]:
360
agg_block_ids: set[int] = set()
@@ -563,6 +582,9 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
582
for c in self._connectors:
583
yield from c.take_events()
584
585
+ def has_pending_block_frees(self) -> bool:
586
+ return any(c.has_pending_block_frees() for c in self._connectors)
587
+
588
def has_pending_push_work(self) -> bool:
589
return any(c.has_pending_push_work() for c in self._connectors)
590
@@ -630,7 +652,7 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
652
653
return MultiKVConnectorStats(data=reconstructed_data)
654
633
- def get_kv_connector_stats(self) -> MultiKVConnectorStats | None:
655
+ def get_kv_connector_stats(self) -> KVConnectorStats | None:
656
# Group connector stats by connector type.
657
stats_by_connector: MultiKVConnectorStats | None = None
658
for c in self._connectors:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
+1
-4
@@ -93,10 +93,7 @@ class NixlBaseConnectorScheduler:
93
for g in kv_cache_config.transfer_groups
94
)
95
)
96
- self._has_mamba = any(
97
- isinstance(g.kv_cache_spec, MambaSpec)
98
- for g in kv_cache_config.transfer_groups
99
- )
96
+ self._has_mamba = kv_cache_config.has_mamba_layers
97
98
logger.info("Initializing NIXL Scheduler %s", engine_id)
99
if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+217
-135
@@ -2,8 +2,10 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Base worker-side logic for the NIXL connector."""
4
5
+import contextlib
6
import itertools
7
import logging
8
+import math
9
import os
10
import queue
11
import threading
@@ -31,7 +33,10 @@ from vllm.distributed.kv_transfer.kv_connector.utils import (
33
kv_postprocess_blksize_on_receive,
34
kv_postprocess_layout_on_receive,
35
)
34
-from vllm.distributed.kv_transfer.kv_connector.v1.base import CopyBlocksOp
36
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
37
+ CopyBlocksOp,
38
+ KVConnectorTransferResults,
39
+)
40
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
41
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
42
GET_META_MSG,
@@ -186,23 +191,28 @@ class NixlBaseConnectorWorker:
191
num_blocks = dst_num_blocks
192
if block_size_ratio is not None:
193
num_blocks = int(num_blocks * block_size_ratio)
194
+ num_regions = self.num_regions or len(self.block_len_per_layer)
195
if region_num_blocks is None:
190
- region_num_blocks = [num_blocks] * self.num_regions
196
+ region_num_blocks = [num_blocks] * num_regions
197
elif block_size_ratio is not None:
198
region_num_blocks = [
199
int(count * block_size_ratio) for count in region_num_blocks
200
]
195
- assert len(region_num_blocks) == self.num_regions
201
+ if region_num_blocks:
202
+ num_regions = len(region_num_blocks)
203
+ assert len(region_num_blocks) == num_regions
204
region_offsets = np.cumsum([0, *region_num_blocks[:-1]])
205
num_fa_descs = sum(region_num_blocks)
206
199
- # All-attention fast path: single vectorized broadcast.
207
+ # All-attention fast path.
208
if num_ssm_regions == 0:
209
if region_group_ids is None:
210
region_group_ids = self.region_group_ids
211
if uses_region_group_mapping is None:
212
uses_region_group_mapping = self._uses_region_group_mapping
205
- assert len(region_group_ids) == self.num_regions
213
+ if not region_group_ids and num_regions == 1:
214
+ region_group_ids = [0]
215
+ assert len(region_group_ids) == num_regions
216
if uses_region_group_mapping is None:
217
uses_region_group_mapping = len(set(region_group_ids)) > 1
218
region_group_ids_array = np.asarray(region_group_ids, dtype=np.int32)
@@ -397,8 +407,6 @@ class NixlBaseConnectorWorker:
407
408
# Config.
409
self.vllm_config = vllm_config
400
- # mypy will complain on re-assignment otherwise.
401
- self.block_size: int = cast(int, vllm_config.cache_config.block_size)
410
411
if vllm_config.kv_transfer_config is None:
412
raise ValueError("kv_transfer_config must be set for NixlConnector")
@@ -420,6 +428,17 @@ class NixlBaseConnectorWorker:
428
)
429
430
self.kv_cache_config = kv_cache_config
431
+ transfer_block_sizes = [
432
+ group.kv_cache_spec.block_size
433
+ for group in kv_cache_config.transfer_groups
434
+ if get_representative_spec_type(group.kv_cache_spec)
435
+ not in (MambaSpec, CircularBufferSpec)
436
+ ]
437
+ self.block_size = (
438
+ math.lcm(*transfer_block_sizes)
439
+ if transfer_block_sizes
440
+ else cast(int, vllm_config.cache_config.block_size)
441
+ )
442
# Per-layer specs, unwrapping UniformTypeKVCacheSpecs group wrappers.
443
self._layer_specs: dict[str, KVCacheSpec] = {}
444
for group in kv_cache_config.transfer_groups:
@@ -599,12 +618,15 @@ class NixlBaseConnectorWorker:
618
# Number of NIXL regions. Currently one region per cache
619
# (so 1 per layer for MLA, otherwise 2 per layer)
620
self.num_regions = 0
621
+ self.region_mem_types: list[str] = []
622
self.region_group_ids: list[int] = []
623
self._uses_region_group_mapping = False
624
self.region_names: list[str] = []
625
self.region_num_blocks: list[int] = []
606
- self.region_mem_types: list[str] = []
626
self._mixed_mem_types = False
627
+ self._desc_is_dram_by_block_size: dict[int, np.ndarray] = {}
628
+ self._desc_pos_by_block_size: dict[int, np.ndarray] = {}
629
+ self._dram_src_handles_by_block_size: dict[int, int] = {}
630
631
# PP>1 (push mode): this worker holds a contiguous layer slice and
632
# transfers into the matching sub-range of a PP=1 remote's regions.
@@ -644,9 +666,6 @@ class NixlBaseConnectorWorker:
666
self.dst_region_group_ids: dict[EngineId, list[int]] = {}
667
self.dst_uses_region_group_mapping: dict[EngineId, bool] = {}
668
self.dst_region_mem_types: dict[EngineId, list[str]] = {}
647
- self._desc_is_dram_by_block_size: dict[int, np.ndarray] = {}
648
- self._desc_pos_by_block_size: dict[int, np.ndarray] = {}
649
- self._dram_src_handles_by_block_size: dict[int, int] = {}
669
self._registered_descs: list[Any] = []
670
671
# In progress transfers.
@@ -667,6 +686,7 @@ class NixlBaseConnectorWorker:
686
# Uses Queue for thread-safe cross-thread coordination with the
687
# background handshake thread, matching the _ready_requests pattern.
688
self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue()
689
+ self._failed_inflight_recvs: set[ReqId] = set()
690
self._pending_recv_notifs: dict[ReqId, list[tuple[str, bytes]]] = {}
691
692
# Handshake metadata of this worker for NIXL transfers.
@@ -690,7 +710,6 @@ class NixlBaseConnectorWorker:
710
"engine_ttl", 3600.0
711
)
712
693
- self.block_size = vllm_config.cache_config.block_size
713
self.model_config = vllm_config.model_config
714
715
self.use_mla = self.model_config.use_mla
@@ -1240,7 +1259,7 @@ class NixlBaseConnectorWorker:
1259
self.use_host_buffer,
1260
)
1261
1243
- registration_ranges: dict[tuple[int, str], tuple[int, int, int, str]] = {}
1262
+ registration_ranges: dict[tuple[int, str], tuple[int, int, int]] = {}
1263
region_mem_types: list[str] = []
1264
seen_base_addresses: list[int] = []
1265
self._ssm_region_indices = []
@@ -1280,10 +1299,16 @@ class NixlBaseConnectorWorker:
1299
# K and V are packed into the content dim, so each attention layer is a
1300
# single NIXL region whose block transfers as one unit. Mamba layers instead
1301
# register separate conv/ssm sub-regions (see `_build_mamba_local`).
1302
+ # P and D may allocate equivalent transferable layers in different
1303
+ # cache-group orders. Keep their region lists aligned without putting
1304
+ # layers.10 before layers.2, which would break PP region slicing.
1305
layer_names = (
1306
xfer_buffers
1307
if self._is_csa_linear
1286
- else sorted(xfer_buffers, key=_region_sort_key)
1308
+ else sorted(
1309
+ xfer_buffers,
1310
+ key=_region_sort_key,
1311
+ )
1312
)
1313
for layer_name in layer_names:
1314
cache = xfer_buffers[layer_name]
@@ -1299,6 +1324,7 @@ class NixlBaseConnectorWorker:
1324
layer_name,
1325
)
1326
continue
1327
+ group_index = self.kv_cache_config.transfer_group_index_by_layer[layer_name]
1328
# `layer_spec.page_size_bytes` only accounts for logical page_size, that is
1329
# the page_size assuming constant `self._logical_num_blocks`.
1330
physical_page_size = (
@@ -1307,17 +1333,48 @@ class NixlBaseConnectorWorker:
1333
else layer_spec.page_size_bytes
1334
// self._physical_blocks_per_logical_kv_block
1335
)
1310
- group_id = self.kv_cache_config.transfer_group_index_by_layer[layer_name]
1336
+ group = self.kv_cache_config.transfer_groups[group_index]
1337
+ if group.host_resident:
1338
+ logical_num_blocks = self.kv_cache_config.hisparse_host_num_blocks
1339
+ assert logical_num_blocks is not None
1340
+ else:
1341
+ logical_num_blocks = self.kv_cache_config.num_blocks
1342
+ group_id = group_index
1343
num_blocks = (
1312
- self._logical_num_blocks
1344
+ logical_num_blocks
1345
if isinstance(layer_spec, MambaSpec)
1314
- else self.num_blocks
1346
+ else logical_num_blocks * self._physical_blocks_per_logical_kv_block
1347
+ )
1348
+ base_addr = cache.data_ptr()
1349
+ is_mla_region = isinstance(
1350
+ layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
1351
)
1352
logger.debug(
1353
"Registering layer %s with cache shape: %s", layer_name, cache.shape
1354
)
1355
storage = cache.untyped_storage()
1356
storage_addr = storage.data_ptr()
1357
+ is_host_resident = group.host_resident
1358
+ if cache.device.type == "cpu":
1359
+ mem_type = "DRAM"
1360
+ region_device_id = 0
1361
+ else:
1362
+ mem_type = self.nixl_memory_type
1363
+ region_device_id = max(cache.get_device(), 0)
1364
+ self.device_id = region_device_id
1365
+ if is_host_resident:
1366
+ registration_base = cache.data_ptr()
1367
+ registration_len = cache.nbytes
1368
+ else:
1369
+ registration_base = storage_addr
1370
+ registration_len = storage.nbytes()
1371
+ storage_key = (registration_base, mem_type)
1372
+ if storage_key not in registration_ranges:
1373
+ registration_ranges[storage_key] = (
1374
+ registration_base,
1375
+ registration_base + registration_len,
1376
+ region_device_id,
1377
+ )
1378
1379
if isinstance(layer_spec, KpoolTailSpec):
1380
compressed_owner = compressed_region_owners.get(cache.data_ptr())
@@ -1345,22 +1402,6 @@ class NixlBaseConnectorWorker:
1402
)
1403
continue
1404
1348
- if cache.device.type == "cpu":
1349
- mem_type = "DRAM"
1350
- region_device_id = 0
1351
- else:
1352
- mem_type = self.nixl_memory_type
1353
- region_device_id = max(cache.get_device(), 0)
1354
- self.device_id = region_device_id
1355
- registration_ranges.setdefault(
1356
- (storage_addr, mem_type),
1357
- (storage_addr, storage.nbytes(), region_device_id, ""),
1358
- )
1359
-
1360
- is_mla_region = isinstance(
1361
- layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
1362
- )
1363
-
1405
if isinstance(layer_spec, MambaSpec):
1406
physical_ratio = self._physical_blocks_per_logical_kv_block
1407
block_len = physical_page_size // physical_ratio
@@ -1381,14 +1422,13 @@ class NixlBaseConnectorWorker:
1422
block_stride = physical_page_size
1423
else:
1424
block_stride = cache.stride(0) * cache.element_size()
1384
- storage_is_block_major = num_blocks * block_stride == storage.nbytes()
1385
- # A layer whose [H, N, C] interior is dense addresses its own page
1386
- # as one chunk. Otherwise the block's whole row is the only
1387
- # contiguous transfer unit.
1388
- hnc_contiguous = (
1389
- cache.ndim == 4
1390
- and cache.stride(2) == cache.shape[3]
1391
- and cache.stride(1) == cache.shape[2] * cache.shape[3]
1425
+ storage_is_block_major = num_blocks * block_stride == registration_len
1426
+ # A view with a contiguous page can own a strided region even
1427
+ # when sibling layers share the same backing allocation.
1428
+ page_contiguous = (
1429
+ cache.ndim > 1
1430
+ and cache[0].is_contiguous()
1431
+ and cache[0].nbytes == physical_page_size
1432
)
1433
virtual_transfer_pages = _uses_dense_virtual_transfer_pages(
1434
layer_spec, cache, physical_page_size, num_blocks
@@ -1399,13 +1439,14 @@ class NixlBaseConnectorWorker:
1439
(cache.data_ptr(), physical_page_size, physical_page_size)
1440
]
1441
elif storage_is_block_major and (
1402
- (packed_storage and is_mla_region)
1403
- or (not hnc_contiguous and not self._is_csa_linear)
1442
+ not page_contiguous
1443
+ and ((packed_storage and is_mla_region) or not self._is_csa_linear)
1444
):
1405
- # Packed MLA layouts transfer the complete storage row.
1406
- storage_block_len = storage.nbytes() // num_blocks
1445
+ # TODO(Lucas): handle TP slicing for packed_storage; for now
1446
+ # restrict to MLA (DSv4) where kv is replicated.
1447
+ storage_block_len = registration_len // num_blocks
1448
region_specs = [
1408
- (storage_addr, storage_block_len, storage_block_len)
1449
+ (registration_base, storage_block_len, storage_block_len)
1450
]
1451
elif storage_is_block_major:
1452
region_specs = [
@@ -1440,6 +1481,10 @@ class NixlBaseConnectorWorker:
1481
region_index = seen_base_addresses.index(base_addr)
1482
assert region_mem_types[region_index] == mem_type
1483
self._region_is_mla[region_index] |= is_mla_region
1484
+ if is_mla_region:
1485
+ self.block_len_per_layer[region_index] = block_len
1486
+ self.block_stride_per_layer[region_index] = block_stride
1487
+ self.region_num_blocks[region_index] = num_blocks
1488
if self.region_group_ids[region_index] != group_id:
1489
self.region_group_ids[region_index] = _SHARED_REGION_GROUP_ID
1490
else:
@@ -1450,8 +1495,8 @@ class NixlBaseConnectorWorker:
1495
self.region_group_ids.append(group_id)
1496
self.region_names.append(layer_name)
1497
self.region_num_blocks.append(num_blocks)
1453
- region_mem_types.append(mem_type)
1498
self._region_is_mla.append(is_mla_region)
1499
+ region_mem_types.append(mem_type)
1500
1501
if isinstance(layer_spec, MambaSpec):
1502
if layer_spec.tp_replicated:
@@ -1496,7 +1541,6 @@ class NixlBaseConnectorWorker:
1541
== len(self.region_group_ids)
1542
== len(self.region_names)
1543
== len(self.region_num_blocks)
1499
- == len(region_mem_types)
1544
)
1545
# Descriptor ids must be region-ordered, matching the remote side.
1546
self._scratch_region_indices.sort()
@@ -1521,16 +1565,23 @@ class NixlBaseConnectorWorker:
1565
# Total local FA descriptors (boundary between FA and mamba descs).
1566
self.num_descs = sum(self.region_num_blocks)
1567
1568
+ self._mixed_mem_types = len(set(region_mem_types)) > 1
1569
+ if self._mixed_mem_types:
1570
+ assert self.use_mla and not self._has_mamba, (
1571
+ "Mixed-device KV registration is only supported for MLA "
1572
+ "models without Mamba layers."
1573
+ )
1574
for mem_type in sorted(set(region_mem_types)):
1525
- ranges = [
1526
- registration
1527
- for (
1528
- _,
1529
- registration_mem_type,
1530
- ), registration in registration_ranges.items()
1531
- if registration_mem_type == mem_type
1575
+ ranges_for_mem_type = [
1576
+ (start, end - start, device_id, "")
1577
+ for (_, cache_mem_type), (
1578
+ start,
1579
+ end,
1580
+ device_id,
1581
+ ) in registration_ranges.items()
1582
+ if cache_mem_type == mem_type
1583
]
1533
- descs = self.nixl_wrapper.get_reg_descs(ranges, mem_type)
1584
+ descs = self.nixl_wrapper.get_reg_descs(ranges_for_mem_type, mem_type)
1585
self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
1586
self._registered_descs.append(descs)
1587
@@ -1854,7 +1905,11 @@ class NixlBaseConnectorWorker:
1905
if self._mixed_mem_types:
1906
desc_is_dram = np.concatenate(
1907
[
1857
- np.full(count * block_size_ratio, mem_type == "DRAM", dtype=bool)
1908
+ np.full(
1909
+ count * block_size_ratio,
1910
+ mem_type == "DRAM",
1911
+ dtype=bool,
1912
+ )
1913
for mem_type, count in zip(
1914
self.region_mem_types,
1915
self.region_num_blocks,
@@ -1864,24 +1919,29 @@ class NixlBaseConnectorWorker:
1919
)
1920
assert len(desc_is_dram) == len(blocks_data)
1921
desc_pos = np.empty(len(desc_is_dram), dtype=np.int64)
1867
- dram_indices = np.flatnonzero(desc_is_dram)
1868
- device_indices = np.flatnonzero(~desc_is_dram)
1869
- desc_pos[dram_indices] = np.arange(len(dram_indices), dtype=np.int64)
1870
- desc_pos[device_indices] = np.arange(len(device_indices), dtype=np.int64)
1922
+ dram_idx = np.where(desc_is_dram)[0]
1923
+ vram_idx = np.where(~desc_is_dram)[0]
1924
+ desc_pos[dram_idx] = np.arange(len(dram_idx), dtype=np.int64)
1925
+ desc_pos[vram_idx] = np.arange(len(vram_idx), dtype=np.int64)
1926
self._desc_is_dram_by_block_size[block_size] = desc_is_dram
1927
self._desc_pos_by_block_size[block_size] = desc_pos
1873
- blocks_data[desc_is_dram, 2] = 0
1874
- dram_descs = self.nixl_wrapper.get_xfer_descs(
1875
- blocks_data[desc_is_dram], "DRAM"
1876
- )
1928
+
1929
+ # DRAM descriptors are registered under CPU device 0.
1930
+ blocks_data = [
1931
+ (addr, length, 0) if is_dram else (addr, length, dev)
1932
+ for (addr, length, dev), is_dram in zip(
1933
+ blocks_data, desc_is_dram, strict=True
1934
+ )
1935
+ ]
1936
+ dram_blocks = [blocks_data[i] for i in dram_idx]
1937
+ vram_blocks = [blocks_data[i] for i in vram_idx]
1938
+ dram_descs = self.nixl_wrapper.get_xfer_descs(dram_blocks, "DRAM")
1939
self._dram_src_handles_by_block_size[block_size] = (
1940
self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", dram_descs)
1941
)
1880
- device_descs = self.nixl_wrapper.get_xfer_descs(
1881
- blocks_data[~desc_is_dram], self.nixl_memory_type
1882
- )
1942
+ descs = self.nixl_wrapper.get_xfer_descs(vram_blocks, self.nixl_memory_type)
1943
return (
1884
- self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", device_descs),
1944
+ self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs),
1945
blocks_data,
1946
)
1947
@@ -1968,10 +2028,6 @@ class NixlBaseConnectorWorker:
2028
]
2029
nixl_agent_meta.block_lens = nixl_agent_meta.block_lens[start:end]
2030
nixl_agent_meta.block_strides = nixl_agent_meta.block_strides[start:end]
1971
- if nixl_agent_meta.region_num_blocks is not None:
1972
- nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[
1973
- start:end
1974
- ]
2031
if nixl_agent_meta.region_group_ids is not None:
2032
nixl_agent_meta.region_group_ids = nixl_agent_meta.region_group_ids[
2033
start:end
@@ -1982,6 +2038,10 @@ class NixlBaseConnectorWorker:
2038
nixl_agent_meta.region_mem_types = nixl_agent_meta.region_mem_types[
2039
start:end
2040
]
2041
+ if nixl_agent_meta.region_num_blocks is not None:
2042
+ nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[
2043
+ start:end
2044
+ ]
2045
2046
### Register remote engine in TransferTopology (idempotent).
2047
assert self.transfer_topo is not None
@@ -2026,11 +2086,12 @@ class NixlBaseConnectorWorker:
2086
nixl_agent_meta.region_num_blocks
2087
or [nixl_agent_meta.num_blocks] * num_remote_regions
2088
)
2029
- self.dst_region_group_ids[engine_id] = nixl_agent_meta.region_group_ids or (
2089
+ remote_region_group_ids = nixl_agent_meta.region_group_ids or (
2090
self.region_group_ids
2091
if len(self.region_group_ids) == num_remote_regions
2092
else [0] * num_remote_regions
2093
)
2094
+ self.dst_region_group_ids[engine_id] = remote_region_group_ids
2095
self.dst_uses_region_group_mapping[engine_id] = (
2096
len(set(self.dst_region_group_ids[engine_id])) > 1
2097
)
@@ -2138,8 +2199,6 @@ class NixlBaseConnectorWorker:
2199
mamba = self._build_mamba_remote(nixl_agent_meta, tp_ratio, transfer_info)
2200
blocks_data = np.concatenate([blocks_data, mamba])
2201
2141
- # A mixed local destination is split into multiple READs below. The
2142
- # producer remains a single descriptor list in this prerequisite.
2202
remote_mem_types = set(self.dst_region_mem_types[engine_id])
2203
if len(remote_mem_types) != 1:
2204
raise NotImplementedError(
@@ -2540,9 +2599,10 @@ class NixlBaseConnectorWorker:
2599
indices=indices,
2600
)
2601
2543
- def get_finished(self) -> tuple[set[str], set[str]]:
2602
+ def get_transfer_results(self) -> KVConnectorTransferResults:
2603
"""
2545
- Get requests that are done sending or recving on this specific worker.
2604
+ Get transfers that completed on this specific worker.
2605
+
2606
The scheduler process (via the MultiprocExecutor) will use this output
2607
to track which workers are done.
2608
"""
@@ -2564,7 +2624,6 @@ class NixlBaseConnectorWorker:
2624
# Add failed requests to done_recving for scheduler tracking
2625
# (blocks are already marked invalid, scheduler will handle recompute)
2626
done_recving.update(failed_recv_reqs)
2567
-
2627
if len(done_sending) > 0 or len(done_recving) > 0:
2628
logger.debug(
2629
"Rank %s, get_finished: %s requests done sending "
@@ -2577,6 +2636,7 @@ class NixlBaseConnectorWorker:
2636
2637
block_ids_for_blocksize_post_process = defaultdict(list)
2638
block_ids_for_heterogeneous_attn_post_process = list[list[int]]()
2639
+ direct_device_recving = set[str]()
2640
for req_id in done_recving:
2641
# clean up metadata for completed requests
2642
meta = self._recving_metadata.pop(req_id, None)
@@ -2594,6 +2654,8 @@ class NixlBaseConnectorWorker:
2654
if self.use_host_buffer:
2655
self.sync_recved_kv_to_device(req_id, meta)
2656
2657
+ direct_device_recving.add(req_id)
2658
+
2659
# Post processing for heteroblocksize/layout, and for blocks the
2660
# transfer clipped. The latter happens either at remote-block
2661
# granularity (block_size_ratio > 1) or at kernel-block
@@ -2642,7 +2704,7 @@ class NixlBaseConnectorWorker:
2704
for block_ids in block_ids_for_heterogeneous_attn_post_process:
2705
self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids)
2706
2645
- self._sync_device_after_mamba_recv(done_recving, failed_recv_reqs)
2707
+ self._sync_device_after_direct_recv(direct_device_recving)
2708
2709
# Handle timeout to avoid stranding blocks on remote.
2710
now = time.perf_counter()
@@ -2664,20 +2726,21 @@ class NixlBaseConnectorWorker:
2726
del self._reqs_to_send[req_id]
2727
done_sending.add(req_id)
2728
2667
- return done_sending, done_recving
2729
+ return KVConnectorTransferResults(
2730
+ finished_sending=done_sending,
2731
+ finished_recving=done_recving,
2732
+ failed_recving=failed_recv_reqs,
2733
+ )
2734
2669
- def _sync_device_after_mamba_recv(
2670
- self,
2671
- done_recving: set[str],
2672
- failed_recv_reqs: set[str],
2673
- ) -> None:
2674
- """Synchronize ROCm direct-GPU Mamba receives before model execution."""
2675
- if (
2676
- not current_platform.is_rocm()
2677
- or not self._has_mamba
2678
- or self.use_host_buffer
2679
- or not (done_recving - failed_recv_reqs)
2680
- ):
2735
+ def get_finished(self) -> tuple[set[str], set[str]]:
2736
+ """Compatibility wrapper for the legacy completion API."""
2737
+ results = self.get_transfer_results()
2738
+ return results.finished_sending, results.finished_recving
2739
+
2740
+ def _sync_device_after_direct_recv(self, done_recving: set[str]) -> None:
2741
+ """Make direct NIXL writes visible before model execution."""
2742
+ requires_sync = current_platform.is_rocm() and self._has_mamba
2743
+ if self.use_host_buffer or not done_recving or not requires_sync:
2744
return
2745
2746
torch.accelerator.synchronize()
@@ -2731,11 +2794,11 @@ class NixlBaseConnectorWorker:
2794
self.nixl_wrapper.release_xfer_handle(handle)
2795
elif xfer_state == "PROC":
2796
in_progress.append(handle)
2734
- continue
2797
else:
2798
self._log_failure(
2799
failure_type="transfer_failed",
2738
- msg="Marking blocks as invalid",
2800
+ msg="Deferring request completion until its last "
2801
+ "xfer is terminal",
2802
req_id=req_id,
2803
xfer_state=xfer_state,
2804
)
@@ -2743,21 +2806,20 @@ class NixlBaseConnectorWorker:
2806
except Exception as e:
2807
self._log_failure(
2808
failure_type="transfer_exception",
2746
- msg="Marking blocks as invalid",
2809
+ msg="Handle is unpollable; treating it as terminal",
2810
req_id=req_id,
2811
error=e,
2812
)
2813
self._handle_failed_transfer(req_id, handle)
2814
2815
if not in_progress:
2753
- # Only report request as completed when all transfers are done.
2754
- # A request failed in an earlier poll was already reported via
2755
- # _failed_recv_reqs and its metadata popped by get_finished();
2756
- # don't report it again, just drop the remaining handles.
2757
- if req_id in self._recving_metadata:
2758
- done_req_ids.add(req_id)
2759
- self._send_pending_recv_notifs(req_id)
2816
del transfers[req_id]
2817
+ done_req_ids.add(req_id)
2818
+ if req_id in self._failed_inflight_recvs:
2819
+ self._failed_inflight_recvs.remove(req_id)
2820
+ self._report_failed_recv(req_id)
2821
+ else:
2822
+ self._send_pending_recv_notifs(req_id)
2823
else:
2824
transfers[req_id] = in_progress
2825
return done_req_ids
@@ -2771,22 +2833,25 @@ class NixlBaseConnectorWorker:
2833
req_id: The request ID.
2834
handle: The transfer handle.
2835
"""
2774
- # (multi-read) One handle is created per remote rank, and they do not
2775
- # all fail in the same _pop_done_transfers poll. The request is
2776
- # reported failed on the first one, which pops its metadata in
2777
- # get_finished(); on the later failures only the handle cleanup is left.
2778
- # TODO (NickLucche) handle failed transfer for HMA.
2779
- # A split READ's notification is sent only after every handle succeeds.
2836
+ # A sibling READ may still be writing these blocks. Retain metadata
2837
+ # and defer invalidation until every handle is terminal.
2838
self._pending_recv_notifs.pop(req_id, None)
2839
+ if req_id in self._recving_transfers:
2840
+ self._failed_inflight_recvs.add(req_id)
2841
+ else:
2842
+ self._report_failed_recv(req_id)
2843
+ if handle is not None:
2844
+ self.nixl_wrapper.release_xfer_handle(handle)
2845
+ self.xfer_stats.record_failed_transfer()
2846
+
2847
+ def _report_failed_recv(self, req_id: str) -> None:
2848
if (meta := self._recving_metadata.get(req_id)) is not None:
2849
if not self._is_hma_required:
2850
self._invalid_block_ids.put(set(meta.local_block_ids[0]))
2851
self._failed_recv_reqs.put(req_id)
2785
- if handle is not None:
2786
- self.nixl_wrapper.release_xfer_handle(handle)
2787
- self.xfer_stats.record_failed_transfer()
2852
2853
def _send_pending_recv_notifs(self, req_id: str) -> None:
2854
+ """Send notifications deferred by split DRAM/VRAM reads."""
2855
for agent_name, notif_id in self._pending_recv_notifs.pop(req_id, []):
2856
try:
2857
self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id)
@@ -2909,10 +2974,10 @@ class NixlBaseConnectorWorker:
2974
return block_ids
2975
block_arange = np.arange(0, ratio).reshape(1, -1)
2976
# Mamba blocks have no logical<>physical discrepancy (block-size=1)
2912
- group_specs = self.kv_cache_config.transfer_groups
2977
physical_block_ids = []
2978
for i, group in enumerate(block_ids):
2915
- if _is_ssm_spec(get_representative_spec_type(group_specs[i].kv_cache_spec)):
2979
+ spec = self.kv_cache_config.transfer_groups[i].kv_cache_spec
2980
+ if _is_ssm_spec(get_representative_spec_type(spec)):
2981
physical_block_ids.append(group)
2982
else:
2983
physical_block_ids.append(
@@ -3215,9 +3280,44 @@ class NixlBaseConnectorWorker:
3280
)
3281
3282
def __del__(self):
3218
- self.shutdown()
3283
+ with contextlib.suppress(Exception):
3284
+ self.shutdown()
3285
3220
- def shutdown(self):
3286
+ def _finish_shutdown(self) -> None:
3287
+ self._recving_transfers.clear()
3288
+ try:
3289
+ for handle in self.src_xfer_handles_by_block_size.values():
3290
+ self.nixl_wrapper.release_dlist_handle(handle)
3291
+ for handles in self.src_xfer_handles_by_tp_ratio.values():
3292
+ for handle in handles:
3293
+ self.nixl_wrapper.release_dlist_handle(handle)
3294
+ for handles in self._dram_src_handles_by_tp_ratio.values():
3295
+ for handle in handles:
3296
+ self.nixl_wrapper.release_dlist_handle(handle)
3297
+ for handle in self._dram_src_handles_by_block_size.values():
3298
+ self.nixl_wrapper.release_dlist_handle(handle)
3299
+ except Exception:
3300
+ logger.exception("NIXL dlist-handle release failed at shutdown.")
3301
+ self.src_xfer_handles_by_block_size.clear()
3302
+ self.src_xfer_handles_by_tp_ratio.clear()
3303
+ self._dram_src_handles_by_tp_ratio.clear()
3304
+ self._dram_src_handles_by_block_size.clear()
3305
+ try:
3306
+ for engine_id in list(self._remote_agents):
3307
+ self._cleanup_remote_engine(engine_id, log_eviction=False)
3308
+ except Exception:
3309
+ logger.exception("NIXL remote-engine cleanup failed at shutdown.")
3310
+ try:
3311
+ for desc in self._registered_descs:
3312
+ self.nixl_wrapper.deregister_memory(desc)
3313
+ finally:
3314
+ self._registered_descs.clear()
3315
+ # Drop cache references before their owners release registered
3316
+ # host memory; handshake futures may outlive model-runner shutdown.
3317
+ self.device_kv_caches = {}
3318
+ self.host_xfer_buffers = {}
3319
+
3320
+ def shutdown(self) -> None:
3321
"""Shutdown the connector worker."""
3322
if not hasattr(self, "_handshake_initiation_executor"):
3323
# error happens during init, no need to shutdown
@@ -3227,22 +3327,4 @@ class NixlBaseConnectorWorker:
3327
for handle in handles:
3328
self.nixl_wrapper.release_xfer_handle(handle)
3329
self._recving_transfers.clear()
3230
- for handle in self.src_xfer_handles_by_block_size.values():
3231
- self.nixl_wrapper.release_dlist_handle(handle)
3232
- self.src_xfer_handles_by_block_size.clear()
3233
- for handles in self.src_xfer_handles_by_tp_ratio.values():
3234
- for handle in handles:
3235
- self.nixl_wrapper.release_dlist_handle(handle)
3236
- self.src_xfer_handles_by_tp_ratio.clear()
3237
- for handles in self._dram_src_handles_by_tp_ratio.values():
3238
- for handle in handles:
3239
- self.nixl_wrapper.release_dlist_handle(handle)
3240
- self._dram_src_handles_by_tp_ratio.clear()
3241
- for handle in self._dram_src_handles_by_block_size.values():
3242
- self.nixl_wrapper.release_dlist_handle(handle)
3243
- self._dram_src_handles_by_block_size.clear()
3244
- for engine_id in list(self._remote_agents):
3245
- self._cleanup_remote_engine(engine_id, log_eviction=False)
3246
- for desc in self._registered_descs:
3247
- self.nixl_wrapper.deregister_memory(desc)
3248
- self._registered_descs.clear()
3330
+ self._finish_shutdown()
vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py
+13
@@ -26,6 +26,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
26
KVConnectorHandshakeMetadata,
27
KVConnectorMetadata,
28
KVConnectorRole,
29
+ KVConnectorTransferResults,
30
SupportsHMA,
31
)
32
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
@@ -234,6 +235,18 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
235
assert self.connector_worker is not None
236
return self.connector_worker.get_finished()
237
238
+ def get_transfer_results(
239
+ self, finished_req_ids: set[str]
240
+ ) -> KVConnectorTransferResults:
241
+ assert self.connector_worker is not None
242
+ results = self.connector_worker.get_transfer_results()
243
+ if (
244
+ self.kv_transfer_config.kv_role == "kv_producer"
245
+ and self.connector_worker.pcp_rank > 0
246
+ ):
247
+ results.finished_sending.clear()
248
+ return results
249
+
250
def get_block_ids_with_load_errors(self) -> set[int]:
251
"""Get block IDs that failed to load via NIXL."""
252
assert self.connector_worker is not None
vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
+3
-1
@@ -305,7 +305,9 @@ class NixlConnectorMetadata(KVConnectorMetadata):
305
local_num_computed_blocks: tuple[int, ...] = (),
306
):
307
req = self._add_new_req(
308
- local_block_ids, kv_transfer_params, local_num_computed_blocks
308
+ local_block_ids,
309
+ kv_transfer_params,
310
+ local_num_computed_blocks,
311
)
312
req.remote = RemoteMeta(
313
block_ids=kv_transfer_params["remote_block_ids"],
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+17
-11
@@ -166,11 +166,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
166
167
dcp_active = self.dcp_size > 1 or remote_info.remote_dcp_size > 1
168
local_block_ids = meta.local_physical_block_ids
169
- local_region_groups = getattr(self, "region_group_ids", [])
170
- remote_region_groups = getattr(self, "dst_region_group_ids", {}).get(
171
- engine_id, local_region_groups
172
- )
173
- if local_region_groups != remote_region_groups:
169
+ remote_region_groups = self.dst_region_group_ids[engine_id]
170
+ local_region_groups = self.region_group_ids or remote_region_groups
171
+ groups_differ = local_region_groups != remote_region_groups
172
+ if groups_differ:
173
if not self.use_mla or self._has_mamba:
174
raise NotImplementedError(
175
"Different NIXL cache-group layouts are only supported for "
@@ -292,6 +291,9 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
291
spec.remote_rank
292
]
293
294
+ # Once a read routes the request to failure reporting, the
295
+ # scheduler may free and reuse its blocks, so no sibling READs
296
+ # may be posted (and P must not be notified).
297
if not self._read_blocks(
298
read_spec=spec,
299
request_id=req_id,
@@ -329,6 +331,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
331
"""
332
Post a READ point-to-point xfer request from a single local worker to
333
a single remote worker.
334
+
335
+ Returns True when the read was posted (or was unnecessary), False
336
+ when the request was routed to failure reporting — the caller must
337
+ not post further transfers for it.
338
"""
339
assert self.transfer_topo is not None
340
remote_rank = read_spec.remote_rank
@@ -416,11 +422,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
422
dst_num_blocks=self.dst_num_blocks[dst_engine_id],
423
block_size_ratio=None,
424
physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical,
419
- region_num_blocks=self.dst_region_num_blocks[dst_engine_id],
425
+ region_num_blocks=(self.dst_region_num_blocks.get(dst_engine_id) or None),
426
region_group_ids=(
427
list(range(self.num_regions))
428
if read_spec.block_ids_by_region
423
- else self.dst_region_group_ids[dst_engine_id]
429
+ else (self.dst_region_group_ids.get(dst_engine_id) or None)
430
),
431
uses_region_group_mapping=(
432
self.num_regions > 1
@@ -433,11 +439,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
439
dst_num_blocks=self.dst_num_blocks[self.engine_id],
440
block_size_ratio=block_size_ratio,
441
physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block,
436
- region_num_blocks=self.dst_region_num_blocks[self.engine_id],
442
+ region_num_blocks=(self.dst_region_num_blocks.get(self.engine_id) or None),
443
region_group_ids=(
444
list(range(self.num_regions))
445
if read_spec.block_ids_by_region
440
- else self.region_group_ids
446
+ else (self.region_group_ids or None)
447
),
448
uses_region_group_mapping=(
449
self.num_regions > 1
@@ -537,11 +543,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
543
self._pending_recv_notifs.setdefault(request_id, []).append(
544
(notif_agent, notif_id)
545
)
540
- for index, handle in enumerate(handles):
546
+ for i, handle in enumerate(handles):
547
try:
548
self.nixl_wrapper.transfer(handle)
549
except Exception:
544
- for unstarted in handles[index:]:
550
+ for unstarted in handles[i:]:
551
self.nixl_wrapper.release_xfer_handle(unstarted)
552
raise
553
self._recving_transfers[request_id].append(handle)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
+18
-5
@@ -41,6 +41,9 @@ from typing import TYPE_CHECKING, Any
41
import msgspec
42
43
from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
44
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
45
+ KVConnectorTransferResults,
46
+)
47
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
48
NixlBaseConnectorWorker,
49
)
@@ -145,7 +148,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
148
self._push_writer_thread.start()
149
logger.info("nixl-push-writer thread started (rank=%d)", self.tp_rank)
150
148
- def shutdown(self):
151
+ def shutdown(self) -> None:
152
self._push_writer_stop.set()
153
# Unblock the writer if it's waiting in the no-active-state branch.
154
self._push_writer_wake.set()
@@ -522,7 +525,14 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
525
)
526
remote_block_ids = meta.remote.block_ids
527
local_block_ids = meta.local_physical_block_ids
525
- num_groups = len(local_block_ids)
528
+ local_region_groups = self.region_group_ids
529
+ remote_region_groups = self.dst_region_group_ids[engine_id]
530
+ groups_differ = local_region_groups != remote_region_groups
531
+ if groups_differ:
532
+ raise NotImplementedError(
533
+ "NixlPushConnector does not support different producer and "
534
+ "consumer cache-group layouts"
535
+ )
536
537
# MLA latent is replicated across D's TP ranks: the tp-mapping
538
# collapses it to one rank (fine for reads), but push must WRITE every
@@ -537,6 +547,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
547
else:
548
write_ranks = list(plan.all_source_ranks)
549
550
+ num_groups = len(local_block_ids)
551
+
552
def group_ids(block_ids: BlockIds, rank: int) -> BlockIds:
553
return [
554
list(block_ids[g])
@@ -777,13 +789,14 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
789
self._reqs_to_send.pop(req_id, None)
790
return notified_req_ids
791
780
- def get_finished(self) -> tuple[set[str], set[str]]:
792
+ def get_transfer_results(self) -> KVConnectorTransferResults:
793
# Engine main thread asking for completions: also wake the writer
794
# so it gets a chance to drain NIXL notifs (heartbeats, completion
795
# notifs, late PUSH_REGs) even if it had been parked.
796
self._push_writer_wake.set()
797
786
- done_sending, done_recving = super().get_finished()
798
+ results = super().get_transfer_results()
799
+ done_sending = results.finished_sending
800
801
# ``_pop_done_transfers`` mutates ``_sending_transfers``; the
802
# writer thread also appends to it, so guard the pop.
@@ -803,4 +816,4 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
816
if done_sending:
817
self._push_writer_wake.set()
818
806
- return done_sending, done_recving
819
+ return results
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+29
-2
@@ -12,6 +12,7 @@ from vllm.v1.core.kv_cache_utils import (
12
from vllm.v1.kv_cache_interface import (
13
AttentionSpec,
14
FullAttentionSpec,
15
+ KVCacheGroupRole,
16
KVCacheSpec,
17
MLAAttentionSpec,
18
SlidingWindowMLASpec,
@@ -31,6 +32,16 @@ if TYPE_CHECKING:
32
from vllm.v1.kv_cache_interface import KVCacheConfig
33
34
35
+def get_offloading_group_ids(kv_cache_config: "KVCacheConfig") -> tuple[int, ...]:
36
+ if kv_cache_config.hisparse_host_num_blocks is None:
37
+ return tuple(range(len(kv_cache_config.kv_cache_groups)))
38
+ return tuple(
39
+ group_id
40
+ for group_id, group in enumerate(kv_cache_config.kv_cache_groups)
41
+ if group.role is KVCacheGroupRole.HISPARSE_INDEXER
42
+ )
43
+
44
+
45
def build_offloading_config(
46
vllm_config: "VllmConfig",
47
kv_cache_config: "KVCacheConfig",
@@ -43,15 +54,22 @@ def build_offloading_config(
54
engine_id = kv_transfer_config.engine_id
55
56
parallel_config = vllm_config.parallel_config
57
+ selected_groups = tuple(
58
+ (group_id, kv_cache_config.kv_cache_groups[group_id])
59
+ for group_id in get_offloading_group_ids(kv_cache_config)
60
+ )
61
+ if not selected_groups:
62
+ raise ValueError("KV offloading found no eligible cache groups.")
63
groups = tuple(
64
OffloadingGroupConfig(
65
+ group_id=group_id,
66
tokens_per_block=resolve_dcp_kv_block_size(
67
group.kv_cache_spec,
68
parallel_config.decode_context_parallel_size,
69
),
70
layer_names=tuple(group.layer_names),
71
)
54
- for group in kv_cache_config.kv_cache_groups
72
+ for group_id, group in selected_groups
73
)
74
75
_, tokens_per_hash = resolve_kv_cache_block_sizes(kv_cache_config, vllm_config)
@@ -103,11 +121,20 @@ def build_offloading_config(
121
)
122
123
worker_kv_bytes_per_block = 0
106
- if kv_cache_config.num_blocks > 0 and kv_cache_config.kv_cache_tensors:
124
+ all_groups_selected = len(selected_groups) == len(kv_cache_config.kv_cache_groups)
125
+ if (
126
+ all_groups_selected
127
+ and kv_cache_config.num_blocks > 0
128
+ and kv_cache_config.kv_cache_tensors
129
+ ):
130
# Every KVCacheTensor describes placement within the same backing allocation,
131
# so its size is the total, not a per-tensor share.
132
total_gpu_kv_bytes = kv_cache_config.kv_cache_tensors[0].size
133
worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks
134
+ elif kv_cache_config.num_blocks > 0:
135
+ worker_kv_bytes_per_block = sum(
136
+ group.kv_cache_spec.page_size_bytes for _, group in selected_groups
137
+ )
138
139
single_group_spec = (
140
kv_cache_config.kv_cache_groups[0].kv_cache_spec
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+70
-33
@@ -152,8 +152,9 @@ def resolve_mamba_align_size(
152
such groups agree on the same value.
153
"""
154
mamba_align_size: int | None = None
155
- for idx, tokens_per_block in enumerate(spec.tokens_per_block):
156
- kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec
155
+ for group in spec.config.groups:
156
+ tokens_per_block = group.tokens_per_block
157
+ kv_spec = kv_cache_config.kv_cache_groups[group.group_id].kv_cache_spec
158
if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode in (
159
"align",
160
"all",
@@ -188,8 +189,9 @@ class SchedulerOffloadConfig(NamedTuple):
189
# each segment can never serve a load hit. Relevant for hybrid
190
# architectures like DeepSeek V4 (MLA + SWA groups).
191
full_attn_tokens_per_chunk: set[int] = set()
191
- for idx, tokens_per_block in enumerate(spec.tokens_per_block):
192
- kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec
192
+ for group in spec.config.groups:
193
+ tokens_per_block = group.tokens_per_block
194
+ kv_spec = kv_cache_config.kv_cache_groups[group.group_id].kv_cache_spec
195
sw = get_sliding_window_size_in_chunks(
196
kv_spec, tokens_per_block * spec.blocks_per_chunk
197
)
@@ -214,9 +216,9 @@ class SchedulerOffloadConfig(NamedTuple):
216
retention_interval = vllm_config.cache_config.prefix_cache_retention_interval
217
218
eagle_groups = {
217
- idx
218
- for idx, g in enumerate(kv_cache_config.kv_cache_groups)
219
- if g.is_eagle_group
219
+ group.group_id
220
+ for group in spec.config.groups
221
+ if kv_cache_config.kv_cache_groups[group.group_id].is_eagle_group
222
}
223
224
use_eagle_block_drop = (
@@ -249,8 +251,10 @@ class SchedulerOffloadConfig(NamedTuple):
251
)
252
253
kv_group_configs_list: list[GroupOffloadConfig] = []
252
- for idx, tokens_per_block in enumerate(spec.tokens_per_block):
253
- kv_cache_group = kv_cache_config.kv_cache_groups[idx]
254
+ for group in spec.config.groups:
255
+ group_id = group.group_id
256
+ tokens_per_block = group.tokens_per_block
257
+ kv_cache_group = kv_cache_config.kv_cache_groups[group_id]
258
kv_spec = kv_cache_group.kv_cache_spec
259
sw = get_sliding_window_size_in_chunks(
260
kv_spec, tokens_per_block * spec.blocks_per_chunk
@@ -261,7 +265,7 @@ class SchedulerOffloadConfig(NamedTuple):
265
)
266
kv_group_configs_list.append(
267
GroupOffloadConfig(
264
- group_idx=idx,
268
+ group_idx=group_id,
269
tokens_per_block=tokens_per_block,
270
tokens_per_chunk=tokens_per_block * spec.blocks_per_chunk,
271
hashes_per_chunk=(
@@ -272,7 +276,7 @@ class SchedulerOffloadConfig(NamedTuple):
276
kv_cache_spec=kv_spec,
277
manager_cls=manager_cls,
278
kv_event_group_spec=get_offloading_event_group_spec(kv_cache_group),
275
- is_eagle_group=idx in eagle_groups,
279
+ is_eagle_group=group_id in eagle_groups,
280
requires_cow_source=(
281
isinstance(kv_spec, MambaSpec)
282
and kv_spec.mamba_cache_mode == "align"
@@ -408,9 +412,10 @@ class RequestOffloadState:
412
if new_block_id_groups is None:
413
return
414
411
- assert len(new_block_id_groups) == len(self.group_states)
412
- for group_state, new_blocks in zip(self.group_states, new_block_id_groups):
413
- group_state.block_ids.extend(new_blocks)
415
+ for group_config, group_state in zip(
416
+ self.config.kv_group_configs, self.group_states
417
+ ):
418
+ group_state.block_ids.extend(new_block_id_groups[group_config.group_idx])
419
420
def storable_chunks(
421
self,
@@ -541,11 +546,11 @@ class OffloadingConnectorScheduler:
546
547
full_attention_groups: list[int] = []
548
sliding_window_groups: list[int] = []
544
- for group_config in self.config.kv_group_configs:
549
+ for config_idx, group_config in enumerate(self.config.kv_group_configs):
550
if group_config.sliding_window_size_in_chunks is None:
546
- full_attention_groups.append(group_config.group_idx)
551
+ full_attention_groups.append(config_idx)
552
else:
548
- sliding_window_groups.append(group_config.group_idx)
553
+ sliding_window_groups.append(config_idx)
554
555
# sort sliding window groups by window size in decreasing order
556
def _sliding_window_sort_key(i: int) -> int:
@@ -737,7 +742,11 @@ class OffloadingConnectorScheduler:
742
req_status.req_context,
743
)
744
740
- def _lookup_complete_chunks(self, req_status: RequestOffloadState) -> int | None:
745
+ def _lookup_complete_chunks(
746
+ self,
747
+ req_status: RequestOffloadState,
748
+ max_num_new_tokens: int | None = None,
749
+ ) -> int | None:
750
"""
751
Find how many tokens beyond num_locally_computed_tokens can be loaded.
752
@@ -748,6 +757,10 @@ class OffloadingConnectorScheduler:
757
"""
758
num_computed_tokens = req_status.num_locally_computed_tokens
759
max_hit_size_tokens: int = req_status.req.num_tokens
760
+ if max_num_new_tokens is not None:
761
+ max_hit_size_tokens = min(
762
+ max_hit_size_tokens, num_computed_tokens + max_num_new_tokens
763
+ )
764
if self._sliding_window_groups:
765
# the last prompt token has to be recomputed to get the logprobs
766
# for sliding window attention, we must reduce by 1 to make sure
@@ -929,8 +942,12 @@ class OffloadingConnectorScheduler:
942
hash_idx = boundary_tokens // self.config.tokens_per_hash - 1
943
return make_offload_key(request.block_hashes[hash_idx], group_idx)
944
932
- def _lookup(self, req_status: RequestOffloadState) -> int | None:
933
- complete_hit = self._lookup_complete_chunks(req_status)
945
+ def _lookup(
946
+ self,
947
+ req_status: RequestOffloadState,
948
+ max_num_new_tokens: int | None = None,
949
+ ) -> int | None:
950
+ complete_hit = self._lookup_complete_chunks(req_status, max_num_new_tokens)
951
req_status.partial_tail_boundary = None
952
if complete_hit is None or not self.config.supports_partial_tail:
953
return complete_hit
@@ -939,9 +956,10 @@ class OffloadingConnectorScheduler:
956
complete_boundary = local_tokens + complete_hit
957
tokens_per_hash = self.config.tokens_per_hash
958
block_end = complete_boundary + self._partial_tail_block_size
942
- max_boundary = round_down(
943
- min(req_status.req.num_prompt_tokens - 1, block_end - 1), tokens_per_hash
944
- )
959
+ max_boundary = min(req_status.req.num_prompt_tokens - 1, block_end - 1)
960
+ if max_num_new_tokens is not None:
961
+ max_boundary = min(max_boundary, local_tokens + max_num_new_tokens)
962
+ max_boundary = round_down(max_boundary, tokens_per_hash)
963
if max_boundary <= complete_boundary:
964
return complete_hit
965
@@ -990,7 +1008,10 @@ class OffloadingConnectorScheduler:
1008
self._req_status[request.request_id] = req_status
1009
1010
def get_num_new_matched_tokens(
993
- self, request: Request, num_computed_tokens: int
1011
+ self,
1012
+ request: Request,
1013
+ num_computed_tokens: int,
1014
+ max_num_new_tokens: int | None = None,
1015
) -> tuple[int | None, bool]:
1016
"""
1017
Get number of new tokens that can be loaded beyond the
@@ -1030,7 +1051,7 @@ class OffloadingConnectorScheduler:
1051
num_hit_tokens = 0
1052
else:
1053
lookup_start = time.monotonic()
1033
- num_hit_tokens = self._lookup(req_status)
1054
+ num_hit_tokens = self._lookup(req_status, max_num_new_tokens)
1055
self._connector_stats.observe_histogram(
1056
_ConnectorMetricName.LOOKUP_SYNC_DELAY,
1057
time.monotonic() - lookup_start,
@@ -1065,11 +1086,11 @@ class OffloadingConnectorScheduler:
1086
# per group
1087
group_sizes: list[int] = []
1088
block_indices: list[int] = []
1068
- for group_config, group_state, group_blocks in zip(
1089
+ for group_config, group_state in zip(
1090
self.config.kv_group_configs,
1091
req_status.group_states,
1071
- blocks.blocks,
1092
):
1093
+ group_blocks = blocks.blocks[group_config.group_idx]
1094
self._current_batch_allocated_block_ids.update(
1095
block.block_id for block in group_blocks if block.block_id != 0
1096
)
@@ -1184,7 +1205,8 @@ class OffloadingConnectorScheduler:
1205
for grp_idx in self._sliding_window_groups
1206
)
1207
req_status.update_block_id_groups(new_block_id_groups)
1187
- for new_blocks in new_block_id_groups:
1208
+ for group_config in self.config.kv_group_configs:
1209
+ new_blocks = new_block_id_groups[group_config.group_idx]
1210
for bid in new_blocks:
1211
if bid != 0:
1212
self._current_batch_allocated_block_ids.add(bid)
@@ -1218,6 +1240,10 @@ class OffloadingConnectorScheduler:
1240
) -> dict[int, TransferJob]:
1241
store_jobs: dict[int, TransferJob] = {}
1242
num_groups = len(self.config.kv_group_configs)
1243
+ config_idx_by_group = {
1244
+ config.group_idx: idx
1245
+ for idx, config in enumerate(self.config.kv_group_configs)
1246
+ }
1247
for req_id, entries in handoffs.items():
1248
req_status = self._req_status.get(req_id)
1249
if req_status is None:
@@ -1225,7 +1251,10 @@ class OffloadingConnectorScheduler:
1251
req = req_status.req
1252
max_boundary = self._calc_num_offloadable_tokens(req_status, req.num_tokens)
1253
for group_idx, block_id, boundary in entries:
1228
- group_config = self.config.kv_group_configs[group_idx]
1254
+ config_idx = config_idx_by_group.get(group_idx)
1255
+ if config_idx is None:
1256
+ continue
1257
+ group_config = self.config.kv_group_configs[config_idx]
1258
if (
1259
block_id == 0
1260
or boundary > max_boundary
@@ -1254,9 +1283,11 @@ class OffloadingConnectorScheduler:
1283
fenced_block_ids=[block_id],
1284
)
1285
group_sizes = [0] * num_groups
1257
- group_sizes[group_idx] = 1
1286
+ group_sizes[config_idx] = 1
1287
block_indices = [0] * num_groups
1259
- block_indices[group_idx] = boundary // group_config.tokens_per_block - 1
1288
+ block_indices[config_idx] = (
1289
+ boundary // group_config.tokens_per_block - 1
1290
+ )
1291
store_jobs[job_id] = TransferJob(
1292
req_id=req_id,
1293
src_spec=GPULoadStoreSpec(
@@ -1300,6 +1331,12 @@ class OffloadingConnectorScheduler:
1331
assert len(boundaries) == 1
1332
boundary = boundaries.pop()
1333
req = req_status.req
1334
+ group_states = {
1335
+ group.group_idx: state
1336
+ for group, state in zip(
1337
+ self.config.kv_group_configs, req_status.group_states
1338
+ )
1339
+ }
1340
max_boundary = min(
1341
req.num_prompt_tokens,
1342
req_status.max_offload_tokens or req.num_prompt_tokens,
@@ -1315,7 +1352,7 @@ class OffloadingConnectorScheduler:
1352
block_idx = boundary // self._partial_tail_block_size
1353
if any(
1354
group.group_idx not in self._cow_source_groups
1318
- and block_idx >= len(req_status.group_states[group.group_idx].block_ids)
1355
+ and block_idx >= len(group_states[group.group_idx].block_ids)
1356
for group in self.config.kv_group_configs
1357
):
1358
continue
@@ -1326,7 +1363,7 @@ class OffloadingConnectorScheduler:
1363
block_ids = [
1364
cow_blocks[group.group_idx]
1365
if group.group_idx in self._cow_source_groups
1329
- else req_status.group_states[group.group_idx].block_ids[block_idx]
1366
+ else group_states[group.group_idx].block_ids[block_idx]
1367
for group in self.config.kv_group_configs
1368
]
1369
assert all(block_id != 0 for block_id in block_ids)
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+16
-8
@@ -14,6 +14,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
14
OffloadingWorkerMetadata,
15
ReqId,
16
)
17
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
18
+ get_offloading_group_ids,
19
+)
20
from vllm.logger import init_logger
21
from vllm.v1.kv_cache_interface import (
22
AttentionSpec,
@@ -65,7 +68,10 @@ class OffloadingConnectorWorker:
68
69
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
70
kv_cache_config = self.kv_cache_config
68
- num_blocks = kv_cache_config.num_blocks
71
+ selected_group_ids = get_offloading_group_ids(kv_cache_config)
72
+ selected_groups = tuple(
73
+ kv_cache_config.kv_cache_groups[group_id] for group_id in selected_group_ids
74
+ )
75
mappings = derive_canonical_mappings(
76
self.vllm_config, kv_cache_config, kv_caches
77
)
@@ -76,7 +82,9 @@ class OffloadingConnectorWorker:
82
unpadded_page_size_bytes: dict[str, int] = {}
83
# layer_name -> size of page in bytes
84
page_size_bytes: dict[str, int] = {}
79
- for kv_cache_group in kv_cache_config.kv_cache_groups:
85
+ for kv_cache_group in selected_groups:
86
+ assert not kv_cache_group.host_resident
87
+ num_blocks = kv_cache_config.num_blocks
88
group_layer_names = kv_cache_group.layer_names
89
group_kv_cache_spec = kv_cache_group.kv_cache_spec
90
if isinstance(group_kv_cache_spec, UniformTypeKVCacheSpecs):
@@ -131,8 +139,11 @@ class OffloadingConnectorWorker:
139
),
140
None,
141
)
134
- if packed_layer_name is not None:
142
+ if packed_layer_name is not None and len(selected_groups) == len(
143
+ kv_cache_config.kv_cache_groups
144
+ ):
145
(tensor,) = tensors_per_block[packed_layer_name]
146
+ num_blocks = tensor.shape[0]
147
block_stride = tensor.stride(0)
148
packed_tensor = tensor.as_strided(
149
(num_blocks, block_stride),
@@ -142,10 +153,7 @@ class OffloadingConnectorWorker:
153
self._init_worker(
154
CanonicalKVCaches(
155
[CanonicalKVCacheTensor(packed_tensor, block_stride)],
145
- [
146
- [CanonicalKVCacheRef(0, block_stride)]
147
- for _ in kv_cache_config.kv_cache_groups
148
- ],
156
+ [[CanonicalKVCacheRef(0, block_stride)] for _ in selected_groups],
157
)
158
)
159
return
@@ -192,7 +200,7 @@ class OffloadingConnectorWorker:
200
)
201
202
group_data_refs: list[list[CanonicalKVCacheRef]] = []
195
- for kv_cache_group in kv_cache_config.kv_cache_groups:
203
+ for kv_cache_group in selected_groups:
204
group_refs: list[CanonicalKVCacheRef] = []
205
for layer_name in kv_cache_group.layer_names:
206
group_refs += block_data_refs[layer_name]
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+46
-1
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Iterable
4
+from functools import cached_property
5
from typing import Any
6
7
import torch
@@ -25,6 +26,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
26
)
27
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
28
build_offloading_config,
29
+ get_offloading_group_ids,
30
)
31
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
32
OffloadingConnectorStats,
@@ -38,6 +40,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import (
40
)
41
from vllm.forward_context import ForwardContext
42
from vllm.v1.attention.backend import AttentionMetadata
43
+from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator
44
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
45
from vllm.v1.core.sched.output import SchedulerOutput
46
from vllm.v1.kv_cache_interface import KVCacheConfig
@@ -47,6 +50,25 @@ from vllm.v1.request import Request
50
51
52
class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
53
+ @cached_property
54
+ def _bounding_group_ids(self) -> tuple[int, ...]:
55
+ """Prefix-cacheable groups this connector does not offload.
56
+
57
+ Never offer tokens for a range some group cannot cover: past such a
58
+ group's own cached prefix its KV would be left unwritten.
59
+ """
60
+ offloaded = set(get_offloading_group_ids(self._kv_cache_config))
61
+ return tuple(
62
+ group_id
63
+ for group_id, group in enumerate(self._kv_cache_config.kv_cache_groups)
64
+ if group.kv_cache_spec.prefix_cacheable and group_id not in offloaded
65
+ )
66
+
67
+ @property
68
+ def scheduler(self) -> OffloadingConnectorScheduler:
69
+ assert self.connector_scheduler is not None
70
+ return self.connector_scheduler
71
+
72
@property
73
def requires_kv_delivery(self) -> bool:
74
# Runs as kv_both, but is a best-effort cache: a dropped save is just a
@@ -138,8 +160,29 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
160
) -> tuple[int | None, bool]:
161
assert self.connector_scheduler is not None
162
return self.connector_scheduler.get_num_new_matched_tokens(
141
- request, num_computed_tokens
163
+ request,
164
+ num_computed_tokens,
165
+ max_num_new_tokens=self._max_loadable_tokens(request, num_computed_tokens),
166
+ )
167
+
168
+ def _max_loadable_tokens(
169
+ self, request: "Request", num_computed_tokens: int
170
+ ) -> int | None:
171
+ """How far past ``num_computed_tokens`` a load may reach, if bounded.
172
+
173
+ Bounded by the deepest prefix the groups this connector does not
174
+ offload already hold, since nothing refills them beyond it.
175
+ """
176
+ if not self._bounding_group_ids:
177
+ return None
178
+ assert self._kv_cache_manager is not None
179
+ coordinator = self._kv_cache_manager.coordinator
180
+ assert isinstance(coordinator, HybridKVCacheCoordinator)
181
+ _, per_group_hits = coordinator.find_longest_cache_hit_per_group(
182
+ request.block_hashes, request.num_tokens - 1
183
)
184
+ bound = min(per_group_hits[group_id] for group_id in self._bounding_group_ids)
185
+ return max(0, bound - num_computed_tokens)
186
187
def update_state_after_alloc(
188
self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
@@ -185,6 +228,8 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
228
229
@classmethod
230
def get_required_kvcache_layout(cls, vllm_config: VllmConfig) -> str | None:
231
+ if vllm_config.attention_config.hisparse_config is not None:
232
+ return "BLHNC"
233
return "LBHNC"
234
235
def reset_cache(self) -> bool | None:
vllm/distributed/kv_transfer/kv_transfer_state.py
+3
-5
@@ -78,13 +78,11 @@ def ensure_kv_transfer_initialized(
78
79
global _KV_CONNECTOR_AGENT
80
81
- if vllm_config.kv_transfer_config is None:
81
+ if _KV_CONNECTOR_AGENT is not None:
82
return
83
84
- if (
85
- vllm_config.kv_transfer_config.is_kv_transfer_instance
86
- and _KV_CONNECTOR_AGENT is None
87
- ):
84
+ kv_transfer_config = vllm_config.kv_transfer_config
85
+ if kv_transfer_config is not None and kv_transfer_config.is_kv_transfer_instance:
86
# NIXL P/D requires an interleave_size equal to block_size.
87
vllm_config.adjust_dcp_kv_cache_interleave_size(kv_cache_config)
88
_sync_engine_id_across_tp(vllm_config)
vllm/model_executor/layers/attention/mla_attention.py
+100
-31
@@ -207,7 +207,7 @@ from collections.abc import Callable
207
from dataclasses import dataclass, replace
208
from enum import Enum
209
from math import lcm
210
-from typing import ClassVar, Generic, TypeVar, cast
210
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast
211
212
import numpy as np
213
import torch
@@ -284,6 +284,7 @@ from vllm.v1.attention.backend import (
284
CommonAttentionMetadata,
285
MLAAttentionImpl,
286
)
287
+from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
288
from vllm.v1.attention.backends.mla.prefill import (
289
MLAPrefillBackend,
290
get_mla_prefill_backend,
@@ -309,6 +310,11 @@ from vllm.v1.kv_cache_interface import (
310
get_kv_quant_mode,
311
)
312
313
+if TYPE_CHECKING:
314
+ from vllm.v1.attention.backends.mla.index_group import (
315
+ SparseMLAIndexGroupBuilder,
316
+ )
317
+
318
logger = init_logger(__name__)
319
320
_FP8_DTYPE = current_platform.fp8_dtype()
@@ -433,6 +439,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
439
use_sparse: bool = False,
440
indexer: object | None = None,
441
topk_indices_buffer: torch.Tensor | None = None,
442
+ index_group_builder: "SparseMLAIndexGroupBuilder | None" = None,
443
non_causal_multi_token_decode: bool = False,
444
sliding_window: int | None = None,
445
prefill_backend_cls: type[MLAPrefillBackend] | None = None,
@@ -538,9 +545,11 @@ class MLAAttention(nn.Module, AttentionLayerBase):
545
# explicitly so backbone "skip" layers (indexer=None) still find it.
546
if use_sparse:
547
extra_impl_args["topk_indices_buffer"] = topk_indices_buffer
548
+ if index_group_builder is not None:
549
+ extra_impl_args["index_group_builder"] = index_group_builder
550
551
impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls())
543
- self.impl = impl_cls( # type: ignore[assignment] # impl_cls always returns an MLAAttentionImpl subclass
552
+ impl = impl_cls(
553
num_heads=self.num_heads,
554
head_size=self.head_size,
555
scale=self.scale,
@@ -562,6 +571,13 @@ class MLAAttention(nn.Module, AttentionLayerBase):
571
indexer=indexer,
572
**extra_impl_args,
573
)
574
+ self.impl = impl # type: ignore[assignment]
575
+ index_group = getattr(impl, "index_group", None)
576
+ self.hisparse_cache = (
577
+ index_group.cache(cast(Any, impl).index_group_index)
578
+ if isinstance(index_group, HiSparseMLAIndexGroup)
579
+ else None
580
+ )
581
self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None)
582
self.is_amx_bmm_enabled = getattr(self.impl, "uses_amx_bmm", False)
583
# MLA reads this weight directly to build W_UK/W_UV, so a backend must
@@ -723,6 +739,61 @@ class MLAAttention(nn.Module, AttentionLayerBase):
739
)
740
return self._chunked_prefill_workspace_size
741
742
+ def update_kv_cache(
743
+ self,
744
+ kv_c_normed: torch.Tensor,
745
+ k_pe: torch.Tensor,
746
+ kv_cache: torch.Tensor,
747
+ slot_mapping: torch.Tensor | None,
748
+ attn_metadata: "MLACommonMetadata | None",
749
+ kv_cache_dtype: str,
750
+ k_scale: torch.Tensor,
751
+ ) -> None:
752
+ cache = self.hisparse_cache
753
+ if slot_mapping is None or (cache is not None and cache.dummy_batch):
754
+ return
755
+ kv_c_normed, k_pe, slot_mapping = maybe_gather_mla_latent_cache_inputs(
756
+ kv_c_normed,
757
+ k_pe,
758
+ slot_mapping,
759
+ attn_metadata.num_decode_tokens if attn_metadata is not None else None,
760
+ self.use_pcp,
761
+ )
762
+ assert slot_mapping is not None
763
+ if cache is not None:
764
+ kv_cache, slot_mapping, num_rows = cache.write_target(
765
+ kv_c_normed.shape[0], slot_mapping.numel()
766
+ )
767
+ kv_c_normed = kv_c_normed[:num_rows]
768
+ k_pe = k_pe[:num_rows]
769
+ self.impl.do_kv_cache_update( # type: ignore[attr-defined]
770
+ kv_c_normed,
771
+ k_pe,
772
+ kv_cache,
773
+ slot_mapping,
774
+ kv_cache_dtype,
775
+ k_scale,
776
+ )
777
+ if cache is not None:
778
+ mirror_target = cache.mirror_write_target(kv_c_normed.shape[0])
779
+ if mirror_target is not None:
780
+ mirror_cache, mirror_slots = mirror_target
781
+ self.impl.do_kv_cache_update( # type: ignore[attr-defined]
782
+ kv_c_normed,
783
+ k_pe,
784
+ mirror_cache,
785
+ mirror_slots,
786
+ kv_cache_dtype,
787
+ k_scale,
788
+ )
789
+
790
+ def prepare_kv_cache_update(
791
+ self, attn_metadata: "MLACommonMetadata | None"
792
+ ) -> None:
793
+ cache = self.hisparse_cache
794
+ if cache is not None and cache.runtime.is_group_leader:
795
+ cache.prepare_group_for_batch(attn_metadata)
796
+
797
def forward(
798
self,
799
q: torch.Tensor,
@@ -750,25 +821,18 @@ class MLAAttention(nn.Module, AttentionLayerBase):
821
f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
822
)
823
layer_slot_mapping = slot_mapping.get(self.layer_name)
753
- kv_for_cache, kpe_for_cache, layer_slot_mapping = (
754
- maybe_gather_mla_latent_cache_inputs(
755
- kv_c_normed,
756
- k_pe,
757
- layer_slot_mapping,
758
- attn_metadata.num_decode_tokens
759
- if attn_metadata is not None
760
- else None,
761
- self.use_pcp,
762
- )
763
- )
764
- self.impl.do_kv_cache_update( # type: ignore[attr-defined]
765
- kv_for_cache,
766
- kpe_for_cache,
824
+ self.prepare_kv_cache_update(attn_metadata)
825
+ self.update_kv_cache(
826
+ kv_c_normed,
827
+ k_pe,
828
self_kv_cache,
829
layer_slot_mapping,
830
+ attn_metadata,
831
self.kv_cache_dtype,
832
self._k_scale,
833
)
834
+ if self.hisparse_cache is not None:
835
+ self.hisparse_cache.finish_kv_update()
836
output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
837
self.forward_impl(
838
q,
@@ -1113,6 +1177,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
1177
return output_padded
1178
1179
def _use_sparse_mha(self, attn_metadata: "MLACommonMetadata") -> bool:
1180
+ if self.hisparse_cache is not None:
1181
+ return False
1182
prefill = attn_metadata.prefill
1183
if prefill is None:
1184
return False
@@ -1296,6 +1362,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
1362
)
1363
return MLAAttentionSpec(
1364
**common_kwargs,
1365
+ is_index_group_leader=self.indexer is not None,
1366
non_causal_multi_token_decode=self.non_causal_multi_token_decode,
1367
)
1368
@@ -1349,19 +1416,14 @@ def unified_mla_kv_cache_update(
1416
attn_metadata, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(
1417
layer_name
1418
)
1419
+ attn_layer.prepare_kv_cache_update(attn_metadata)
1420
if layer_slot_mapping is not None:
1353
- kv_c_normed, k_pe, layer_slot_mapping = maybe_gather_mla_latent_cache_inputs(
1354
- kv_c_normed,
1355
- k_pe,
1356
- layer_slot_mapping,
1357
- attn_metadata.num_decode_tokens if attn_metadata is not None else None,
1358
- attn_layer.use_pcp,
1359
- )
1360
- attn_layer.impl.do_kv_cache_update( # type: ignore[attr-defined]
1421
+ attn_layer.update_kv_cache(
1422
kv_c_normed,
1423
k_pe,
1424
kv_cache,
1425
layer_slot_mapping,
1426
+ attn_metadata,
1427
kv_cache_dtype,
1428
k_scale,
1429
)
@@ -1409,6 +1471,8 @@ def unified_mla_attention_with_output(
1471
del kv_cache_dummy_dep
1472
layer_name = _resolve_layer_name(layer_name)
1473
attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name)
1474
+ if layer.hisparse_cache is not None:
1475
+ layer.hisparse_cache.finish_kv_update()
1476
layer.forward_impl(
1477
q,
1478
kv_c_normed,
@@ -1501,6 +1565,7 @@ class MLACommonBackend(AttentionBackend):
1565
def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec":
1566
"""Per-token-head modes pack an inline fp32 scale pair after the
1567
latent data (single-sided: ``head_size_v == 0`` for MLA)."""
1568
+ spec = super().customize_spec(spec)
1569
mode = spec.kv_quant_mode
1570
if spec.state_content_bytes is not None or not mode.is_per_token_head:
1571
return spec
@@ -2090,7 +2155,11 @@ def build_mla_chunked_context_metadata(
2155
local_token_offset += num_local_tokens
2156
2157
seq_lens_cpu = _flat_int32(seq_lens_flat)
2093
- starts = _flat_int32(starts_flat).to(device, non_blocking=True)
2158
+ starts_and_context_lens = _flat_int32(starts_flat + context_lens).to(
2159
+ device, non_blocking=True
2160
+ )
2161
+ starts = starts_and_context_lens[: len(starts_flat)]
2162
+ context_lens_gpu = starts_and_context_lens[len(starts_flat) :]
2163
cu_seq_lens = _flat_int32(cu_seq_lens_flat).to(device, non_blocking=True)
2164
cu_seqlens_q = _flat_int32(cu_seqlens_q_flat).to(device, non_blocking=True)
2165
token_to_seq = _flat_int32(np.concatenate(token_to_seq_parts)).to(
@@ -2145,7 +2214,7 @@ def build_mla_chunked_context_metadata(
2214
chunks.append(chunk)
2215
2216
return MLACommonPrefillMetadata.ChunkedContextMetadata(
2148
- context_lens=context_lens_cpu.to(device, non_blocking=True),
2217
+ context_lens=context_lens_gpu,
2218
workspace=chunked_prefill_workspace,
2219
chunks=chunks,
2220
context_lens_list=context_lens,
@@ -2835,12 +2904,12 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2904
output_lse = None
2905
for chunk in chunked_context.chunks:
2906
toks = chunk.num_context_tokens
2838
- block_table = prefill_metadata.block_table[chunk.request_slice]
2907
+ chunk_block_table = prefill_metadata.block_table[chunk.request_slice]
2908
if self.kv_cache_dtype == "fp8_ds_mla":
2909
ops.cp_gather_and_upconvert_fp8_kv_cache(
2910
src_cache=kv_c_and_k_pe_cache,
2911
dst=workspace[:toks],
2843
- block_table=block_table,
2912
+ block_table=chunk_block_table,
2913
workspace_starts=chunk.cu_seq_lens,
2914
batch_size=chunk.num_requests,
2915
seq_starts=chunk.starts,
@@ -2853,7 +2922,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2922
ops.gather_mla_context_cache_cpu(
2923
src_cache=kv_c_and_k_pe_cache,
2924
dst=workspace[:toks],
2856
- block_table=block_table,
2925
+ block_table=chunk_block_table,
2926
starts=chunk.starts,
2927
cu_seq_lens=chunk.cu_seq_lens,
2928
)
@@ -2861,7 +2930,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2930
ops.gather_and_maybe_dequant_cache(
2931
src_cache=kv_c_and_k_pe_cache,
2932
dst=workspace,
2864
- block_table=block_table,
2933
+ block_table=chunk_block_table,
2934
cu_seq_lens=chunk.cu_seq_lens,
2935
token_to_seq=chunk.token_to_seq,
2936
num_tokens=toks,
@@ -2874,7 +2943,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2943
ops.cp_gather_cache(
2944
src_cache=kv_c_and_k_pe_cache,
2945
dst=workspace[:toks],
2877
- block_table=block_table,
2946
+ block_table=chunk_block_table,
2947
cu_seq_lens=chunk.cu_seq_lens,
2948
batch_size=chunk.num_requests,
2949
seq_starts=chunk.starts,
vllm/model_executor/layers/attention/sparse_mla_attention.py
+166
-25
@@ -3,13 +3,15 @@
3
"""Shared MHA implementation and metadata builder for sparse MLA backends."""
4
5
import math
6
+from dataclasses import dataclass
7
from shutil import which
7
-from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar, cast
8
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar
9
10
import numpy as np
11
import torch
12
13
from vllm import _custom_ops as ops
14
+from vllm.config import get_current_vllm_config
15
from vllm.distributed import (
16
get_dcp_group,
17
get_tensor_model_parallel_world_size,
@@ -17,6 +19,7 @@ from vllm.distributed import (
19
from vllm.logger import init_logger
20
from vllm.model_executor.layers.attention.mla_attention import (
21
MLACommonBaseImpl,
22
+ MLACommonDecodeMetadata,
23
MLACommonMetadata,
24
MLACommonPrefillMetadata,
25
accumulate_mla_context_chunk,
@@ -33,11 +36,19 @@ from vllm.utils.torch_utils import (
36
is_quantized_kv_cache,
37
np_to_pinned_tensor,
38
)
36
-from vllm.v1.attention.backend import AttentionMetadata, AttentionMetadataBuilder
39
+from vllm.v1.attention.backend import AttentionMetadataBuilder
40
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
41
+from vllm.v1.attention.backends.mla.index_group import (
42
+ SparseMLAIndexGroup,
43
+ SparseMLAIndexGroupBuilder,
44
+)
45
from vllm.v1.attention.backends.utils import split_decodes_and_prefills
46
from vllm.v1.attention.ops.dcp import MLADCPManager
47
from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
48
+from vllm.v1.hisparse.runtime import (
49
+ HiSparsePrefillStagingPlan,
50
+ build_hisparse_prefill_staging_plan,
51
+)
52
53
if TYPE_CHECKING:
54
from vllm.config import VllmConfig
@@ -47,8 +58,6 @@ if TYPE_CHECKING:
58
59
logger = init_logger(__name__)
60
50
-T = TypeVar("T", bound=AttentionMetadata)
51
-
61
GLOBAL_TOPK_MASK_MAX_BYTES = 128 * 1024 * 1024 # 128 MiB
62
63
@@ -116,9 +125,45 @@ def _is_masked_mha_available(
125
return fa_version == 4 and not is_quantized_kv_cache(kv_cache_dtype)
126
127
128
+def _use_dense_mha_prefill(
129
+ vllm_config: "VllmConfig", prefill_max_seq_len: int, topk_tokens: int
130
+) -> bool:
131
+ attention_config = vllm_config.attention_config
132
+ return (
133
+ attention_config.hisparse_config is None
134
+ and not attention_config.sparse_mla_force_mqa
135
+ and prefill_max_seq_len <= topk_tokens
136
+ )
137
+
138
+
139
+@dataclass
140
+class SparseMLAPrefillMetadata(MLACommonPrefillMetadata):
141
+ host_staging_plan: HiSparsePrefillStagingPlan | None = None
142
+
143
+
144
+@dataclass(kw_only=True)
145
+class SparseMLACommonMetadata(MLACommonMetadata[MLACommonDecodeMetadata]):
146
+ block_table: torch.Tensor
147
+ req_id_per_token: torch.Tensor
148
+ seq_lens: torch.Tensor | None = None
149
+ block_size: int = 64
150
+ topk_tokens: int = 2048
151
+ num_decodes: int = 0
152
+ num_prefills: int = 0
153
+ num_decode_tokens: int = 0
154
+ decode_max_query_len: int = 0
155
+ prefill_max_seq_len: int = 0
156
+ prefill: SparseMLAPrefillMetadata | None = None
157
+ cp_kv_cache_interleave_size: int = 1
158
+
159
+
160
+T = TypeVar("T", bound=SparseMLACommonMetadata)
161
+
162
+
163
class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
164
metadata_cls: type[T]
165
require_uniform_decodes: ClassVar[bool] = False
166
+ hisparse_supports_multi_token_decode: ClassVar[bool] = False
167
168
def __init__(
169
self,
@@ -138,6 +183,12 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
183
dtype=torch.int32,
184
device=device,
185
)
186
+ self._context_lens_cpu = torch.empty(
187
+ vllm_config.scheduler_config.max_num_seqs,
188
+ dtype=torch.int32,
189
+ device="cpu",
190
+ pin_memory=PIN_MEMORY,
191
+ )
192
parallel_config = vllm_config.parallel_config
193
self.use_pcp = parallel_config.prefill_context_parallel_size > 1
194
try:
@@ -197,6 +248,22 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
248
layer_prefill_backend.clone() if layer_prefill_backend is not None else None
249
)
250
251
+ def _init_reorder_batch_threshold(
252
+ self,
253
+ reorder_batch_threshold: int | None = 1,
254
+ supports_spec_as_decode: bool = False,
255
+ supports_dcp_with_varlen: bool = False,
256
+ ) -> None:
257
+ if self.vllm_config.attention_config.hisparse_config is not None:
258
+ reorder_batch_threshold = 1
259
+ if not self.hisparse_supports_multi_token_decode:
260
+ supports_spec_as_decode = False
261
+ super()._init_reorder_batch_threshold(
262
+ reorder_batch_threshold,
263
+ supports_spec_as_decode,
264
+ supports_dcp_with_varlen,
265
+ )
266
+
267
@staticmethod
268
def determine_chunked_prefill_workspace_size(vllm_config: "VllmConfig") -> int:
269
scheduler_config = vllm_config.scheduler_config
@@ -245,9 +312,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
312
313
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
314
assert seq_lens_cpu is not None
248
- context_lens_cpu = torch.empty(
249
- num_prefills, dtype=seq_lens_cpu.dtype, pin_memory=PIN_MEMORY
250
- )
315
+ context_lens_cpu = self._context_lens_cpu[:num_prefills]
316
torch.subtract(
317
seq_lens_cpu[num_decodes : num_decodes + num_prefills],
318
prefill_query_lens_cpu,
@@ -289,34 +354,66 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
354
prefill_max_query_len,
355
prefill_query_lens_cpu,
356
) = self._build_prefill_fields(common_attn_metadata, num_decodes, num_prefills)
357
+ decode_max_query_len = 0
358
+ if num_decodes:
359
+ if common_attn_metadata.max_logits_per_req is not None:
360
+ decode_max_query_len = common_attn_metadata.max_logits_per_req
361
+ elif num_prefills == 0:
362
+ decode_max_query_len = common_attn_metadata.max_query_len
363
+ else:
364
+ decode_query_lens_cpu = torch.diff(
365
+ common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]
366
+ )
367
+ decode_max_query_len = int(decode_query_lens_cpu.max().item())
368
369
prefill_max_seq_len = 0
294
- prefill: MLACommonPrefillMetadata | None = None
370
+ prefill: SparseMLAPrefillMetadata | None = None
371
if num_prefills > 0 and self._prefill_backend is not None:
372
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
373
assert seq_lens_cpu is not None
374
prefill_max_seq_len = int(
375
seq_lens_cpu[num_decodes : num_decodes + num_prefills].max().item()
376
)
301
- prefill = MLACommonPrefillMetadata(
302
- block_table=common_attn_metadata.block_table_tensor[num_decodes:, ...],
377
+ block_table = common_attn_metadata.block_table_tensor[num_decodes:, ...]
378
+ chunked_context = self._build_chunked_context_fields(
379
+ common_attn_metadata,
380
+ num_decodes,
381
+ num_prefills,
382
+ prefill_query_lens_cpu,
383
+ )
384
+ staging_plan = None
385
+ if self.vllm_config.attention_config.hisparse_config is not None:
386
+ prefill_seq_lens_cpu = seq_lens_cpu[
387
+ num_decodes : num_decodes + num_prefills
388
+ ]
389
+ staging_block_capacity = int(
390
+ (
391
+ (prefill_seq_lens_cpu + self.kv_cache_spec.block_size - 1)
392
+ // self.kv_cache_spec.block_size
393
+ ).sum()
394
+ )
395
+ staging_plan = build_hisparse_prefill_staging_plan(
396
+ block_table,
397
+ common_attn_metadata.seq_lens[num_decodes:],
398
+ self.kv_cache_spec.block_size,
399
+ staging_block_capacity,
400
+ )
401
+ if chunked_context is not None:
402
+ block_table = staging_plan.block_table
403
+ prefill = SparseMLAPrefillMetadata(
404
+ block_table=block_table,
405
query_start_loc=prefill_query_start_loc,
406
max_query_len=prefill_max_query_len,
407
query_lens_cpu=prefill_query_lens_cpu,
306
- chunked_context=self._build_chunked_context_fields(
307
- common_attn_metadata,
308
- num_decodes,
309
- num_prefills,
310
- prefill_query_lens_cpu,
311
- ),
408
+ chunked_context=chunked_context,
409
q_data_type=self.model_config.dtype,
410
output_dtype=self.model_config.dtype,
411
prefill_backend=self._prefill_backend,
315
- use_dense_mha=(
316
- prefill_max_seq_len <= self.topk_tokens
317
- and not self.vllm_config.attention_config.sparse_mla_force_mqa
412
+ use_dense_mha=_use_dense_mha_prefill(
413
+ self.vllm_config, prefill_max_seq_len, self.topk_tokens
414
),
415
topk_mask_workspace=self.topk_mask_workspace,
416
+ host_staging_plan=staging_plan,
417
)
418
self._prefill_backend.prepare_metadata(prefill)
419
@@ -335,6 +432,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
432
num_decodes=num_decodes,
433
num_prefills=num_prefills,
434
num_decode_tokens=num_decode_tokens,
435
+ decode_max_query_len=decode_max_query_len,
436
prefill_max_seq_len=prefill_max_seq_len,
437
prefill=prefill,
438
cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
@@ -560,6 +658,7 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic
658
kv_b_proj: "ColumnParallelLinear",
659
indexer: object | None = None,
660
topk_indices_buffer: torch.Tensor | None = None,
661
+ index_group_builder: SparseMLAIndexGroupBuilder | None = None,
662
q_pad_num_heads: int | None = None,
663
) -> None:
664
super().__init__(
@@ -577,6 +676,21 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic
676
)
677
678
self.init_topk_indices_buffer(indexer, topk_indices_buffer)
679
+ self.index_group: SparseMLAIndexGroup | None = None
680
+ self.index_group_index = 0
681
+ if index_group_builder is None and self.topk_indices_buffer is not None:
682
+ index_group_builder = SparseMLAIndexGroupBuilder(self.topk_indices_buffer)
683
+ if index_group_builder is not None:
684
+ vllm_config = get_current_vllm_config()
685
+ self.index_group, self.index_group_index = (
686
+ index_group_builder.register_layer(
687
+ indexer is not None,
688
+ vllm_config,
689
+ head_size=head_size,
690
+ kv_cache_dtype=kv_cache_dtype,
691
+ )
692
+ )
693
+
694
self._use_flashinfer_concat_mla_k = (
695
has_flashinfer()
696
and which("ninja") is not None
@@ -593,6 +707,33 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic
707
kv_cache_dtype=kv_cache_dtype,
708
)
709
710
+ def record_logical_topk_ready(self) -> None:
711
+ if self.index_group is not None:
712
+ self.index_group.set_logical_topk_ready(self.index_group_index)
713
+
714
+ def prepare_for_batch(self, attn_metadata: T | None) -> None:
715
+ if self.index_group is not None:
716
+ self.index_group.prepare_for_batch(self.index_group_index, attn_metadata)
717
+
718
+ def _convert_logical_to_physical_topk(
719
+ self,
720
+ logical_topk_indices: torch.Tensor,
721
+ attn_metadata: Any,
722
+ *,
723
+ block_stride_rows: int | None,
724
+ return_valid_counts: bool,
725
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
726
+ group = self.index_group
727
+ assert group is not None
728
+ assert self.dcp_world_size == 1
729
+ return group.convert_logical_to_physical_topk(
730
+ self.index_group_index,
731
+ logical_topk_indices,
732
+ attn_metadata,
733
+ block_stride_rows=block_stride_rows,
734
+ return_valid_counts=return_valid_counts,
735
+ )
736
+
737
@staticmethod
738
def masked_mha_workspace_fits(prefill: MLACommonPrefillMetadata) -> bool:
739
"""Whether this prefill batch's top-k masks fit the workspace."""
@@ -853,8 +994,10 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic
994
output: torch.Tensor,
995
output_scale: torch.Tensor | None = None,
996
) -> None:
856
- prefill_max_seq_len = attn_metadata.prefill_max_seq_len # type: ignore[attr-defined]
857
- topk_tokens = attn_metadata.topk_tokens # type: ignore[attr-defined]
997
+ prefill_metadata = attn_metadata.prefill
998
+ assert prefill_metadata is not None
999
+ prefill_max_seq_len = attn_metadata.prefill_max_seq_len
1000
+ topk_tokens = attn_metadata.topk_tokens
1001
force_dense = getattr(self, "_sparse_mla_force_dense_mha", False)
1002
force_masked = getattr(self, "_sparse_mla_force_masked_mha", False)
1003
if force_dense or (prefill_max_seq_len <= topk_tokens and not force_masked):
@@ -863,7 +1006,7 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic
1006
kv_c_normed,
1007
k_pe,
1008
kv_c_and_k_pe_cache,
866
- cast(MLACommonMetadata, attn_metadata),
1009
+ attn_metadata,
1010
k_scale,
1011
output,
1012
output_scale,
@@ -871,13 +1014,11 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic
1014
1015
assert output_scale is None
1016
assert self.masked_mha_available
874
- prefill_metadata = attn_metadata.prefill # type: ignore[attr-defined]
875
- assert prefill_metadata is not None
1017
assert prefill_metadata.query_lens_cpu is not None
1018
assert self.topk_indices_buffer is not None
1019
1020
q_lens = prefill_metadata.query_lens_cpu.tolist()
880
- num_decode_tokens = attn_metadata.num_decode_tokens # type: ignore[attr-defined]
1021
+ num_decode_tokens = attn_metadata.num_decode_tokens
1022
topk_all = self.topk_indices_buffer[
1023
num_decode_tokens : num_decode_tokens + q.shape[0]
1024
]
vllm/model_executor/layers/mla.py
+10
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from dataclasses import dataclass
4
+from typing import TYPE_CHECKING
5
6
import torch
7
@@ -11,6 +12,11 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
12
from vllm.models.common.ops import fused_q_kv_rmsnorm
13
from vllm.platforms import current_platform
14
15
+if TYPE_CHECKING:
16
+ from vllm.v1.attention.backends.mla.index_group import (
17
+ SparseMLAIndexGroupBuilder,
18
+ )
19
+
20
21
@dataclass
22
class MLAModules:
@@ -30,6 +36,7 @@ class MLAModules:
36
topk_indices_buffer: torch.Tensor | None
37
indexer_rotary_emb: torch.nn.Module | None = None
38
g_proj: torch.nn.Module | None = None
39
+ index_group_builder: "SparseMLAIndexGroupBuilder | None" = None
40
41
42
# --8<-- [start:multi_head_latent_attention]
@@ -128,6 +135,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
135
use_sparse=self.is_sparse,
136
indexer=self.indexer,
137
topk_indices_buffer=mla_modules.topk_indices_buffer,
138
+ index_group_builder=mla_modules.index_group_builder,
139
non_causal_multi_token_decode=non_causal_multi_token_decode,
140
)
141
indexer_op = getattr(self.indexer, "indexer_op", None)
@@ -223,6 +231,8 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
231
232
if self.indexer and self.is_sparse and not self.skip_topk:
233
self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb)
234
+ if self.is_sparse:
235
+ self.mla_attn.impl.record_logical_topk_ready() # type: ignore[attr-defined]
236
237
if llama_4_scaling is not None:
238
q *= llama_4_scaling
vllm/model_executor/layers/sparse_attn_indexer.py
+1
-1
@@ -425,7 +425,6 @@ def sparse_attn_indexer(
425
quant_block_size,
426
scale_fmt,
427
)
428
-
428
# The indexer and main MLA may classify the same short extend differently
429
# because they use independent decode thresholds. Only the main MLA route
430
# can determine whether the top-k indices will be consumed.
@@ -955,6 +954,7 @@ class SparseAttnIndexer(CustomOp):
954
self.dcp_rank,
955
self.dcp_world_size,
956
self.cp_kv_cache_interleave_size,
957
+ False,
958
candidate_blocks=self.candidate_blocks,
959
candidate_block_size=self.candidate_block_size,
960
candidate_write=self.candidate_write,
vllm/model_executor/models/deepseek_v2.py
+24
-1
@@ -103,10 +103,14 @@ from vllm.platforms import current_platform
103
from vllm.sequence import IntermediateTensors
104
from vllm.utils.torch_utils import direct_register_custom_op
105
from vllm.v1.attention.backend import AttentionBackend, AttentionType
106
+from vllm.v1.attention.backends.mla.index_group import (
107
+ SparseMLAIndexGroupBuilder,
108
+ get_sparse_mla_index_group_max_rows,
109
+)
110
from vllm.v1.attention.backends.mla.indexer import (
111
DeepseekV32IndexerBackend,
112
)
109
-from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
113
+from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec, SparseCacheRole
114
115
from .interfaces import (
116
MixtureOfExperts,
@@ -656,6 +660,7 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):
660
num_kv_heads=1,
661
head_size=self.head_dim,
662
dtype=self.dtype,
663
+ cache_role=SparseCacheRole.INDEXER,
664
) # Only has one vector instead of K + V
665
666
def forward(self): ...
@@ -1012,6 +1017,7 @@ class DeepseekV2MLAAttention(nn.Module):
1017
quant_config: QuantizationConfig | None = None,
1018
prefix: str = "",
1019
topk_indices_buffer: torch.Tensor | None = None,
1020
+ index_group_builder: SparseMLAIndexGroupBuilder | None = None,
1021
input_size: int | None = None,
1022
reduce_results: bool = True,
1023
non_causal_multi_token_decode: bool = False,
@@ -1199,6 +1205,7 @@ class DeepseekV2MLAAttention(nn.Module):
1205
indexer_rotary_emb=self.indexer_rope_emb,
1206
is_sparse=self.is_v32,
1207
topk_indices_buffer=topk_indices_buffer,
1208
+ index_group_builder=index_group_builder,
1209
)
1210
1211
self.mla_attn = MultiHeadLatentAttentionWrapper(
@@ -1242,6 +1249,7 @@ class DeepseekV2DecoderLayer(nn.Module):
1249
prefix: str,
1250
config: DeepseekV2Config | None = None,
1251
topk_indices_buffer: torch.Tensor | None = None,
1252
+ index_group_builder: SparseMLAIndexGroupBuilder | None = None,
1253
) -> None:
1254
super().__init__()
1255
@@ -1290,6 +1298,11 @@ class DeepseekV2DecoderLayer(nn.Module):
1298
and parallel_config.pipeline_parallel_size == 1
1299
and is_moe_layer
1300
)
1301
+ attn_kwargs = (
1302
+ {"index_group_builder": index_group_builder}
1303
+ if attn_cls is DeepseekV2MLAAttention
1304
+ else {}
1305
+ )
1306
self.self_attn = attn_cls(
1307
vllm_config=vllm_config,
1308
config=config,
@@ -1306,6 +1319,7 @@ class DeepseekV2DecoderLayer(nn.Module):
1319
prefix=f"{prefix}.self_attn",
1320
topk_indices_buffer=topk_indices_buffer,
1321
reduce_results=not self.use_sequence_parallel_moe,
1322
+ **attn_kwargs,
1323
)
1324
1325
if is_moe_layer:
@@ -1429,6 +1443,14 @@ class DeepseekV2Model(nn.Module):
1443
)
1444
else:
1445
topk_indices_buffer = None
1446
+ index_group_builder = (
1447
+ SparseMLAIndexGroupBuilder(
1448
+ topk_indices_buffer,
1449
+ get_sparse_mla_index_group_max_rows(vllm_config),
1450
+ )
1451
+ if topk_indices_buffer is not None
1452
+ else None
1453
+ )
1454
1455
if get_pp_group().is_first_rank:
1456
self.embed_tokens = VocabParallelEmbedding(
@@ -1445,6 +1467,7 @@ class DeepseekV2Model(nn.Module):
1467
vllm_config=vllm_config,
1468
prefix=prefix,
1469
topk_indices_buffer=topk_indices_buffer,
1470
+ index_group_builder=index_group_builder,
1471
),
1472
prefix=f"{prefix}.layers",
1473
)
vllm/model_executor/models/glm4_moe_lite.py
+20
@@ -67,6 +67,10 @@ from vllm.model_executor.models.glm4_moe import (
67
)
68
from vllm.platforms import current_platform
69
from vllm.sequence import IntermediateTensors
70
+from vllm.v1.attention.backends.mla.index_group import (
71
+ SparseMLAIndexGroupBuilder,
72
+ get_sparse_mla_index_group_max_rows,
73
+)
74
75
from .interfaces import SupportsLoRA, SupportsPP
76
from .utils import (
@@ -108,6 +112,7 @@ class Glm4MoeLiteDecoderLayer(nn.Module):
112
prefix: str,
113
config: "Glm4MoeLiteConfig | None" = None,
114
topk_indices_buffer: torch.Tensor | None = None,
115
+ index_group_builder: SparseMLAIndexGroupBuilder | None = None,
116
) -> None:
117
super().__init__()
118
@@ -136,6 +141,11 @@ class Glm4MoeLiteDecoderLayer(nn.Module):
141
attn_cls = Glm4MoeLiteMLAAttention
142
else:
143
attn_cls = Glm4MoeLiteAttention
144
+ attn_kwargs: dict[str, typing.Any] = (
145
+ {"index_group_builder": index_group_builder}
146
+ if attn_cls is Glm4MoeLiteMLAAttention
147
+ else {}
148
+ )
149
150
self.self_attn = attn_cls(
151
vllm_config=vllm_config,
@@ -152,6 +162,7 @@ class Glm4MoeLiteDecoderLayer(nn.Module):
162
quant_config=quant_config,
163
prefix=f"{prefix}.self_attn",
164
topk_indices_buffer=topk_indices_buffer,
165
+ **attn_kwargs,
166
)
167
168
if (
@@ -234,6 +245,14 @@ class Glm4MoeLiteModel(nn.Module):
245
)
246
else:
247
topk_indices_buffer = None
248
+ index_group_builder = (
249
+ SparseMLAIndexGroupBuilder(
250
+ topk_indices_buffer,
251
+ get_sparse_mla_index_group_max_rows(vllm_config),
252
+ )
253
+ if topk_indices_buffer is not None
254
+ else None
255
+ )
256
257
if get_pp_group().is_first_rank:
258
self.embed_tokens = VocabParallelEmbedding(
@@ -252,6 +271,7 @@ class Glm4MoeLiteModel(nn.Module):
271
config=config,
272
prefix=prefix,
273
topk_indices_buffer=topk_indices_buffer,
274
+ index_group_builder=index_group_builder,
275
),
276
prefix=f"{prefix}.layers",
277
)
vllm/model_executor/models/glm4_moe_lite_mtp.py
+13
@@ -51,6 +51,10 @@ from vllm.model_executor.model_loader.weight_utils import (
51
)
52
from vllm.platforms import current_platform
53
from vllm.sequence import IntermediateTensors
54
+from vllm.v1.attention.backends.mla.index_group import (
55
+ SparseMLAIndexGroupBuilder,
56
+ get_sparse_mla_index_group_max_rows,
57
+)
58
59
from .glm4_moe_lite import (
60
Glm4MixtureOfExperts,
@@ -108,6 +112,14 @@ class Glm4MoeLiteMultiTokenPredictorLayer(nn.Module):
112
)
113
else:
114
topk_indices_buffer = None
115
+ index_group_builder = (
116
+ SparseMLAIndexGroupBuilder(
117
+ topk_indices_buffer,
118
+ get_sparse_mla_index_group_max_rows(vllm_config),
119
+ )
120
+ if topk_indices_buffer is not None
121
+ else None
122
+ )
123
124
self.shared_head = SharedHead(
125
config=config, prefix=prefix, quant_config=quant_config
@@ -117,6 +129,7 @@ class Glm4MoeLiteMultiTokenPredictorLayer(nn.Module):
129
prefix=prefix,
130
config=self.config,
131
topk_indices_buffer=topk_indices_buffer,
132
+ index_group_builder=index_group_builder,
133
)
134
135
def forward(
vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py
+29
@@ -47,6 +47,35 @@ _FLASHINFER_SM120_SPARSE_MLA_DECODE_LABELS = {
47
_SPARSE_MLA_MIXED_WARMUP_TOKENS = 16
48
49
50
+def autotune_hisparse_flashinfer_attention(runner: "GPUModelRunner") -> None:
51
+ """Autotune each HiSparse FlashInfer sparse-MLA configuration."""
52
+ from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
53
+ FlashInferMLASparseImpl,
54
+ )
55
+
56
+ tuned: set[tuple[object, ...]] = set()
57
+ for layer in runner.vllm_config.compilation_config.static_forward_context.values():
58
+ impl = getattr(layer, "impl", None)
59
+ if not isinstance(impl, FlashInferMLASparseImpl):
60
+ continue
61
+ if getattr(layer, "hisparse_cache", None) is None:
62
+ continue
63
+ if impl.topk_indices_buffer is None:
64
+ continue
65
+ key = (
66
+ impl.kv_cache_dtype,
67
+ impl.num_heads,
68
+ impl.qk_nope_head_dim,
69
+ impl.qk_rope_head_dim,
70
+ impl.kv_lora_rank,
71
+ impl.topk_indices_buffer.shape[1],
72
+ )
73
+ if key in tuned:
74
+ continue
75
+ impl.autotune_hisparse_decode(layer)
76
+ tuned.add(key)
77
+
78
+
79
def _attention_backend_name(backend: object) -> str | None:
80
get_name = getattr(backend, "get_name", None)
81
if get_name is None:
vllm/model_executor/warmup/kernel_warmup.py
+14
-2
@@ -25,6 +25,7 @@ from vllm.model_executor.warmup.flashinfer_autotune_cache import (
25
write_flashinfer_autotune_cache,
26
)
27
from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import (
28
+ autotune_hisparse_flashinfer_attention,
29
deepseek_v4_sparse_mla_attention_warmup,
30
flashinfer_sparse_mla_decode_autotune_warmup,
31
)
@@ -353,9 +354,12 @@ def _flashinfer_autotune_token_counts(runner: "GPUModelRunner") -> tuple[int, ..
354
return tuple(dict.fromkeys(token_counts))
355
356
356
-def _run_flashinfer_autotune_dummy_runs(runner: "GPUModelRunner") -> None:
357
+def _run_flashinfer_autotune_dummy_runs(
358
+ runner: "GPUModelRunner", *, skip_attn: bool = False
359
+) -> None:
360
import vllm.utils.flashinfer as fi_utils
361
362
+ dummy_run_kwargs = {"skip_attn": True} if skip_attn else {}
363
for num_tokens in _flashinfer_autotune_token_counts(runner):
364
tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens)
365
logger.info(
@@ -369,6 +373,7 @@ def _run_flashinfer_autotune_dummy_runs(runner: "GPUModelRunner") -> None:
373
skip_eplb=True,
374
is_profile=True,
375
randomize_inputs=True,
376
+ **dummy_run_kwargs,
377
)
378
379
@@ -431,7 +436,14 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
436
torch.inference_mode(),
437
fi_utils.autotune(tune_mode=True, **autotune_kwargs),
438
):
434
- _run_flashinfer_autotune_dummy_runs(runner)
439
+ hisparse_enabled = (
440
+ runner.vllm_config.attention_config.hisparse_config is not None
441
+ )
442
+ if hisparse_enabled:
443
+ # HiSparse hot-buffer attention is bounded by decode batch
444
+ # size, not the prefill-sized batch used for the full model.
445
+ autotune_hisparse_flashinfer_attention(runner)
446
+ _run_flashinfer_autotune_dummy_runs(runner, skip_attn=hisparse_enabled)
447
replayssm_autotune_warmup(runner)
448
_autotune_kimi_k3_kda_qkvg(runner.get_model())
449
finally:
vllm/models/deepseek_v32/attention.py
+36
-3
@@ -42,6 +42,9 @@ from vllm.v1.attention.ops.pcp import (
42
43
if TYPE_CHECKING:
44
from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata
45
+ from vllm.v1.attention.backends.mla.index_group import (
46
+ SparseMLAIndexGroupBuilder,
47
+ )
48
49
50
class DeepseekV32Indexer(nn.Module):
@@ -126,6 +129,7 @@ class DeepseekV32Attention(MLAAttention):
129
prefix: str,
130
topk_indices_buffer: torch.Tensor | None = None,
131
attn_backend: "type | None" = None,
132
+ index_group_builder: "SparseMLAIndexGroupBuilder | None" = None,
133
) -> None:
134
quant_config = vllm_config.quant_config
135
cache_config = vllm_config.cache_config
@@ -209,6 +213,7 @@ class DeepseekV32Attention(MLAAttention):
213
use_sparse=True,
214
indexer=indexer,
215
topk_indices_buffer=topk_indices_buffer,
216
+ index_group_builder=index_group_builder,
217
attn_backend=attn_backend,
218
)
219
@@ -304,6 +309,14 @@ class DeepseekV32Attention(MLAAttention):
309
slot_mapping = forward_context.slot_mapping
310
assert isinstance(slot_mapping, dict)
311
mla_slot = slot_mapping.get(self.layer_name)
312
+ indexer_slot = (
313
+ slot_mapping.get(self.indexer.k_cache.prefix)
314
+ if self.indexer is not None
315
+ else None
316
+ )
317
+ hisparse_cache = self.hisparse_cache
318
+ layer_attn_metadata, _, _, _ = get_attention_context(self.layer_name)
319
+ self.impl.prepare_for_batch(layer_attn_metadata)
320
321
if self.indexer is not None and not self.skip_topk:
322
has_indexer = True
@@ -331,12 +344,17 @@ class DeepseekV32Attention(MLAAttention):
344
mla_k_scale = None
345
indexer_k_cache = None
346
mla_slot = None
347
+ indexer_slot = None
348
else:
335
- mla_kv_cache = self.kv_cache
349
+ mla_kv_cache = None if hisparse_cache is not None else self.kv_cache
350
mla_k_scale = self._k_scale
351
338
- kv_c_out = torch.empty_like(kv_c)
339
- k_pe_out = torch.empty_like(k_pe)
352
+ if self.use_pcp or hisparse_cache is None:
353
+ kv_c_out = torch.empty_like(kv_c)
354
+ k_pe_out = torch.empty_like(k_pe)
355
+ else:
356
+ kv_c_out = torch.empty_like(kv_c)
357
+ k_pe_out = torch.empty_like(k_pe)
358
q_c = fused_norm_rope(
359
positions,
360
q_c,
@@ -354,6 +372,7 @@ class DeepseekV32Attention(MLAAttention):
372
indexer_k_rope_cos_sin_cache,
373
self.topk_indices_buffer,
374
slot_mapping=mla_slot,
375
+ indexer_slot_mapping=indexer_slot,
376
indexer_k_cache=indexer_k_cache,
377
mla_kv_cache=mla_kv_cache,
378
mla_kv_cache_dtype=self.kv_cache_dtype,
@@ -365,6 +384,18 @@ class DeepseekV32Attention(MLAAttention):
384
index_k_out=index_k_out,
385
)
386
387
+ if hisparse_cache is not None and mla_slot is not None:
388
+ assert kv_c_out is not None and k_pe_out is not None
389
+ self.update_kv_cache(
390
+ kv_c_out,
391
+ k_pe_out,
392
+ self.kv_cache,
393
+ mla_slot,
394
+ layer_attn_metadata,
395
+ self.kv_cache_dtype,
396
+ self._k_scale,
397
+ )
398
+
399
q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim)
400
q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
401
ql_nope = torch.bmm(q_nope.transpose(0, 1), self.W_UK_T).transpose(0, 1)
@@ -459,6 +490,7 @@ class DeepseekV32Attention(MLAAttention):
490
),
491
skip_topk_buffer_clear=True,
492
)
493
+ self.impl.record_logical_topk_ready() # type: ignore[attr-defined]
494
495
attn_metadata, _, kv_cache, layer_slot_mapping = get_attention_context(
496
self.layer_name
@@ -467,6 +499,7 @@ class DeepseekV32Attention(MLAAttention):
499
output.zero_()
500
return
501
attn_metadata = cast("MLACommonMetadata", attn_metadata)
502
+ self.impl.prepare_for_batch(attn_metadata)
503
504
if self.use_pcp:
505
assert kv_c is not None and k_pe is not None
vllm/models/deepseek_v32/common/kernels.py
+12
-7
@@ -68,7 +68,7 @@ def _fp8_quant_and_cache_write(
68
kv_cache_ptr,
69
kv_cache_scale_ptr,
70
cache_block_size,
71
- cache_stride,
71
+ cache_block_stride,
72
offsets,
73
HEAD_DIM: tl.constexpr,
74
):
@@ -76,7 +76,7 @@ def _fp8_quant_and_cache_write(
76
77
block_idx = slot_idx // cache_block_size
78
block_offset = slot_idx % cache_block_size
79
- block_start = block_idx * cache_block_size * cache_stride
79
+ block_start = block_idx * cache_block_stride
80
81
tl.store(
82
kv_cache_ptr + block_start + block_offset * HEAD_DIM + offsets,
@@ -129,8 +129,9 @@ def _fused_norm_rope_kernel(
129
index_k_out_ptr,
130
index_k_out_stride,
131
INDEX_K_HALF_ROT_DIM: tl.constexpr,
132
- # Cache params (shared by indexer K and MLA)
132
+ # Cache params
133
slot_mapping_ptr,
134
+ indexer_slot_mapping_ptr,
135
# Index K FP8 cache
136
indexer_cache_ptr,
137
indexer_cache_scale_ptr,
@@ -400,8 +401,8 @@ def _fused_norm_rope_kernel(
401
)
402
403
# PCP inserts index K after gathering; other paths write it directly.
403
- if indexer_cache_ptr is not None and slot_mapping_ptr is not None:
404
- slot_idx = tl.load(slot_mapping_ptr + tok_idx)
404
+ if indexer_cache_ptr is not None and indexer_slot_mapping_ptr is not None:
405
+ slot_idx = tl.load(indexer_slot_mapping_ptr + tok_idx)
406
_fp8_quant_and_cache_write(
407
result,
408
index_k_mask,
@@ -431,8 +432,9 @@ def fused_norm_rope(
432
index_k_layer_norm_eps: float,
433
index_k_rope_cos_sin_cache: torch.Tensor | None,
434
topk_indices_buffer: torch.Tensor,
434
- # Cache params for fused writes (single slot_mapping for both caches)
435
+ # Cache params for fused writes
436
slot_mapping: torch.Tensor | None = None,
437
+ indexer_slot_mapping: torch.Tensor | None = None,
438
indexer_k_cache: torch.Tensor | None = None,
439
mla_kv_cache: torch.Tensor | None = None,
440
mla_kv_cache_dtype: str = "auto",
@@ -476,9 +478,11 @@ def fused_norm_rope(
478
# --- Indexer K cache setup ---
479
if indexer_k_cache is not None:
480
assert slot_mapping is not None
481
+ if indexer_slot_mapping is None:
482
+ indexer_slot_mapping = slot_mapping
483
idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32)
484
idx_cache_block_size = indexer_k_cache.shape[1]
481
- idx_cache_stride = indexer_k_cache.shape[2]
485
+ idx_cache_stride = indexer_k_cache.stride(0)
486
if indexer_k_cache.dtype == torch.uint8:
487
indexer_k_cache = indexer_k_cache.view(torch.float8_e4m3fn)
488
else:
@@ -578,6 +582,7 @@ def fused_norm_rope(
582
index_k_rope_cos_sin_cache.shape[-1] // 2,
583
# Cache params
584
slot_mapping,
585
+ indexer_slot_mapping,
586
indexer_k_cache,
587
idx_cache_scale_view,
588
idx_cache_block_size,
vllm/models/deepseek_v32/nvidia/model.py
+11
@@ -46,6 +46,10 @@ from vllm.models.common.ops.sequence_parallel import (
46
)
47
from vllm.models.deepseek_v32.attention import DeepseekV32Attention
48
from vllm.sequence import IntermediateTensors
49
+from vllm.v1.attention.backends.mla.index_group import (
50
+ SparseMLAIndexGroupBuilder,
51
+ get_sparse_mla_index_group_max_rows,
52
+)
53
54
from .glm52_low_latency_gemm import enable_glm52_low_latency_gemm
55
@@ -57,6 +61,7 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
61
prefix: str,
62
config=None,
63
topk_indices_buffer: torch.Tensor | None = None,
64
+ index_group_builder: SparseMLAIndexGroupBuilder | None = None,
65
) -> None:
66
super().__init__()
67
@@ -80,6 +85,7 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
85
config=config,
86
prefix=f"{prefix}.self_attn",
87
topk_indices_buffer=topk_indices_buffer,
88
+ index_group_builder=index_group_builder,
89
)
90
91
if (
@@ -189,6 +195,10 @@ class DeepseekV32Model(torch.nn.Module):
195
dtype=torch.int32,
196
device=self.device,
197
)
198
+ index_group_builder = SparseMLAIndexGroupBuilder(
199
+ self.topk_indices_buffer,
200
+ get_sparse_mla_index_group_max_rows(vllm_config),
201
+ )
202
203
if get_pp_group().is_first_rank:
204
self.embed_tokens = make_input_embedding(
@@ -209,6 +219,7 @@ class DeepseekV32Model(torch.nn.Module):
219
vllm_config=vllm_config,
220
prefix=prefix,
221
topk_indices_buffer=self.topk_indices_buffer,
222
+ index_group_builder=index_group_builder,
223
),
224
prefix=f"{prefix}.layers",
225
)
vllm/models/deepseek_v32/nvidia/mtp.py
+9
@@ -48,6 +48,10 @@ from vllm.models.common.ops.sequence_parallel import (
48
from vllm.models.deepseek_v32.common.kernels import fused_eh_norm
49
from vllm.platforms import current_platform
50
from vllm.sequence import IntermediateTensors
51
+from vllm.v1.attention.backends.mla.index_group import (
52
+ SparseMLAIndexGroupBuilder,
53
+ get_sparse_mla_index_group_max_rows,
54
+)
55
56
from .glm52_low_latency_gemm import (
57
build_glm52_plan,
@@ -80,6 +84,10 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
84
dtype=torch.int32,
85
device=current_platform.device_type,
86
)
87
+ index_group_builder = SparseMLAIndexGroupBuilder(
88
+ topk_indices_buffer,
89
+ get_sparse_mla_index_group_max_rows(vllm_config),
90
+ )
91
self.shared_head = SharedHead(
92
config=config, prefix=prefix, quant_config=quant_config
93
)
@@ -88,6 +96,7 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
96
prefix,
97
config=config,
98
topk_indices_buffer=topk_indices_buffer,
99
+ index_group_builder=index_group_builder,
100
)
101
102
def forward(
vllm/models/minimax_m3/common/indexer.py
+2
@@ -55,6 +55,7 @@ from vllm.v1.kv_cache_interface import (
55
AttentionSpec,
56
KVCacheSpec,
57
MLAAttentionSpec,
58
+ SparseCacheRole,
59
)
60
61
logger = init_logger(__name__)
@@ -148,6 +149,7 @@ class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase):
149
num_kv_heads=1,
150
head_size=self.head_dim,
151
dtype=self.dtype,
152
+ cache_role=SparseCacheRole.INDEXER,
153
)
154
155
def forward(self) -> None: ...
vllm/v1/attention/backend.py
+5
@@ -22,6 +22,7 @@ if TYPE_CHECKING:
22
from vllm.model_executor.layers.linear import ColumnParallelLinear
23
from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
24
from vllm.platforms.interface import DeviceCapability
25
+ from vllm.v1.hisparse.runtime import HiSparseCacheHandle
26
from vllm.v1.kv_cache_interface import (
27
AttentionSpec,
28
KVCacheLayout,
@@ -838,6 +839,9 @@ class AttentionImplBase(ABC, Generic[T]):
839
def process_weights_after_loading(self, act_dtype: torch.dtype):
840
pass
841
842
+ def prepare_for_batch(self, attn_metadata: T | None) -> None:
843
+ """Prepare implementation-specific state for the current batch."""
844
+
845
846
class AttentionImpl(AttentionImplBase[T], Generic[T]):
847
"""Standard attention implementation with forward method."""
@@ -956,6 +960,7 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]):
960
class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):
961
"""MLA attention implementation with forward_mqa and forward_mha methods."""
962
963
+ hisparse_cache: "HiSparseCacheHandle | None" = None
964
supports_pcp: bool = True
965
966
@abstractmethod
vllm/v1/attention/backends/mla/flashattn_mla_sparse.py
+103
-24
@@ -8,9 +8,9 @@ import torch
8
9
from vllm.config import VllmConfig
10
from vllm.config.cache import CacheDType
11
-from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata
11
from vllm.model_executor.layers.attention.sparse_mla_attention import (
12
SparseMLACommonImpl,
13
+ SparseMLACommonMetadata,
14
SparseMLACommonMetadataBuilder,
15
)
16
from vllm.platforms.interface import DeviceCapability
@@ -18,11 +18,11 @@ from vllm.v1.attention.backend import (
18
AttentionBackend,
19
AttentionCGSupport,
20
AttentionLayer,
21
- AttentionMetadata,
21
MLAAttentionImpl,
22
MultipleOf,
23
)
24
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
25
+from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
26
from vllm.v1.attention.backends.mla.sparse_utils import (
27
flat_kv_row_view,
28
triton_convert_req_index_to_global_index,
@@ -106,26 +106,8 @@ class FlashAttnMLASparseBackend(AttentionBackend):
106
107
108
@dataclass
109
-class FlashAttnMLASparseMetadata(AttentionMetadata):
110
- num_reqs: int
111
- max_query_len: int
112
- max_seq_len: int
113
-
114
- num_actual_tokens: int
115
- query_start_loc: torch.Tensor
116
- slot_mapping: torch.Tensor
117
-
118
- block_table: torch.Tensor
119
- req_id_per_token: torch.Tensor
120
- seq_lens: torch.Tensor
121
- block_size: int = 64
122
- topk_tokens: int = 2048
123
- num_decodes: int = 0
124
- num_prefills: int = 0
125
- num_decode_tokens: int = 0
126
- prefill_max_seq_len: int = 0
127
- prefill: MLACommonPrefillMetadata | None = None
128
- cp_kv_cache_interleave_size: int = 1
109
+class FlashAttnMLASparseMetadata(SparseMLACommonMetadata):
110
+ pass
111
112
113
class FlashAttnMLASparseMetadataBuilder(
@@ -214,6 +196,76 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
196
197
assert self.topk_indices_buffer is not None
198
topk_indices = self.topk_indices_buffer[:num_actual_toks]
199
+ index_group = self.index_group
200
+ if isinstance(index_group, HiSparseMLAIndexGroup):
201
+ num_decode_tokens = attn_metadata.num_decode_tokens
202
+ outputs = []
203
+ if num_decode_tokens:
204
+ physical_topk, valid_counts = (
205
+ index_group.convert_decode_logical_to_physical_topk(
206
+ self.index_group_index,
207
+ topk_indices[:num_decode_tokens],
208
+ attn_metadata,
209
+ return_valid_counts=True,
210
+ )
211
+ )
212
+ outputs.append(
213
+ self._run_mqa_kernel(
214
+ q_nope[:num_decode_tokens],
215
+ q_rope[:num_decode_tokens],
216
+ index_group.physical_kv_cache(self.index_group_index).view(
217
+ kv_c_and_k_pe_cache.dtype
218
+ ),
219
+ physical_topk,
220
+ valid_counts,
221
+ attn_metadata.block_size,
222
+ )
223
+ )
224
+ if num_decode_tokens < num_actual_toks:
225
+ cache = index_group.cache(self.index_group_index)
226
+ if num_decode_tokens == 0 and cache.all_context_pages_resident:
227
+ physical_topk, valid_counts = (
228
+ index_group.convert_logical_to_physical_topk(
229
+ self.index_group_index,
230
+ topk_indices,
231
+ attn_metadata,
232
+ block_stride_rows=None,
233
+ return_valid_counts=True,
234
+ )
235
+ )
236
+ prefill_cache = index_group.physical_kv_cache(
237
+ self.index_group_index
238
+ ).view(kv_c_and_k_pe_cache.dtype)
239
+ else:
240
+ prefill_cache, block_table, req_ids = (
241
+ index_group.stage_prefill_rows(
242
+ self.index_group_index,
243
+ kv_c_and_k_pe_cache,
244
+ attn_metadata,
245
+ )
246
+ )
247
+ physical_topk, valid_counts = (
248
+ triton_convert_req_index_to_global_index(
249
+ req_ids,
250
+ block_table,
251
+ topk_indices[num_decode_tokens:],
252
+ BLOCK_SIZE=attn_metadata.block_size,
253
+ NUM_TOPK_TOKENS=topk_indices.shape[1],
254
+ return_valid_counts=True,
255
+ )
256
+ )
257
+ outputs.append(
258
+ self._run_mqa_kernel(
259
+ q_nope[num_decode_tokens:],
260
+ q_rope[num_decode_tokens:],
261
+ prefill_cache,
262
+ physical_topk,
263
+ valid_counts,
264
+ attn_metadata.block_size,
265
+ )
266
+ )
267
+ return torch.cat(outputs) if len(outputs) > 1 else outputs[0], None
268
+
269
kv_rows, block_stride_rows = flat_kv_row_view(
270
kv_c_and_k_pe_cache, attn_metadata.block_size
271
)
@@ -226,9 +278,36 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
278
NUM_TOPK_TOKENS=topk_indices.shape[1],
279
return_valid_counts=True,
280
)
281
+ return (
282
+ self._run_mqa_kernel(
283
+ q_nope,
284
+ q_rope,
285
+ kv_rows,
286
+ topk_indices,
287
+ valid_counts,
288
+ attn_metadata.block_size,
289
+ cache_is_flat=True,
290
+ ),
291
+ None,
292
+ )
293
+
294
+ def _run_mqa_kernel(
295
+ self,
296
+ q_nope: torch.Tensor,
297
+ q_rope: torch.Tensor,
298
+ kv_cache: torch.Tensor,
299
+ topk_indices: torch.Tensor,
300
+ valid_counts: torch.Tensor,
301
+ block_size: int,
302
+ *,
303
+ cache_is_flat: bool = False,
304
+ ) -> torch.Tensor:
305
+ kv_rows = (
306
+ kv_cache if cache_is_flat else flat_kv_row_view(kv_cache, block_size)[0]
307
+ )
308
309
cu_seqlens_q = torch.arange(
231
- 0, num_actual_toks + 1, dtype=torch.int32, device=q_rope.device
310
+ 0, q_rope.shape[0] + 1, dtype=torch.int32, device=q_rope.device
311
)
312
k_cache = kv_rows[:, self.kv_lora_rank :].unsqueeze(1).unsqueeze(1)
313
v_cache = kv_rows[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1)
@@ -247,4 +326,4 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
326
causal=True,
327
fa_version=3,
328
)
250
- return out, None
329
+ return out
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+154
-41
@@ -11,22 +11,23 @@ from vllm import envs
11
from vllm.config import VllmConfig
12
from vllm.config.cache import CacheDType
13
from vllm.logger import init_logger
14
-from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata
14
from vllm.model_executor.layers.attention.sparse_mla_attention import (
15
SparseMLACommonImpl,
16
+ SparseMLACommonMetadata,
17
SparseMLACommonMetadataBuilder,
18
)
19
+from vllm.platforms import current_platform
20
from vllm.platforms.interface import DeviceCapability
21
from vllm.utils.torch_utils import is_quantized_kv_cache
22
from vllm.v1.attention.backend import (
23
AttentionBackend,
24
AttentionCGSupport,
25
AttentionLayer,
25
- AttentionMetadata,
26
AttentionType,
27
MLAAttentionImpl,
28
MultipleOf,
29
)
30
+from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
31
from vllm.v1.attention.backends.mla.sparse_utils import (
32
flat_kv_row_view,
33
triton_convert_req_index_to_global_index,
@@ -224,33 +225,9 @@ class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase):
225
226
227
@dataclass
227
-class FlashInferMLASparseMetadata(AttentionMetadata):
228
+class FlashInferMLASparseMetadata(SparseMLACommonMetadata):
229
"""Attention metadata for FlashInfer MLA Sparse backend."""
230
230
- num_reqs: int
231
- max_query_len: int
232
- max_seq_len: int
233
- num_actual_tokens: int
234
-
235
- # Query start locations
236
- query_start_loc: torch.Tensor
237
- slot_mapping: torch.Tensor
238
- block_table: torch.Tensor
239
- req_id_per_token: torch.Tensor
240
-
241
- # Sequence lengths for all requests (context + query)
242
- seq_lens: torch.Tensor
243
- num_decodes: int
244
- num_prefills: int
245
- num_decode_tokens: int
246
- prefill_max_seq_len: int = 0
247
- prefill: MLACommonPrefillMetadata | None = None
248
-
249
- # Sparse-specific
250
- block_size: int = 64
251
- topk_tokens: int = 2048
252
- cp_kv_cache_interleave_size: int = 1
253
-
231
232
class FlashInferMLASparseMetadataBuilder(
233
SparseMLACommonMetadataBuilder[FlashInferMLASparseMetadata]
@@ -286,6 +263,7 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde
263
"""Metadata builder for the SM100 TRT-LLM sparse MLA kernel."""
264
265
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
266
+ hisparse_supports_multi_token_decode: ClassVar[bool] = True
267
268
def _build_req_id_per_token(
269
self,
@@ -402,6 +380,78 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
380
assert self.topk_indices_buffer is not None
381
topk_indices = self.topk_indices_buffer[:num_actual_toks]
382
383
+ self._prepare_mqa_kernel(layer, q.device)
384
+
385
+ index_group = self.index_group
386
+ if isinstance(index_group, HiSparseMLAIndexGroup):
387
+ num_decode_tokens = attn_metadata.num_decode_tokens
388
+ decode_out: torch.Tensor | None = None
389
+ decode_lse: torch.Tensor | None = None
390
+ if num_decode_tokens > 0:
391
+ physical_topk, valid_counts = (
392
+ index_group.convert_decode_logical_to_physical_topk(
393
+ self.index_group_index,
394
+ topk_indices[:num_decode_tokens],
395
+ attn_metadata,
396
+ return_valid_counts=True,
397
+ )
398
+ )
399
+ decode_out, decode_lse = self._run_mqa_kernel(
400
+ q[:num_decode_tokens],
401
+ index_group.physical_kv_cache(self.index_group_index).view(
402
+ kv_c_and_k_pe_cache.dtype
403
+ ),
404
+ physical_topk,
405
+ valid_counts,
406
+ )
407
+ if num_decode_tokens == num_actual_toks:
408
+ return decode_out, decode_lse
409
+
410
+ cache = index_group.cache(self.index_group_index)
411
+ if num_decode_tokens == 0 and cache.all_context_pages_resident:
412
+ physical_topk, valid_counts = (
413
+ index_group.convert_logical_to_physical_topk(
414
+ self.index_group_index,
415
+ topk_indices,
416
+ attn_metadata,
417
+ block_stride_rows=None,
418
+ return_valid_counts=True,
419
+ )
420
+ )
421
+ return self._run_mqa_kernel(
422
+ q,
423
+ index_group.physical_kv_cache(self.index_group_index).view(
424
+ kv_c_and_k_pe_cache.dtype
425
+ ),
426
+ physical_topk,
427
+ valid_counts,
428
+ )
429
+
430
+ prefill_cache, block_table, req_ids = index_group.stage_prefill_rows(
431
+ self.index_group_index, kv_c_and_k_pe_cache, attn_metadata
432
+ )
433
+ prefill_indices, prefill_lens = triton_convert_req_index_to_global_index(
434
+ req_ids,
435
+ block_table,
436
+ topk_indices[num_decode_tokens:],
437
+ BLOCK_SIZE=attn_metadata.block_size,
438
+ NUM_TOPK_TOKENS=topk_indices.shape[1],
439
+ return_valid_counts=True,
440
+ )
441
+ prefill_out, prefill_lse = self._run_mqa_kernel(
442
+ q[num_decode_tokens:],
443
+ prefill_cache,
444
+ prefill_indices,
445
+ prefill_lens,
446
+ )
447
+ if decode_out is None:
448
+ return prefill_out, prefill_lse
449
+ output = torch.cat((decode_out, prefill_out))
450
+ if decode_lse is None:
451
+ return output, None
452
+ assert prefill_lse is not None
453
+ return output, torch.cat((decode_lse, prefill_lse))
454
+
455
_, block_stride_rows = flat_kv_row_view(
456
kv_c_and_k_pe_cache, attn_metadata.block_size
457
)
@@ -420,18 +470,27 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
470
return_valid_counts=True,
471
)
472
else:
423
- topk_indices_physical, seq_lens = triton_convert_req_index_to_global_index(
424
- attn_metadata.req_id_per_token[:num_actual_toks],
425
- attn_metadata.block_table,
473
+ topk_indices_physical, seq_lens = self._convert_logical_to_physical_topk(
474
topk_indices,
427
- BLOCK_SIZE=attn_metadata.block_size,
428
- BLOCK_STRIDE_ROWS=block_stride_rows,
429
- NUM_TOPK_TOKENS=topk_indices.shape[1],
475
+ attn_metadata,
476
+ block_stride_rows=block_stride_rows,
477
return_valid_counts=True,
478
)
479
480
+ return self._run_mqa_kernel(
481
+ q,
482
+ kv_c_and_k_pe_cache,
483
+ topk_indices_physical,
484
+ seq_lens,
485
+ )
486
+
487
+ def _prepare_mqa_kernel(
488
+ self,
489
+ layer: AttentionLayer,
490
+ device: torch.device,
491
+ ) -> None:
492
if self._workspace_buffer is None:
434
- self._workspace_buffer = _get_workspace_buffer(q.device)
493
+ self._workspace_buffer = _get_workspace_buffer(device)
494
495
if self.bmm1_scale is None:
496
self.bmm1_scale = self.scale
@@ -442,15 +501,71 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
501
if is_quantized_kv_cache(self.kv_cache_dtype):
502
self.bmm2_scale *= layer._k_scale_float
503
504
+ def autotune_hisparse_decode(self, layer: AttentionLayer) -> None:
505
+ """Autotune the largest legal HiSparse decode batch."""
506
+ assert isinstance(self.index_group, HiSparseMLAIndexGroup)
507
+ cache = self.index_group.cache(self.index_group_index)
508
+ assert self.topk_indices_buffer is not None
509
+
510
+ runtime = cache.runtime
511
+ kv_cache = runtime.hot.attention_cache
512
+ num_tokens = runtime.max_num_reqs
513
+ topk_tokens = self.topk_indices_buffer.shape[1]
514
+ self._prepare_mqa_kernel(layer, kv_cache.device)
515
+
516
+ q_dtype = (
517
+ current_platform.fp8_dtype()
518
+ if is_quantized_kv_cache(self.kv_cache_dtype)
519
+ else kv_cache.dtype
520
+ )
521
+ q = torch.zeros(
522
+ (
523
+ num_tokens,
524
+ self.num_heads,
525
+ self.kv_lora_rank + self.qk_rope_head_dim,
526
+ ),
527
+ dtype=q_dtype,
528
+ device=kv_cache.device,
529
+ )
530
+ topk_indices = (
531
+ torch.arange(
532
+ topk_tokens,
533
+ dtype=torch.int32,
534
+ device=kv_cache.device,
535
+ )
536
+ .expand(num_tokens, -1)
537
+ .contiguous()
538
+ )
539
+ seq_lens = torch.full(
540
+ (num_tokens,),
541
+ topk_tokens,
542
+ dtype=torch.int32,
543
+ device=kv_cache.device,
544
+ )
545
+ self._run_mqa_kernel(q, kv_cache, topk_indices, seq_lens)
546
+
547
+ def _run_mqa_kernel(
548
+ self,
549
+ q: torch.Tensor,
550
+ kv_cache: torch.Tensor,
551
+ topk_indices: torch.Tensor,
552
+ seq_lens: torch.Tensor,
553
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
554
+ assert self._workspace_buffer is not None
555
+ assert self.bmm1_scale is not None
556
+ assert self.bmm2_scale is not None
557
+
558
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
559
560
+ kv_cache = kv_cache.view(q.dtype)
561
+
562
# Single-token sparse decode. trtllm-gen requires the q_len_per_request
563
# dim, but the sparse attention mask is fully per-token (each query token
564
# carries its own top-k index row), so unsqueeze is sufficient and
565
# correct. The MTP/multi-token q_len grouping is a perf-only layout and is
566
# deferred until MTP is validated end-to-end for this backend.
567
query = q.unsqueeze(1)
453
- block_tables = topk_indices_physical.unsqueeze(1)
568
+ block_tables = topk_indices.unsqueeze(1)
569
seq_lens_arg = seq_lens
570
571
# page_table width = topk buffer width, which kpool widens past
@@ -460,7 +575,7 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
575
# count), so the -1 padding slots past seq_lens are never attended to.
576
# Use the actual buffer width instead of the fixed topk_tokens, which
577
# mismatches the page_table when index_kpool > 1.
463
- sparse_topk_capacity = topk_indices_physical.shape[1]
578
+ sparse_topk_capacity = topk_indices.shape[1]
579
580
extra_kwargs: dict[str, torch.Tensor] = {}
581
empty_rows: torch.Tensor | None = None
@@ -473,14 +588,12 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
588
# valid indices into a contiguous prefix, which is what the kernel
589
# requires of the page table.
590
empty_rows = seq_lens == 0
476
- topk_indices_physical[:, 0] = topk_indices_physical[:, 0].masked_fill(
477
- empty_rows, 0
478
- )
591
+ topk_indices[:, 0] = topk_indices[:, 0].masked_fill(empty_rows, 0)
592
extra_kwargs["sparse_mla_top_k_lens"] = seq_lens.clamp(min=1)
593
594
kernel_out = trtllm_batch_decode_with_kv_cache_mla(
595
query=query,
483
- kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
596
+ kv_cache=kv_cache.unsqueeze(1),
597
workspace_buffer=self._workspace_buffer,
598
qk_nope_head_dim=self.qk_nope_head_dim,
599
kv_lora_rank=self.kv_lora_rank,
@@ -506,7 +619,7 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
619
if lse is not None:
620
lse = self._normalize_lse(lse, out.shape[0], out.shape[1])
621
if empty_rows is None and lse is not None:
509
- empty_rows = (topk_indices_physical == -1).all(dim=-1)
622
+ empty_rows = (topk_indices == -1).all(dim=-1)
623
if empty_rows is not None:
624
out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0)
625
if lse is not None:
vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py
+106
-25
@@ -7,17 +7,14 @@ from typing import TYPE_CHECKING, cast
7
import torch
8
9
from vllm.model_executor.layers.attention.sparse_mla_attention import (
10
- SharedTopkIndicesBuffer,
11
-)
12
-from vllm.v1.attention.backend import (
13
- AttentionLayer,
14
- AttentionType,
15
- MLAAttentionImpl,
10
+ SparseMLACommonImpl,
11
)
12
+from vllm.v1.attention.backend import AttentionLayer, AttentionType
13
from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
14
FlashInferMLASparseMetadata,
15
_get_workspace_buffer,
16
)
17
+from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
18
from vllm.v1.attention.backends.mla.sparse_utils import (
19
triton_convert_req_index_to_global_index,
20
)
@@ -32,9 +29,7 @@ def _kv_scale_format_for_model(model_type: str | None) -> str:
29
return "pow2_fp32"
30
31
35
-class FlashInferMLASparseSM120Impl(
36
- MLAAttentionImpl[FlashInferMLASparseMetadata], SharedTopkIndicesBuffer
37
-):
32
+class FlashInferMLASparseSM120Impl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
33
"""SM120 FlashInfer sparse-MLA implementation."""
34
35
is_sparse = True
@@ -65,20 +60,28 @@ class FlashInferMLASparseSM120Impl(
60
"FLASHINFER_MLA_SPARSE_SM120 only supports decoder self-attention"
61
)
62
68
- self.num_heads = num_heads
69
- self.head_size = head_size
70
- self.scale = float(scale)
71
- self.num_kv_heads = num_kv_heads
72
- self.kv_cache_dtype = kv_cache_dtype
73
- if self.kv_cache_dtype != "fp8_ds_mla":
63
+ if kv_cache_dtype != "fp8_ds_mla":
64
raise NotImplementedError(
65
"FLASHINFER_MLA_SPARSE_SM120 requires the packed fp8_ds_mla "
66
f"KV cache layout; got kv_cache_dtype={kv_cache_dtype!r}."
67
)
68
79
- self.kv_lora_rank: int = mla_args["kv_lora_rank"]
80
- self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"]
81
- self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]
69
+ topk_indices_buffer = mla_args.pop("topk_indices_buffer", None)
70
+ super().__init__(
71
+ num_heads,
72
+ head_size,
73
+ scale,
74
+ num_kv_heads,
75
+ alibi_slopes,
76
+ sliding_window,
77
+ kv_cache_dtype,
78
+ logits_soft_cap,
79
+ attn_type,
80
+ kv_sharing_target_layer_name,
81
+ indexer=indexer,
82
+ topk_indices_buffer=topk_indices_buffer,
83
+ **mla_args,
84
+ )
85
from vllm.config import get_current_vllm_config
86
87
vllm_config = get_current_vllm_config()
@@ -89,9 +92,6 @@ class FlashInferMLASparseSM120Impl(
92
)
93
self.kv_scale_format = _kv_scale_format_for_model(model_type)
94
92
- # Skip-topk layers are built with indexer=None and get the shared
93
- # buffer via mla_args instead (cf. FLASHMLA_SPARSE).
94
- self.init_topk_indices_buffer(indexer, mla_args.get("topk_indices_buffer"))
95
from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
96
97
if not has_flashinfer_sparse_mla_sm120():
@@ -118,6 +118,70 @@ class FlashInferMLASparseSM120Impl(
118
119
assert self.topk_indices_buffer is not None
120
topk_indices = self.topk_indices_buffer[:num_actual_toks]
121
+ index_group = self.index_group
122
+ if isinstance(index_group, HiSparseMLAIndexGroup):
123
+ num_decode_tokens = attn_metadata.num_decode_tokens
124
+ outputs = []
125
+ if num_decode_tokens:
126
+ topk_indices_physical = cast(
127
+ torch.Tensor,
128
+ index_group.convert_decode_logical_to_physical_topk(
129
+ self.index_group_index,
130
+ topk_indices[:num_decode_tokens],
131
+ attn_metadata,
132
+ return_valid_counts=False,
133
+ ),
134
+ )
135
+ outputs.append(
136
+ self._run_mqa_kernel(
137
+ q[:num_decode_tokens],
138
+ index_group.physical_kv_cache(self.index_group_index),
139
+ topk_indices_physical,
140
+ )
141
+ )
142
+ if num_decode_tokens < num_actual_toks:
143
+ cache = index_group.cache(self.index_group_index)
144
+ if num_decode_tokens == 0 and cache.all_context_pages_resident:
145
+ topk_indices_physical = cast(
146
+ torch.Tensor,
147
+ index_group.convert_logical_to_physical_topk(
148
+ self.index_group_index,
149
+ topk_indices,
150
+ attn_metadata,
151
+ block_stride_rows=None,
152
+ return_valid_counts=False,
153
+ ),
154
+ )
155
+ prefill_cache = index_group.physical_kv_cache(
156
+ self.index_group_index
157
+ )
158
+ else:
159
+ prefill_cache, block_table, req_ids = (
160
+ index_group.stage_prefill_rows(
161
+ self.index_group_index,
162
+ kv_c_and_k_pe_cache,
163
+ attn_metadata,
164
+ )
165
+ )
166
+ topk_indices_physical = cast(
167
+ torch.Tensor,
168
+ triton_convert_req_index_to_global_index(
169
+ req_ids,
170
+ block_table,
171
+ topk_indices[num_decode_tokens:],
172
+ BLOCK_SIZE=attn_metadata.block_size,
173
+ NUM_TOPK_TOKENS=topk_indices.shape[1],
174
+ ),
175
+ )
176
+ outputs.append(
177
+ self._run_mqa_kernel(
178
+ q[num_decode_tokens:],
179
+ prefill_cache,
180
+ topk_indices_physical,
181
+ )
182
+ )
183
+ output = torch.cat(outputs) if len(outputs) > 1 else outputs[0]
184
+ return output, None
185
186
topk_indices_physical = cast(
187
torch.Tensor,
@@ -129,6 +193,22 @@ class FlashInferMLASparseSM120Impl(
193
NUM_TOPK_TOKENS=topk_indices.shape[1],
194
),
195
)
196
+ return (
197
+ self._run_mqa_kernel(
198
+ q,
199
+ kv_c_and_k_pe_cache,
200
+ topk_indices_physical,
201
+ ),
202
+ None,
203
+ )
204
+
205
+ def _run_mqa_kernel(
206
+ self,
207
+ q: torch.Tensor,
208
+ kv_cache: torch.Tensor,
209
+ topk_indices_physical: torch.Tensor,
210
+ ) -> torch.Tensor:
211
+ num_actual_toks = q.shape[0]
212
213
output = q.new_empty(
214
(num_actual_toks, self.num_heads, self.kv_lora_rank),
@@ -142,20 +222,21 @@ class FlashInferMLASparseSM120Impl(
222
flashinfer_trtllm_batch_decode_with_kv_cache_mla,
223
)
224
225
+ sparse_capacity = topk_indices_physical.shape[1]
226
out = flashinfer_trtllm_batch_decode_with_kv_cache_mla(
227
query=q.unsqueeze(1),
147
- kv_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(1),
228
+ kv_cache=kv_cache.view(torch.uint8).unsqueeze(1),
229
workspace_buffer=self._workspace_buffer,
230
qk_nope_head_dim=self.qk_nope_head_dim,
231
kv_lora_rank=self.kv_lora_rank,
232
qk_rope_head_dim=self.qk_rope_head_dim,
233
block_tables=topk_indices_physical.unsqueeze(1),
234
seq_lens=None,
154
- max_seq_len=attn_metadata.topk_tokens,
235
+ max_seq_len=sparse_capacity,
236
out=output.unsqueeze(1),
237
bmm1_scale=self.scale,
238
bmm2_scale=1.0,
158
- sparse_mla_top_k=attn_metadata.topk_tokens,
239
+ sparse_mla_top_k=sparse_capacity,
240
kv_scale_format=self.kv_scale_format,
241
)
161
- return out.squeeze(1), None
242
+ return out.squeeze(1)
vllm/v1/attention/backends/mla/flashmla_sparse.py
+225
-91
@@ -9,23 +9,23 @@ from vllm import _custom_ops as ops
9
from vllm.config import VllmConfig, get_current_vllm_config
10
from vllm.config.cache import CacheDType
11
from vllm.logger import init_logger
12
-from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata
12
from vllm.model_executor.layers.attention.sparse_mla_attention import (
13
SparseMLACommonImpl,
14
+ SparseMLACommonMetadata,
15
SparseMLACommonMetadataBuilder,
16
)
17
from vllm.platforms import current_platform
18
from vllm.platforms.interface import DeviceCapability
19
from vllm.utils.platform_utils import num_compute_units
20
-from vllm.utils.torch_utils import is_quantized_kv_cache
20
+from vllm.utils.torch_utils import current_stream, is_quantized_kv_cache
21
from vllm.v1.attention.backend import (
22
AttentionBackend,
23
AttentionCGSupport,
24
AttentionLayer,
25
- AttentionMetadata,
25
CommonAttentionMetadata,
26
MultipleOf,
27
)
28
+from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
29
from vllm.v1.attention.backends.mla.sparse_utils import (
30
flat_kv_row_view,
31
triton_convert_req_index_to_global_index,
@@ -172,28 +172,7 @@ class FlashMLASparseBackend(AttentionBackend):
172
173
174
@dataclass
175
-class FlashMLASparseMetadata(AttentionMetadata):
176
- num_reqs: int
177
- max_query_len: int
178
- max_seq_len: int
179
-
180
- num_actual_tokens: int # Number of tokens excluding padding.
181
- query_start_loc: torch.Tensor
182
- slot_mapping: torch.Tensor
183
-
184
- block_table: torch.Tensor
185
- req_id_per_token: torch.Tensor
186
- block_size: int = 64
187
- topk_tokens: int = 2048
188
-
189
- num_decodes: int = 0
190
- num_prefills: int = 0
191
- num_decode_tokens: int = 0
192
- seq_lens: torch.Tensor | None = None
193
- prefill_max_seq_len: int = 0
194
- prefill: MLACommonPrefillMetadata | None = None
195
- cp_kv_cache_interleave_size: int = 1
196
-
175
+class FlashMLASparseMetadata(SparseMLACommonMetadata):
176
@dataclass
177
class FP8KernelMetadata:
178
scheduler_metadata: FlashMLASchedMeta
@@ -233,6 +212,7 @@ class FlashMLASparseMetadata(AttentionMetadata):
212
req_start_idx: int
213
workspace_starts: torch.Tensor
214
chunk_tot_seqlen: int
215
+ seq_lens: torch.Tensor | None = None
216
217
chunks: list[Chunk]
218
@@ -263,6 +243,7 @@ class FlashMLASparseMetadataBuilder(
243
):
244
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
245
require_uniform_decodes: ClassVar[bool] = True
246
+ hisparse_supports_multi_token_decode: ClassVar[bool] = True
247
metadata_cls = FlashMLASparseMetadata
248
249
def __init__(
@@ -283,6 +264,9 @@ class FlashMLASparseMetadataBuilder(
264
)
265
else:
266
threshold = {16: 128, 32: 128, 64: 256, 128: 256}.get(num_q_heads, 256)
267
+ self.use_hisparse = vllm_config.attention_config.hisparse_config is not None
268
+ if self.use_hisparse:
269
+ threshold = 1
270
# Varlen decodes are safe under DCP: causality comes from the
271
# indexer's top-k indices, not from the kernel metadata.
272
self._init_reorder_batch_threshold(
@@ -347,7 +331,9 @@ class FlashMLASparseMetadataBuilder(
331
device=device,
332
)
333
350
- self.fp8_use_mixed_batch = self.num_heads < MIN_HEADS_FOR_BF16_PREFILL
334
+ self.fp8_use_mixed_batch = (
335
+ self.num_heads < MIN_HEADS_FOR_BF16_PREFILL and not self.use_hisparse
336
+ )
337
338
if parallel_config.decode_context_parallel_size > 1:
339
if parallel_config.dcp_comm_backend != "ag_rs":
@@ -385,37 +371,21 @@ class FlashMLASparseMetadataBuilder(
371
372
def _build_fp8_mixed_decode_prefill(
373
self,
388
- common_attn_metadata: CommonAttentionMetadata,
389
- ) -> "FlashMLASparseMetadata.FP8KernelMetadata":
374
+ ) -> FlashMLASparseMetadata.FP8KernelMetadata:
375
"""Build FP8 metadata treating MQA tokens as one batch.
376
377
The scheduler initializes lazily from the runtime query shape, which may
378
be the full batch or only decodes when prefills use dense MHA. This avoids
379
the BF16 prefill kernel's head-padding overhead at high TP.
380
"""
396
- num_tokens = common_attn_metadata.num_actual_tokens
397
-
398
- # Use padded head count since that's what the kernel will see
399
- padded_heads = self.fp8_decode_padded_heads
400
-
401
- # Build metadata for all tokens as a single batch
402
- scheduler_metadata, _ = get_mla_metadata(
403
- cache_seqlens=self.topk_tokens_tensor[:1], # Single batch
404
- num_q_tokens_per_head_k=num_tokens * padded_heads,
405
- topk=self.topk_tokens,
406
- num_heads_q=padded_heads,
407
- num_heads_k=1,
408
- is_fp8_kvcache=True,
409
- )
381
411
- fp8_metadata = FlashMLASparseMetadata.FP8KernelMetadata(
382
+ scheduler_metadata, _ = get_mla_metadata()
383
+ return FlashMLASparseMetadata.FP8KernelMetadata(
384
scheduler_metadata=scheduler_metadata,
385
cache_lens=self.max_model_len_tensor[:1],
386
dummy_block_table=self.dummy_block_table[:1],
387
)
388
417
- return fp8_metadata
418
-
389
def _build_fp8_separate_prefill_decode(
390
self,
391
common_attn_metadata: CommonAttentionMetadata,
@@ -520,6 +490,9 @@ class FlashMLASparseMetadataBuilder(
490
chunk_block_table = common_attn_metadata.block_table_tensor[
491
num_decodes + chunk_start : num_decodes + chunk_end
492
]
493
+ chunk_seq_lens = common_attn_metadata.seq_lens[
494
+ num_decodes + chunk_start : num_decodes + chunk_end
495
+ ]
496
497
prefill_chunks.append(
498
FP8Meta.Prefill.Chunk(
@@ -528,6 +501,7 @@ class FlashMLASparseMetadataBuilder(
501
req_start_idx=chunk_start,
502
workspace_starts=chunk_workspace_starts,
503
chunk_tot_seqlen=chunk_tot_seqlen,
504
+ seq_lens=chunk_seq_lens,
505
)
506
)
507
@@ -569,9 +543,7 @@ class FlashMLASparseMetadataBuilder(
543
metadata.fp8_use_mixed_batch = self.fp8_use_mixed_batch
544
if self.use_fp8_kv_cache:
545
if self.fp8_use_mixed_batch:
572
- metadata.fp8_extra_metadata = self._build_fp8_mixed_decode_prefill(
573
- common_attn_metadata
574
- )
546
+ metadata.fp8_extra_metadata = self._build_fp8_mixed_decode_prefill()
547
else:
548
metadata.fp8_extra_metadata = self._build_fp8_separate_prefill_decode(
549
common_attn_metadata, metadata
@@ -679,30 +651,83 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
651
topk_indices: torch.Tensor,
652
attn_metadata: FlashMLASparseMetadata,
653
actual_num_heads: int,
682
- ) -> tuple[torch.Tensor, torch.Tensor]:
683
- # Convert per-request indices to global slots (decode) or workspace
684
- # offsets (prefill). req_id_per_token covers the whole batch; slice it
685
- # to the MQA tokens (q may exclude prefill tokens routed to dense MHA).
654
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
655
+ index_group = self.index_group
656
+ if isinstance(index_group, HiSparseMLAIndexGroup):
657
+ cache = index_group.cache(self.index_group_index)
658
+ else:
659
+ cache = None
660
+ block_table = attn_metadata.block_table
661
+ # req_id_per_token covers the whole batch; slice it to the MQA tokens
662
+ # (q may exclude prefill tokens routed to dense MHA).
663
+ req_id_per_token = attn_metadata.req_id_per_token[: topk_indices.shape[0]]
664
+ decode_out: torch.Tensor | None = None
665
+ if cache is not None:
666
+ assert isinstance(index_group, HiSparseMLAIndexGroup)
667
+ num_decode_tokens = attn_metadata.num_decode_tokens
668
+ if num_decode_tokens > 0:
669
+ decode_topk, decode_lengths = (
670
+ index_group.convert_decode_logical_to_physical_topk(
671
+ self.index_group_index,
672
+ topk_indices[:num_decode_tokens],
673
+ attn_metadata,
674
+ return_valid_counts=True,
675
+ )
676
+ )
677
+ decode_out, _ = self._bf16_flash_mla_kernel(
678
+ q[:num_decode_tokens],
679
+ index_group.physical_kv_cache(self.index_group_index),
680
+ decode_topk,
681
+ decode_lengths,
682
+ actual_num_heads,
683
+ )
684
+ if num_decode_tokens == q.shape[0]:
685
+ return decode_out, None
686
+ q = q[num_decode_tokens:]
687
+ topk_indices = topk_indices[num_decode_tokens:]
688
+ kv_c_and_k_pe_cache, block_table, req_id_per_token = (
689
+ index_group.stage_prefill_rows(
690
+ self.index_group_index, kv_c_and_k_pe_cache, attn_metadata
691
+ )
692
+ )
693
+ # Convert per-request indices to global slots (decode) or workspace offsets.
694
kv_rows, block_stride_rows = flat_kv_row_view(
695
kv_c_and_k_pe_cache, attn_metadata.block_size
696
)
689
- topk_indices, topk_length = triton_convert_req_index_to_global_index(
690
- attn_metadata.req_id_per_token[: topk_indices.shape[0]],
691
- attn_metadata.block_table,
692
- topk_indices,
693
- BLOCK_SIZE=attn_metadata.block_size,
694
- BLOCK_STRIDE_ROWS=block_stride_rows,
695
- NUM_TOPK_TOKENS=topk_indices.shape[1],
696
- return_valid_counts=True,
697
+ decode_only = (
698
+ attn_metadata.num_decode_tokens
699
+ == attn_metadata.num_actual_tokens
700
+ == topk_indices.shape[0]
701
)
702
+ uses_host_cache = isinstance(index_group, HiSparseMLAIndexGroup)
703
+ if not uses_host_cache and decode_only:
704
+ topk_indices, topk_length = self._convert_logical_to_physical_topk(
705
+ topk_indices,
706
+ attn_metadata,
707
+ block_stride_rows=block_stride_rows,
708
+ return_valid_counts=True,
709
+ )
710
+ else:
711
+ topk_indices, topk_length = triton_convert_req_index_to_global_index(
712
+ req_id_per_token,
713
+ block_table,
714
+ topk_indices,
715
+ BLOCK_SIZE=attn_metadata.block_size,
716
+ BLOCK_STRIDE_ROWS=block_stride_rows,
717
+ NUM_TOPK_TOKENS=topk_indices.shape[1],
718
+ return_valid_counts=True,
719
+ )
720
699
- return self._bf16_flash_mla_kernel(
721
+ attn_out, lse = self._bf16_flash_mla_kernel(
722
q,
723
kv_rows,
724
topk_indices,
725
topk_length,
726
actual_num_heads,
727
)
728
+ if decode_out is None:
729
+ return attn_out, lse
730
+ return torch.cat([decode_out, attn_out], dim=0), None
731
732
def _forward_fp8_kv_separate_prefill_decode(
733
self,
@@ -721,6 +746,28 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
746
"FP8 sparse MLA expects either the decode subset or the full batch"
747
)
748
749
+ decode_topk: torch.Tensor | None = None
750
+ index_group = self.index_group
751
+ uses_host_cache = isinstance(index_group, HiSparseMLAIndexGroup)
752
+ if uses_host_cache and num_decode_tokens > 0:
753
+ decode_topk = topk_indices[:num_decode_tokens]
754
+
755
+ prefill_ready = None
756
+ if num_prefill_tokens > 0 and uses_host_cache:
757
+ assert fp8_metadata.prefill is not None
758
+ first_chunk = fp8_metadata.prefill.chunks[0]
759
+ assert isinstance(index_group, HiSparseMLAIndexGroup)
760
+ prefill_ready = index_group.gather_fp8_prefill(
761
+ self.index_group_index,
762
+ kv_c_and_k_pe_cache,
763
+ self.prefill_bf16_workspace[: first_chunk.chunk_tot_seqlen],
764
+ first_chunk.block_table,
765
+ first_chunk.workspace_starts,
766
+ len(first_chunk.block_table),
767
+ attn_metadata,
768
+ first_chunk.req_start_idx,
769
+ )
770
+
771
prefill_request_ids = None
772
prefill_workspace_starts = None
773
has_prefill_workspace = False
@@ -736,17 +783,26 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
783
# For BF16 cache: always use global cache slots (no workspace)
784
# prefill_workspace_starts has been adjusted in-place per chunk so
785
# prefill indices automatically come out chunk-local
739
- topk_indices, topk_length = triton_convert_req_index_to_global_index(
740
- attn_metadata.req_id_per_token[: topk_indices.shape[0]],
741
- attn_metadata.block_table,
742
- topk_indices,
743
- BLOCK_SIZE=attn_metadata.block_size,
744
- NUM_TOPK_TOKENS=topk_indices.shape[1],
745
- HAS_PREFILL_WORKSPACE=has_prefill_workspace,
746
- prefill_workspace_request_ids=prefill_request_ids,
747
- prefill_workspace_starts=prefill_workspace_starts,
748
- return_valid_counts=True,
749
- )
786
+ topk_length = None
787
+ if num_prefill_tokens == 0 and not uses_host_cache:
788
+ topk_indices, topk_length = self._convert_logical_to_physical_topk(
789
+ topk_indices,
790
+ attn_metadata,
791
+ block_stride_rows=None,
792
+ return_valid_counts=True,
793
+ )
794
+ elif num_prefill_tokens > 0:
795
+ topk_indices, topk_length = triton_convert_req_index_to_global_index(
796
+ attn_metadata.req_id_per_token[: topk_indices.shape[0]],
797
+ attn_metadata.block_table,
798
+ topk_indices,
799
+ BLOCK_SIZE=attn_metadata.block_size,
800
+ NUM_TOPK_TOKENS=topk_indices.shape[1],
801
+ HAS_PREFILL_WORKSPACE=has_prefill_workspace,
802
+ prefill_workspace_request_ids=prefill_request_ids,
803
+ prefill_workspace_starts=prefill_workspace_starts,
804
+ return_valid_counts=True,
805
+ )
806
807
fp8_metadata = attn_metadata.fp8_extra_metadata
808
assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8SeparatePrefillDecode)
@@ -755,6 +811,16 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
811
q: torch.Tensor,
812
topk_indices: torch.Tensor,
813
) -> torch.Tensor:
814
+ assert fp8_metadata.decode is not None
815
+ if uses_host_cache:
816
+ return self._host_backed_fp8_decode(
817
+ q,
818
+ topk_indices,
819
+ attn_metadata,
820
+ fp8_metadata.decode.kernel_metadata,
821
+ num_decodes,
822
+ fp8_metadata.decode.decode_query_len,
823
+ )
824
# Reshape q: (num_decode_tokens, num_heads, head_dim)
825
# -> (num_decodes, seq_len, num_heads, head_dim)
826
q = reshape_query_for_spec_decode(q, num_decodes)
@@ -762,7 +828,6 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
828
# Reshape topk_indices: (num_decode_tokens, topk)
829
# -> (num_decodes, seq_len, topk)
830
topk_indices = topk_indices.view(num_decodes, seq_len, -1)
765
- assert fp8_metadata.decode is not None
831
attn_out, _ = self._fp8_flash_mla_kernel(
832
q=q,
833
kv_c_and_k_pe_cache=kv_c_and_k_pe_cache,
@@ -776,7 +841,9 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
841
# Pure decode: direct call without allocation
842
if num_decode_tokens > 0 and num_prefill_tokens == 0:
843
assert fp8_metadata.decode is not None
779
- attn_out = _fp8_decode(q, topk_indices)
844
+ attn_out = _fp8_decode(
845
+ q, decode_topk if decode_topk is not None else topk_indices
846
+ )
847
else:
848
# Mixed or pure prefill: allocate output tensor
849
attn_out = q.new_empty(
@@ -788,13 +855,30 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
855
if num_decode_tokens > 0:
856
attn_out[:num_decode_tokens] = _fp8_decode(
857
q[:num_decode_tokens],
791
- topk_indices[:num_decode_tokens],
858
+ decode_topk
859
+ if decode_topk is not None
860
+ else topk_indices[:num_decode_tokens],
861
)
862
863
assert fp8_metadata.prefill is not None
795
- for chunk in fp8_metadata.prefill.chunks:
864
+ for chunk_index, chunk in enumerate(fp8_metadata.prefill.chunks):
865
chunk_workspace = self.prefill_bf16_workspace[: chunk.chunk_tot_seqlen]
797
- if self.kv_cache_dtype == "fp8_ds_mla":
866
+ if uses_host_cache and chunk_index > 0:
867
+ assert isinstance(index_group, HiSparseMLAIndexGroup)
868
+ prefill_ready = index_group.gather_fp8_prefill(
869
+ self.index_group_index,
870
+ kv_c_and_k_pe_cache,
871
+ chunk_workspace,
872
+ chunk.block_table,
873
+ chunk.workspace_starts,
874
+ len(chunk.block_table),
875
+ attn_metadata,
876
+ chunk.req_start_idx,
877
+ )
878
+ if uses_host_cache:
879
+ assert prefill_ready is not None
880
+ current_stream().wait_event(prefill_ready)
881
+ elif self.kv_cache_dtype == "fp8_ds_mla":
882
ops.cp_gather_and_upconvert_fp8_kv_cache(
883
kv_c_and_k_pe_cache,
884
chunk_workspace,
@@ -813,6 +897,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
897
898
chunk_q = q[chunk.tokens_slice]
899
chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice]
900
+ assert topk_length is not None
901
chunk_topk_length = topk_length[chunk.tokens_slice]
902
903
attn_out[chunk.tokens_slice], _ = self._bf16_flash_mla_kernel(
@@ -839,6 +924,17 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
924
925
The lse is only returned when DCP needs it, otherwise None.
926
"""
927
+ assert attn_metadata.fp8_extra_metadata is not None
928
+ assert isinstance(
929
+ attn_metadata.fp8_extra_metadata,
930
+ FlashMLASparseMetadata.FP8KernelMetadata,
931
+ )
932
+ fp8_metadata = attn_metadata.fp8_extra_metadata
933
+
934
+ block_table = attn_metadata.block_table
935
+ # req_id_per_token covers the whole batch; slice it to the MQA tokens
936
+ # (q may exclude prefill tokens routed to dense MHA).
937
+ req_id_per_token = attn_metadata.req_id_per_token[: topk_indices.shape[0]]
938
if self.dcp_world_size > 1:
939
# The indexer emits global token ids; keep this rank's shard and
940
# convert to local slots. compact_valid_to_front=False keeps the
@@ -846,8 +942,8 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
942
# empty-row neutralization below relies on. req_id is sliced to
943
# topk_indices rows (the converter grids from req_id).
944
topk_indices = triton_filter_and_convert_dcp_index(
849
- attn_metadata.req_id_per_token[: topk_indices.shape[0]],
850
- attn_metadata.block_table,
945
+ req_id_per_token,
946
+ block_table,
947
topk_indices,
948
dcp_size=self.dcp_world_size,
949
dcp_rank=self.dcp_rank,
@@ -859,19 +955,26 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
955
else:
956
# Convert per-request indices to global slots (decode) or workspace
957
# offsets (prefill).
862
- topk_indices = triton_convert_req_index_to_global_index(
863
- attn_metadata.req_id_per_token[: topk_indices.shape[0]],
864
- attn_metadata.block_table,
865
- topk_indices,
866
- BLOCK_SIZE=attn_metadata.block_size,
867
- NUM_TOPK_TOKENS=topk_indices.shape[1],
958
+ decode_only = (
959
+ attn_metadata.num_decode_tokens
960
+ == attn_metadata.num_actual_tokens
961
+ == topk_indices.shape[0]
962
)
869
-
870
- assert attn_metadata.fp8_extra_metadata is not None
871
- assert isinstance(
872
- attn_metadata.fp8_extra_metadata, FlashMLASparseMetadata.FP8KernelMetadata
873
- )
874
- fp8_metadata = attn_metadata.fp8_extra_metadata
963
+ if decode_only:
964
+ topk_indices = self._convert_logical_to_physical_topk(
965
+ topk_indices,
966
+ attn_metadata,
967
+ block_stride_rows=None,
968
+ return_valid_counts=False,
969
+ )
970
+ else:
971
+ topk_indices = triton_convert_req_index_to_global_index(
972
+ req_id_per_token,
973
+ block_table,
974
+ topk_indices,
975
+ BLOCK_SIZE=attn_metadata.block_size,
976
+ NUM_TOPK_TOKENS=topk_indices.shape[1],
977
+ )
978
979
_attn_out, _lse = self._fp8_flash_mla_kernel(
980
q=q.unsqueeze(0), # unsqueeze to add batch_dim: (T, H, D) -> (1, T, H, D)
@@ -937,6 +1040,37 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1040
1041
return out, lse
1042
1043
+ def _host_backed_fp8_decode(
1044
+ self,
1045
+ q: torch.Tensor,
1046
+ topk_indices: torch.Tensor,
1047
+ attn_metadata: FlashMLASparseMetadata,
1048
+ kernel_metadata: FlashMLASparseMetadata.FP8KernelMetadata,
1049
+ num_decodes: int,
1050
+ decode_query_len: int,
1051
+ ) -> torch.Tensor:
1052
+ assert isinstance(self.index_group, HiSparseMLAIndexGroup)
1053
+ physical_topk = self.index_group.convert_decode_logical_to_physical_topk(
1054
+ self.index_group_index,
1055
+ topk_indices,
1056
+ attn_metadata,
1057
+ return_valid_counts=False,
1058
+ num_decodes=num_decodes,
1059
+ decode_query_len=decode_query_len,
1060
+ )
1061
+ assert isinstance(physical_topk, torch.Tensor)
1062
+ q = reshape_query_for_spec_decode(q, num_decodes)
1063
+ physical_topk = physical_topk.view(num_decodes, q.shape[1], -1)
1064
+ output, _ = self._fp8_flash_mla_kernel(
1065
+ q=q,
1066
+ kv_c_and_k_pe_cache=self.index_group.physical_kv_cache(
1067
+ self.index_group_index
1068
+ ),
1069
+ topk_indices=physical_topk,
1070
+ kernel_metadata=kernel_metadata,
1071
+ )
1072
+ return reshape_attn_output_for_spec_decode(output)
1073
+
1074
def _bf16_flash_mla_kernel(
1075
self,
1076
q: torch.Tensor,
vllm/v1/attention/backends/mla/index_group.py
new
+503
@@ -0,0 +1,503 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from __future__ import annotations
5
+
6
+from dataclasses import dataclass, field
7
+from typing import Any
8
+
9
+import torch
10
+
11
+from vllm import _custom_ops as ops
12
+from vllm.config import VllmConfig
13
+from vllm.utils.torch_utils import current_stream
14
+from vllm.v1.attention.backends.mla.sparse_utils import (
15
+ triton_convert_req_index_to_global_index,
16
+)
17
+from vllm.v1.hisparse.runtime import (
18
+ FP8_DS_MLA_ROW_BYTES,
19
+ HiSparseCacheHandle,
20
+ HiSparsePrefillStagingPlan,
21
+ create_hisparse_cache_handle,
22
+)
23
+
24
+
25
+def _create_side_stream(device: torch.device) -> torch.Stream:
26
+ return torch.Stream(device=device)
27
+
28
+
29
+def _create_event() -> torch.Event:
30
+ return torch.Event()
31
+
32
+
33
+@dataclass
34
+class SparseMLAIndexGroup:
35
+ """Layers that consume one sparse-indexer result."""
36
+
37
+ logical_topk_indices: torch.Tensor
38
+ physical_topk_indices: torch.Tensor
39
+ valid_topk_counts: torch.Tensor
40
+ request_ids: torch.Tensor
41
+ side_stream: torch.Stream
42
+ logical_topk_ready: torch.Event
43
+ physical_topk_ready: torch.Event
44
+ has_indexer: bool
45
+ num_layers: int = 0
46
+
47
+ def register_layer(
48
+ self,
49
+ vllm_config: VllmConfig | None = None,
50
+ *,
51
+ head_size: int | None = None,
52
+ kv_cache_dtype: str | None = None,
53
+ ) -> int:
54
+ layer_index = self.num_layers
55
+ self.num_layers += 1
56
+ return layer_index
57
+
58
+ def set_logical_topk_ready(self, layer_index: int) -> None:
59
+ if layer_index == 0 and self.has_indexer:
60
+ self.logical_topk_ready.record(current_stream())
61
+
62
+ def prepare_for_batch(self, layer_index: int, attn_metadata: Any | None) -> None:
63
+ pass
64
+
65
+ def gather_fp8_prefill(
66
+ self,
67
+ layer_index: int,
68
+ source_cache: torch.Tensor,
69
+ dst: torch.Tensor,
70
+ block_table: torch.Tensor,
71
+ workspace_starts: torch.Tensor,
72
+ batch_size: int,
73
+ attn_metadata: Any,
74
+ request_start: int,
75
+ ) -> torch.Event | None:
76
+ ops.cp_gather_and_upconvert_fp8_kv_cache(
77
+ source_cache,
78
+ dst,
79
+ block_table,
80
+ workspace_starts,
81
+ batch_size,
82
+ )
83
+ return None
84
+
85
+ def convert_logical_to_physical_topk(
86
+ self,
87
+ layer_index: int,
88
+ logical_topk_indices: torch.Tensor,
89
+ attn_metadata: Any,
90
+ *,
91
+ block_stride_rows: int | None,
92
+ return_valid_counts: bool,
93
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
94
+ """Convert one index result once for all layers in this group."""
95
+ return self._convert_once(
96
+ layer_index,
97
+ logical_topk_indices,
98
+ attn_metadata.req_id_per_token[: logical_topk_indices.shape[0]],
99
+ attn_metadata.block_table,
100
+ attn_metadata.block_size,
101
+ block_stride_rows=block_stride_rows,
102
+ return_valid_counts=return_valid_counts,
103
+ )
104
+
105
+ def _convert_once(
106
+ self,
107
+ layer_index: int,
108
+ logical_topk_indices: torch.Tensor,
109
+ req_id_per_token: torch.Tensor,
110
+ block_table: torch.Tensor,
111
+ block_size: int,
112
+ *,
113
+ block_stride_rows: int | None,
114
+ return_valid_counts: bool,
115
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
116
+ num_tokens = logical_topk_indices.shape[0]
117
+ if num_tokens > self.physical_topk_indices.shape[0]:
118
+ return triton_convert_req_index_to_global_index(
119
+ req_id_per_token,
120
+ block_table,
121
+ logical_topk_indices,
122
+ BLOCK_SIZE=block_size,
123
+ BLOCK_STRIDE_ROWS=block_stride_rows,
124
+ NUM_TOPK_TOKENS=logical_topk_indices.shape[1],
125
+ return_valid_counts=return_valid_counts,
126
+ )
127
+
128
+ physical_topk_indices = self.physical_topk_indices[:num_tokens]
129
+ valid_topk_counts = self.valid_topk_counts[:num_tokens]
130
+ compute_stream = current_stream()
131
+ if layer_index == 0:
132
+ if self.has_indexer:
133
+ self.side_stream.wait_event(self.logical_topk_ready)
134
+ else:
135
+ self.side_stream.wait_stream(compute_stream)
136
+ with self.side_stream:
137
+ triton_convert_req_index_to_global_index(
138
+ req_id_per_token,
139
+ block_table,
140
+ logical_topk_indices,
141
+ BLOCK_SIZE=block_size,
142
+ BLOCK_STRIDE_ROWS=block_stride_rows,
143
+ NUM_TOPK_TOKENS=logical_topk_indices.shape[1],
144
+ return_valid_counts=return_valid_counts,
145
+ out=physical_topk_indices,
146
+ valid_counts_out=(
147
+ valid_topk_counts if return_valid_counts else None
148
+ ),
149
+ )
150
+ self.physical_topk_ready.record(self.side_stream)
151
+ compute_stream.wait_event(self.physical_topk_ready)
152
+
153
+ if return_valid_counts:
154
+ return physical_topk_indices, valid_topk_counts
155
+ return physical_topk_indices
156
+
157
+
158
+@dataclass
159
+class HiSparseMLAIndexGroup(SparseMLAIndexGroup):
160
+ """Sparse index group with per-layer host-backed KV caches."""
161
+
162
+ caches: list[HiSparseCacheHandle] = field(default_factory=list)
163
+ hisparse_group: Any | None = None
164
+ prefill_stream: torch.Stream | None = None
165
+ prefill_ready_events: list[torch.Event] = field(default_factory=list)
166
+
167
+ def register_layer(
168
+ self,
169
+ vllm_config: VllmConfig | None = None,
170
+ *,
171
+ head_size: int | None = None,
172
+ kv_cache_dtype: str | None = None,
173
+ ) -> int:
174
+ assert vllm_config is not None
175
+ assert head_size is not None
176
+ assert kv_cache_dtype is not None
177
+ layer_index = super().register_layer()
178
+ if self.prefill_stream is None:
179
+ self.prefill_stream = _create_side_stream(self.logical_topk_indices.device)
180
+ if kv_cache_dtype == "fp8_ds_mla":
181
+ row_width = FP8_DS_MLA_ROW_BYTES
182
+ kv_dtype = torch.uint8
183
+ else:
184
+ from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
185
+
186
+ row_width = head_size
187
+ kv_dtype = kv_cache_dtype_str_to_dtype(
188
+ kv_cache_dtype, vllm_config.model_config
189
+ )
190
+ cache = create_hisparse_cache_handle(
191
+ vllm_config,
192
+ self.logical_topk_indices.shape[1],
193
+ is_index_group_leader=layer_index == 0,
194
+ row_width=row_width,
195
+ kv_dtype=kv_dtype,
196
+ index_group=self,
197
+ )
198
+ assert cache is not None
199
+ self.caches.append(cache)
200
+ cache.index_group_caches = self.caches
201
+ self.prefill_ready_events.append(_create_event())
202
+ return layer_index
203
+
204
+ def cache(self, layer_index: int) -> HiSparseCacheHandle:
205
+ return self.caches[layer_index]
206
+
207
+ def physical_kv_cache(self, layer_index: int) -> torch.Tensor:
208
+ cache = self.cache(layer_index)
209
+ return cache.runtime.hot.attention_cache
210
+
211
+ def prepare_for_batch(self, layer_index: int, attn_metadata: Any | None) -> None:
212
+ if layer_index == 0:
213
+ self.cache(0).prepare_group_for_batch(attn_metadata)
214
+
215
+ def convert_logical_to_physical_topk(
216
+ self,
217
+ layer_index: int,
218
+ logical_topk_indices: torch.Tensor,
219
+ attn_metadata: Any,
220
+ *,
221
+ block_stride_rows: int | None,
222
+ return_valid_counts: bool,
223
+ req_id_per_token: torch.Tensor | None = None,
224
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
225
+ cache = self.cache(layer_index)
226
+ num_tokens = logical_topk_indices.shape[0]
227
+ if req_id_per_token is None:
228
+ req_id_per_token = attn_metadata.req_id_per_token[:num_tokens]
229
+ if num_tokens > self.physical_topk_indices.shape[0]:
230
+ # Prefill-sized batches do not fit the decode residency workspace.
231
+ # Non-resident prefills are staged before reaching this path.
232
+ assert cache.all_context_pages_resident
233
+ leader = self.cache(0)
234
+ assert leader.view is not None and leader.block_table is not None
235
+ return self._convert_once(
236
+ layer_index,
237
+ logical_topk_indices,
238
+ req_id_per_token,
239
+ leader.block_table,
240
+ leader.view.block_size,
241
+ block_stride_rows=leader.view.attention_block_stride,
242
+ return_valid_counts=return_valid_counts,
243
+ )
244
+ source_block_table = cache.source_block_table
245
+ assert source_block_table is not None
246
+ return cache.swap_in(
247
+ req_id_per_token,
248
+ block_table=source_block_table,
249
+ logical_topk_indices=logical_topk_indices,
250
+ block_size=attn_metadata.block_size,
251
+ return_valid_counts=return_valid_counts,
252
+ )
253
+
254
+ def convert_decode_logical_to_physical_topk(
255
+ self,
256
+ layer_index: int,
257
+ logical_topk_indices: torch.Tensor,
258
+ attn_metadata: Any,
259
+ *,
260
+ return_valid_counts: bool,
261
+ num_decodes: int | None = None,
262
+ decode_query_len: int | None = None,
263
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
264
+ num_tokens = logical_topk_indices.shape[0]
265
+ if num_decodes is None:
266
+ num_decodes = attn_metadata.num_decodes
267
+ if decode_query_len is None:
268
+ decode_query_len = attn_metadata.decode_max_query_len
269
+ if decode_query_len == 1:
270
+ return self.convert_logical_to_physical_topk(
271
+ layer_index,
272
+ logical_topk_indices,
273
+ attn_metadata,
274
+ block_stride_rows=None,
275
+ return_valid_counts=return_valid_counts,
276
+ req_id_per_token=self.request_ids[:num_decodes],
277
+ )
278
+
279
+ if num_tokens != num_decodes * decode_query_len:
280
+ query_start_loc = attn_metadata.query_start_loc[: num_decodes + 1]
281
+ physical_topk_indices = self.physical_topk_indices[: num_tokens + 1]
282
+ valid_topk_counts = self.valid_topk_counts[: num_tokens + 1]
283
+ request_ids = self.request_ids[:num_decodes]
284
+ cache = self.cache(layer_index)
285
+ source_block_table = cache.source_block_table
286
+ assert source_block_table is not None
287
+ for step in range(decode_query_len):
288
+ token_indices = query_start_loc[:-1] + step
289
+ active = (token_indices < query_start_loc[1:]) & (
290
+ token_indices < num_tokens
291
+ )
292
+ output_indices = torch.where(
293
+ active, token_indices, torch.full_like(token_indices, num_tokens)
294
+ ).long()
295
+ token_indices = token_indices.clamp(0, num_tokens - 1).long()
296
+ step_topk = logical_topk_indices.index_select(
297
+ 0, token_indices
298
+ ).masked_fill(~active.unsqueeze(1), -1)
299
+ step_result = cache.swap_in(
300
+ request_ids,
301
+ block_table=source_block_table,
302
+ logical_topk_indices=step_topk,
303
+ block_size=attn_metadata.block_size,
304
+ return_valid_counts=return_valid_counts,
305
+ )
306
+ if layer_index == 0:
307
+ if return_valid_counts:
308
+ step_indices, step_counts = step_result
309
+ valid_topk_counts.index_copy_(0, output_indices, step_counts)
310
+ else:
311
+ step_indices = step_result
312
+ physical_topk_indices.index_copy_(0, output_indices, step_indices)
313
+ physical_topk_indices = physical_topk_indices[:num_tokens]
314
+ if return_valid_counts:
315
+ return physical_topk_indices, valid_topk_counts[:num_tokens]
316
+ return physical_topk_indices
317
+
318
+ assert num_tokens == num_decodes * decode_query_len
319
+ logical_topk_by_request = logical_topk_indices.view(
320
+ num_decodes, decode_query_len, -1
321
+ )
322
+ physical_topk_by_request = self.physical_topk_indices[:num_tokens].view(
323
+ num_decodes, decode_query_len, -1
324
+ )
325
+ valid_topk_by_request = self.valid_topk_counts[:num_tokens].view(
326
+ num_decodes, decode_query_len
327
+ )
328
+ request_ids = self.request_ids[:num_decodes]
329
+ for step in range(decode_query_len):
330
+ cache = self.cache(layer_index)
331
+ source_block_table = cache.source_block_table
332
+ assert source_block_table is not None
333
+ cache.swap_in(
334
+ request_ids,
335
+ block_table=source_block_table,
336
+ logical_topk_indices=logical_topk_by_request[:, step],
337
+ block_size=attn_metadata.block_size,
338
+ return_valid_counts=return_valid_counts,
339
+ attention_indices_out=(
340
+ physical_topk_by_request[:, step] if layer_index == 0 else None
341
+ ),
342
+ valid_counts_out=(
343
+ valid_topk_by_request[:, step]
344
+ if layer_index == 0 and return_valid_counts
345
+ else None
346
+ ),
347
+ )
348
+ physical_topk_indices = physical_topk_by_request.view(num_tokens, -1)
349
+ if return_valid_counts:
350
+ return physical_topk_indices, valid_topk_by_request.view(num_tokens)
351
+ return physical_topk_indices
352
+
353
+ def stage_prefill_rows(
354
+ self, layer_index: int, kv_cache: torch.Tensor, attn_metadata: Any
355
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
356
+ cache = self.cache(layer_index)
357
+ prefill = attn_metadata.prefill
358
+ staging_plan = prefill.host_staging_plan if prefill is not None else None
359
+ assert staging_plan is not None
360
+ resident_cache = None
361
+ if cache.view is not None and cache.block_table is not None:
362
+ staging_plan.ensure_gpu_sources(
363
+ cache.block_table[attn_metadata.num_decodes :],
364
+ cache.view.block_size,
365
+ )
366
+ resident_cache = cache.view.cache
367
+ staged_cache = cache.runtime.gather_prefill_cache(
368
+ kv_cache,
369
+ staging_plan,
370
+ resident_cache=resident_cache,
371
+ )
372
+ req_ids = attn_metadata.req_id_per_token[attn_metadata.num_decode_tokens :]
373
+ if attn_metadata.num_decodes > 0:
374
+ req_ids = req_ids - attn_metadata.num_decodes
375
+ return staged_cache, staging_plan.block_table, req_ids
376
+
377
+ def _prefill_gather_plan(
378
+ self, layer_index: int, attn_metadata: Any
379
+ ) -> HiSparsePrefillStagingPlan:
380
+ cache = self.cache(layer_index)
381
+ prefill = attn_metadata.prefill
382
+ plan = prefill.host_staging_plan if prefill is not None else None
383
+ assert plan is not None
384
+ assert cache.view is not None and cache.block_table is not None
385
+ plan.ensure_gpu_sources(
386
+ cache.block_table[attn_metadata.num_decodes :],
387
+ cache.view.block_size,
388
+ )
389
+ assert plan.gpu_row_ids is not None
390
+ return plan
391
+
392
+ def gather_fp8_prefill(
393
+ self,
394
+ layer_index: int,
395
+ source_cache: torch.Tensor,
396
+ dst: torch.Tensor,
397
+ block_table: torch.Tensor,
398
+ workspace_starts: torch.Tensor,
399
+ batch_size: int,
400
+ attn_metadata: Any,
401
+ request_start: int,
402
+ ) -> torch.Event | None:
403
+ cache = self.cache(layer_index)
404
+ plan = self._prefill_gather_plan(layer_index, attn_metadata)
405
+ assert cache.view is not None and plan.gpu_row_ids is not None
406
+ row_width = source_cache.shape[-1]
407
+ prefill_stream = self.prefill_stream
408
+ assert prefill_stream is not None
409
+ block_table = plan.block_table[request_start : request_start + batch_size]
410
+ prefill_stream.wait_stream(current_stream())
411
+ with prefill_stream:
412
+ ops.cp_gather_and_upconvert_fp8_kv_cache(
413
+ cache.view.cache,
414
+ dst,
415
+ block_table,
416
+ workspace_starts,
417
+ batch_size,
418
+ host_cache=source_cache.view(-1, row_width),
419
+ host_row_ids=plan.row_ids,
420
+ device_row_ids=plan.gpu_row_ids,
421
+ )
422
+ ready = self.prefill_ready_events[layer_index]
423
+ ready.record(prefill_stream)
424
+ return ready
425
+
426
+
427
+class SparseMLAIndexGroupBuilder:
428
+ """Assign sparse MLA layers to their index-producing layer."""
429
+
430
+ def __init__(
431
+ self, logical_topk_indices: torch.Tensor, max_decode_rows: int | None = None
432
+ ) -> None:
433
+ self.logical_topk_indices = logical_topk_indices
434
+ self.max_decode_rows = (
435
+ logical_topk_indices.shape[0]
436
+ if max_decode_rows is None
437
+ else max_decode_rows
438
+ )
439
+ self.current_group: SparseMLAIndexGroup | None = None
440
+
441
+ def register_layer(
442
+ self,
443
+ is_index_producing_layer: bool,
444
+ vllm_config: VllmConfig | None = None,
445
+ *,
446
+ head_size: int | None = None,
447
+ kv_cache_dtype: str | None = None,
448
+ ) -> tuple[SparseMLAIndexGroup, int]:
449
+ if is_index_producing_layer or self.current_group is None:
450
+ group_cls = (
451
+ HiSparseMLAIndexGroup
452
+ if vllm_config is not None
453
+ and vllm_config.attention_config.hisparse_config is not None
454
+ else SparseMLAIndexGroup
455
+ )
456
+ workspace_rows = self.max_decode_rows
457
+ physical_topk_indices = torch.empty(
458
+ (workspace_rows + 1, self.logical_topk_indices.shape[1]),
459
+ dtype=self.logical_topk_indices.dtype,
460
+ device=self.logical_topk_indices.device,
461
+ )
462
+ self.current_group = group_cls(
463
+ logical_topk_indices=self.logical_topk_indices,
464
+ physical_topk_indices=physical_topk_indices,
465
+ valid_topk_counts=torch.empty(
466
+ workspace_rows + 1,
467
+ dtype=torch.int32,
468
+ device=self.logical_topk_indices.device,
469
+ ),
470
+ request_ids=torch.arange(
471
+ workspace_rows,
472
+ dtype=torch.int32,
473
+ device=self.logical_topk_indices.device,
474
+ ),
475
+ side_stream=_create_side_stream(self.logical_topk_indices.device),
476
+ logical_topk_ready=_create_event(),
477
+ physical_topk_ready=_create_event(),
478
+ has_indexer=is_index_producing_layer,
479
+ )
480
+ group = self.current_group
481
+ group_index = group.register_layer(
482
+ vllm_config,
483
+ head_size=head_size,
484
+ kv_cache_dtype=kv_cache_dtype,
485
+ )
486
+ return group, group_index
487
+
488
+
489
+def get_sparse_mla_index_group_max_rows(vllm_config: VllmConfig) -> int:
490
+ max_query_len = 1
491
+ speculative_config = vllm_config.speculative_config
492
+ if (
493
+ speculative_config is not None
494
+ and speculative_config.num_speculative_tokens is not None
495
+ ):
496
+ max_query_len += speculative_config.num_speculative_tokens * (
497
+ 2 if speculative_config.parallel_drafting else 1
498
+ )
499
+ scheduler_config = vllm_config.scheduler_config
500
+ return min(
501
+ scheduler_config.max_num_batched_tokens,
502
+ scheduler_config.max_num_seqs * max_query_len,
503
+ )
vllm/v1/attention/backends/mla/sparse_utils.py
+28
-6
@@ -407,6 +407,8 @@ def triton_convert_req_index_to_global_index(
407
prefill_workspace_request_ids: torch.Tensor | None = None,
408
prefill_workspace_starts: torch.Tensor | None = None,
409
return_valid_counts: bool = False,
410
+ out: torch.Tensor | None = None,
411
+ valid_counts_out: torch.Tensor | None = None,
412
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
413
"""
414
out[token_id, indice_id] =
@@ -462,16 +464,36 @@ def triton_convert_req_index_to_global_index(
464
# contiguous prefix [0, valid_count) and leaves the tail unwritten, so
465
# pre-fill -1 there. flash_mla_sparse_fwd then bounds attention to
466
# [:topk_length] == exactly the valid set (no dropped tokens).
465
- if return_valid_counts:
466
- out = torch.full_like(token_indices_c, -1)
467
+ if out is None:
468
+ out = (
469
+ torch.full_like(token_indices_c, -1)
470
+ if return_valid_counts
471
+ else torch.empty_like(token_indices_c)
472
+ )
473
else:
468
- out = torch.empty_like(token_indices_c)
474
+ assert out.dtype == token_indices_c.dtype
475
+ assert out.device == token_indices_c.device
476
+ assert out.shape == token_indices_c.shape
477
+ assert out.is_contiguous()
478
+ if return_valid_counts:
479
+ out.fill_(-1)
480
481
valid_counts: torch.Tensor | None = None
482
if return_valid_counts:
472
- # Zero-init only matters for the atomic accumulation path.
473
- alloc = torch.empty if single_tile else torch.zeros
474
- valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device)
483
+ if valid_counts_out is None:
484
+ # Zero-init only matters for the atomic accumulation path.
485
+ alloc = torch.empty if single_tile else torch.zeros
486
+ valid_counts = alloc(
487
+ num_tokens, dtype=torch.int32, device=token_indices.device
488
+ )
489
+ else:
490
+ assert valid_counts_out.dtype == torch.int32
491
+ assert valid_counts_out.device == token_indices.device
492
+ assert valid_counts_out.shape == (num_tokens,)
493
+ assert valid_counts_out.is_contiguous()
494
+ valid_counts = valid_counts_out
495
+ if not single_tile:
496
+ valid_counts.zero_()
497
498
# Prepare prefill pointers
499
if HAS_PREFILL_WORKSPACE:
vllm/v1/core/block_pool.py
+51
-5
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections.abc import Iterable, Sequence
3
+from collections.abc import Callable, Iterable, Sequence
4
from typing import Any
5
6
from vllm.distributed.kv_events import (
@@ -157,6 +157,7 @@ class BlockPool:
157
actual block size can be a multiple of hash_block_size.
158
enable_kv_cache_events: Whether to enable kv cache events.
159
metrics_collector: Optional metrics collector for tracking block residency.
160
+ medium: Storage medium reported in KV cache events.
161
"""
162
163
def __init__(
@@ -166,14 +167,16 @@ class BlockPool:
167
hash_block_size: int,
168
enable_kv_cache_events: bool = False,
169
metrics_collector: KVCacheMetricsCollector | None = None,
170
+ medium: str = MEDIUM_GPU,
171
):
172
assert isinstance(num_gpu_blocks, int) and num_gpu_blocks > 0
173
self.num_gpu_blocks = num_gpu_blocks
174
+ self.medium = medium
175
self.enable_caching = enable_caching
176
self.hash_block_size = hash_block_size
177
# All kv-cache blocks.
178
self.blocks: list[KVCacheBlock] = [
176
- KVCacheBlock(idx) for idx in range(num_gpu_blocks)
179
+ KVCacheBlock(idx, pool=self) for idx in range(num_gpu_blocks)
180
]
181
# Free block queue that constructs and manipulates a doubly linked
182
# list of free blocks (including eviction candidates when caching is
@@ -194,6 +197,9 @@ class BlockPool:
197
self.kv_event_queue: list[KVCacheEvent] = []
198
199
self.metrics_collector = metrics_collector
200
+ # Callbacks for blocks released with ``unpin_blocks`` whose contents
201
+ # are still being read until the pool reuses them.
202
+ self._reuse_watchers: dict[int, Callable[[KVCacheBlock], None]] = {}
203
204
def get_cached_block(
205
self, block_hash: BlockHash, kv_cache_group_ids: list[int]
@@ -364,7 +370,7 @@ class BlockPool:
370
token_ids=request.all_token_ids[start_token_idx:end_token_idx],
371
block_size=block_size,
372
lora_id=request.lora_request.adapter_id if request.lora_request else None,
367
- medium=MEDIUM_GPU,
373
+ medium=self.medium,
374
lora_name=request.lora_request.name if request.lora_request else None,
375
extra_keys=extra_keys_list if extra_keys_list else None,
376
group_idx=kv_cache_group_id,
@@ -543,7 +549,7 @@ class BlockPool:
549
lora_id=request.lora_request.adapter_id
550
if request.lora_request
551
else None,
546
- medium=MEDIUM_GPU,
552
+ medium=self.medium,
553
lora_name=request.lora_request.name
554
if request.lora_request
555
else None,
@@ -610,7 +616,7 @@ class BlockPool:
616
self.kv_event_queue.append(
617
BlockRemoved(
618
block_hashes=[maybe_convert_block_hash(get_block_hash(block_hash))],
613
- medium=MEDIUM_GPU,
619
+ medium=self.medium,
620
group_idx=get_group_id(block_hash),
621
)
622
)
@@ -671,6 +677,9 @@ class BlockPool:
677
678
ret: list[KVCacheBlock] = self.free_block_queue.popleft_n(num_blocks)
679
680
+ if self._reuse_watchers:
681
+ self._notify_reuse(ret)
682
+
683
# In order to only iterate the list once, we duplicated code a bit
684
if self.enable_caching:
685
for block in ret:
@@ -687,6 +696,33 @@ class BlockPool:
696
self.metrics_collector.on_block_allocated(block)
697
return ret
698
699
+ def _notify_reuse(self, blocks: Iterable[KVCacheBlock]) -> None:
700
+ for block in blocks:
701
+ watcher = self._reuse_watchers.pop(block.block_id, None)
702
+ if watcher is not None:
703
+ watcher(block)
704
+
705
+ def unpin_blocks(
706
+ self,
707
+ blocks: Iterable[KVCacheBlock],
708
+ on_reuse: Callable[[KVCacheBlock], None],
709
+ ) -> None:
710
+ """Release references to blocks that stay readable until reused.
711
+
712
+ The blocks become last-resort eviction candidates regardless of prefix
713
+ caching: they join the tail of the free queue and count as free.
714
+ ``on_reuse`` fires when ``get_new_blocks`` hands a block out (or the
715
+ cache is reset), so the caller can stop reading it.
716
+ """
717
+ released: list[KVCacheBlock] = []
718
+ for block in blocks:
719
+ assert block.ref_cnt > 0 and not block.is_null
720
+ block.ref_cnt -= 1
721
+ self._reuse_watchers[block.block_id] = on_reuse
722
+ if block.ref_cnt == 0:
723
+ released.append(block)
724
+ self.free_block_queue.append_n(released)
725
+
726
def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool:
727
"""
728
If a block is cached in `cached_block_hash_to_block`, we reset its hash
@@ -742,7 +778,11 @@ class BlockPool:
778
# Identify blocks with hash (LRU cache) and without it (never match APC)
779
blocks_to_evict_last = []
780
blocks_to_evict_first = []
781
+ other_pools: dict[BlockPool, list[KVCacheBlock]] = {}
782
for block in ordered_blocks:
783
+ if block.pool is not None and block.pool is not self:
784
+ other_pools.setdefault(block.pool, []).append(block)
785
+ continue
786
block.ref_cnt -= 1
787
if block.ref_cnt == 0 and not block.is_null:
788
if block.block_hash is None or not self.enable_caching:
@@ -756,6 +796,8 @@ class BlockPool:
796
self.free_block_queue.prepend_n(blocks_to_evict_first)
797
# Blocks to reuse last are appended to the end of the free queue.
798
self.free_block_queue.append_n(blocks_to_evict_last)
799
+ for pool, blocks in other_pools.items():
800
+ pool.free_blocks(blocks)
801
802
def evict_blocks(self, block_ids: set[int]) -> None:
803
"""evict blocks from the prefix cache by their block IDs.
@@ -797,6 +839,10 @@ class BlockPool:
839
# Remove all hashes so that no new blocks will hit.
840
self.cached_block_hash_to_block = BlockHashToBlockMap()
841
self.cached_block_hashes_by_block.clear()
842
+ if self._reuse_watchers:
843
+ self._notify_reuse(
844
+ [self.blocks[block_id] for block_id in list(self._reuse_watchers)]
845
+ )
846
847
# Remove all hashes from all blocks.
848
for block in self.blocks:
vllm/v1/core/kv_cache_coordinator.py
+39
-8
@@ -139,6 +139,7 @@ class KVCacheCoordinator(ABC):
139
max_in_flight_tokens=max_in_flight_tokens,
140
max_model_len=max_model_len,
141
block_pool=self.block_pool,
142
+ role=kv_cache_group.role,
143
enable_caching=enable_caching,
144
kv_cache_group_id=i,
145
dcp_world_size=dcp_world_size_for_kv_cache_spec(
@@ -179,7 +180,7 @@ class KVCacheCoordinator(ABC):
180
apply_admission_cap: bool = False,
181
) -> int:
182
"""
182
- Get the number of blocks needed to be allocated for the request.
183
+ Get the number of device blocks needed to be allocated for the request.
184
185
Args:
186
request_id: The request ID.
@@ -343,6 +344,8 @@ class KVCacheCoordinator(ABC):
344
"""
345
boundaries = self.get_replay_boundaries(request)
346
for manager in self.single_type_managers:
347
+ if not manager.enable_caching:
348
+ continue
349
# Only cache tokens with finalized KV. The last num_reprefillable_tokens
350
# tokens can be re-prefilled during multi-module MTP.
351
num_tokens_to_cache = max(
@@ -355,6 +358,27 @@ class KVCacheCoordinator(ABC):
358
replay_boundaries=boundaries,
359
)
360
361
+ def emit_cached_block_events(
362
+ self, request: Request, computed_blocks: tuple[list[KVCacheBlock], ...]
363
+ ) -> None:
364
+ for group_idx, group_blocks in enumerate(computed_blocks):
365
+ if group_blocks:
366
+ manager = self.single_type_managers[group_idx]
367
+ group = self.kv_cache_config.kv_cache_groups[group_idx]
368
+ manager.block_pool.emit_cached_block_events(
369
+ request,
370
+ len(group_blocks),
371
+ group.kv_cache_spec.block_size,
372
+ group_idx,
373
+ )
374
+
375
+ def reset_prefix_cache(self) -> bool:
376
+ """Reset each manager's pool once."""
377
+ pools = dict.fromkeys(
378
+ manager.block_pool for manager in self.single_type_managers
379
+ )
380
+ return all([pool.reset_prefix_cache() for pool in pools])
381
+
382
def free(self, request_id: str) -> None:
383
"""
384
Free the blocks for the request.
@@ -728,10 +752,12 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
752
753
# Try to find an existing group with the same spec
754
for idx, group in enumerate(self.attention_groups):
731
- if group.spec == spec:
732
- assert manager_cls is group.manager_cls, (
733
- "Expected same manager class for identical KV cache specs."
734
- )
755
+ if (
756
+ group.spec == spec
757
+ and group.manager_cls is manager_cls
758
+ and self.single_type_managers[group.group_ids[0]].block_pool
759
+ is self.single_type_managers[i].block_pool
760
+ ):
761
group.group_ids.append(i)
762
if use_eagle and not group.use_eagle:
763
self.attention_groups[idx] = group._replace(use_eagle=True)
@@ -782,6 +808,8 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
808
cached_num_computed_tokens = self._align_cacheable(num_computed_tokens)
809
boundaries = self.get_replay_boundaries(request)
810
for manager in self.single_type_managers:
811
+ if not manager.enable_caching:
812
+ continue
813
num_tokens_to_cache = cached_num_computed_tokens
814
# EAGLE groups match one block past each aligned boundary and drop
815
# it, so make that lookahead block eligible to be cached.
@@ -895,7 +923,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
923
block_hashes=block_hashes,
924
max_length=_max_length,
925
kv_cache_group_ids=group_ids,
898
- block_pool=self.block_pool,
926
+ block_pool=self.single_type_managers[first_group_id].block_pool,
927
kv_cache_spec=spec,
928
drop_eagle_block=drop_eagle_block,
929
alignment_tokens=self._cache_hit_alignment_tokens,
@@ -929,7 +957,10 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
957
for group in self.attention_groups:
958
if not isinstance(group.spec, FullAttentionSpec):
959
continue
932
- group_block_size = self.single_type_managers[group.group_ids[0]].block_size
960
+ manager = self.single_type_managers[group.group_ids[0]]
961
+ if manager.retains_longer_hit:
962
+ continue
963
+ group_block_size = manager.block_size
964
num_blocks = cdiv(hit_length, group_block_size)
965
for group_id in group.group_ids:
966
if (blks := hit_blocks_by_group[group_id]) is not None:
@@ -966,7 +997,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
997
block_hashes=block_hashes,
998
max_length=max_cache_hit_length,
999
kv_cache_group_ids=group_ids,
969
- block_pool=self.block_pool,
1000
+ block_pool=manager.block_pool,
1001
kv_cache_spec=spec,
1002
drop_eagle_block=use_eagle,
1003
alignment_tokens=self._cache_hit_alignment_tokens,
vllm/v1/core/kv_cache_manager.py
+20
-13
@@ -183,6 +183,11 @@ class KVCacheManager:
183
manager.fine_grained_prefix_cache = True
184
self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
185
self.block_pool = self.coordinator.block_pool
186
+ self.retained_hit_group_ids = tuple(
187
+ manager.kv_cache_group_id
188
+ for manager in self.coordinator.single_type_managers
189
+ if manager.retains_longer_hit
190
+ )
191
self.kv_cache_config = kv_cache_config
192
193
# Watermark: minimum number of KV cache blocks to keep free when
@@ -287,17 +292,7 @@ class KVCacheManager:
292
and self.enable_kv_cache_events
293
and getattr(request, "kv_cache_report_mode", "incremental") == "full"
294
):
290
- for group_idx, group_blocks in enumerate(computed_blocks):
291
- num_blocks = len(group_blocks)
292
- if num_blocks > 0:
293
- group = self.kv_cache_config.kv_cache_groups[group_idx]
294
- block_size = group.kv_cache_spec.block_size
295
- self.block_pool.emit_cached_block_events(
296
- request,
297
- num_blocks,
298
- block_size,
299
- group_idx,
300
- )
295
+ self.coordinator.emit_cached_block_events(request, computed_blocks)
296
297
# The junction to pin is where the lagging sparse-retention group stops
298
# (``num_new_computed_tokens``) plus the uncached shared prefix -- i.e.
@@ -463,6 +458,16 @@ class KVCacheManager:
458
459
if new_computed_blocks is not None:
460
new_computed_block_list = new_computed_blocks.blocks
461
+ if self.retained_hit_group_ids:
462
+ # Only retain the prefix covered by the local and external hits.
463
+ reused = num_new_computed_tokens + num_external_computed_tokens
464
+ groups = list(new_computed_block_list)
465
+ for group_id in self.retained_hit_group_ids:
466
+ manager = self.coordinator.single_type_managers[group_id]
467
+ groups[group_id] = groups[group_id][
468
+ : cdiv(reused, manager.block_size)
469
+ ]
470
+ new_computed_block_list = tuple(groups)
471
else:
472
new_computed_block_list = self.empty_kv_cache_blocks.blocks
473
@@ -564,7 +569,7 @@ class KVCacheManager:
569
570
# P/D: delay caching blocks if we have to recv from
571
# remote. Update state for locally cached blocks.
567
- if not self.enable_caching or delay_cache_blocks:
572
+ if delay_cache_blocks:
573
return self.create_kv_cache_blocks(new_blocks)
574
575
# NOTE(woosuk): We want to commit (cache) up to num_local_computed_tokens
@@ -639,7 +644,7 @@ class KVCacheManager:
644
bool: True if the prefix cache is successfully reset,
645
False otherwise.
646
"""
642
- if not self.block_pool.reset_prefix_cache():
647
+ if not self.coordinator.reset_prefix_cache():
648
return False
649
if self.log_stats:
650
assert self.prefix_cache_stats is not None
@@ -741,6 +746,8 @@ class KVCacheManager:
746
self.kv_cache_config.kv_cache_groups,
747
self.get_blocks(request.request_id).blocks,
748
):
749
+ if not group.kv_cache_spec.prefix_cacheable:
750
+ continue
751
if isinstance(
752
group.kv_cache_spec,
753
(CrossAttentionSpec, EncoderOnlyAttentionSpec),
vllm/v1/core/kv_cache_utils.py
+74
-10
@@ -8,23 +8,30 @@ import math
8
import os
9
from collections import defaultdict
10
from collections.abc import Callable, Iterable, Iterator, Sequence
11
-from dataclasses import dataclass, replace
11
+from dataclasses import dataclass, field, replace
12
from functools import partial
13
-from typing import Any, NamedTuple, NewType, TypeAlias, cast, overload
13
+from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, cast, overload
14
15
from vllm import envs
16
from vllm.config import VllmConfig
17
from vllm.logger import init_logger
18
from vllm.utils.hashing import xxhash, xxhash_cbor
19
-from vllm.utils.math_utils import cdiv
19
+from vllm.utils.math_utils import cdiv, round_up
20
from vllm.utils.mem_utils import format_gib
21
from vllm.utils.torch_utils import get_dtype_size
22
+from vllm.v1.hisparse.layout import (
23
+ get_hisparse_gpu_memory_usage,
24
+ get_hisparse_host_pool_bytes,
25
+ get_hisparse_kv_cache_config,
26
+ get_hisparse_kv_cache_groups,
27
+)
28
from vllm.v1.kv_cache_interface import (
29
AttentionSpec,
30
ChunkedLocalAttentionSpec,
31
CircularBufferSpec,
32
FullAttentionSpec,
33
HiddenStateCacheSpec,
34
+ HiSparseHotSpec,
35
KpoolTailSpec,
36
KVCacheConfig,
37
KVCacheGroupSpec,
@@ -44,6 +51,10 @@ from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
51
from vllm.v1.request import Request
52
from vllm.v1.utils import tensor_data
53
54
+if TYPE_CHECKING:
55
+ from vllm.v1.core.block_pool import BlockPool
56
+
57
+
58
# BlockHash represents the hash of a single KV-cache block used for
59
# prefix caching. Treating it as a distinct type from `bytes` helps
60
# catch accidental misuse when passing around raw byte strings.
@@ -169,6 +180,10 @@ class KVCacheBlock:
180
block_id: int
181
# Reference count.
182
ref_cnt: int = 0
183
+ # Deferred frees may contain blocks from multiple pools.
184
+ pool: "BlockPool | None" = field(
185
+ default=None, repr=False, compare=False, kw_only=True
186
+ )
187
# The hash key (block hash + group id) of the block, only available
188
# when the block is full and cached.
189
_block_hash: BlockHashWithGroupId | None = None
@@ -1057,21 +1072,32 @@ def get_max_concurrency_for_kv_cache_config(
1072
1073
A request at max_model_len consumes whole blocks from each group's block
1074
table — cdiv(per-request bytes, page bytes) of the group's spec — and all
1060
- groups draw those block ids from one shared pool, so the per-request
1075
+ device groups draw those block ids from one shared pool, so the per-request
1076
total is the sum over groups. The memory/page ratio is identical whether
1077
a group carries an aggregated UniformTypeKVCacheSpecs (worker config) or
1078
a representative per-layer spec (scheduler config), so both capacity
1079
call sites agree.
1080
+
1081
+ Host groups use a separate pool; the smaller concurrency limit applies.
1082
"""
1066
- num_blocks_per_request = sum(
1067
- cdiv(
1083
+ num_blocks_per_request = 0
1084
+ host_blocks_per_request = 0
1085
+ for group in kv_cache_config.kv_cache_groups:
1086
+ required = cdiv(
1087
group.kv_cache_spec.max_memory_usage_bytes(vllm_config),
1088
group.kv_cache_spec.page_size_bytes,
1089
)
1071
- for group in kv_cache_config.kv_cache_groups
1072
- )
1073
- max_concurrency = kv_cache_config.num_blocks / num_blocks_per_request
1074
- return max_concurrency
1090
+ if group.host_resident:
1091
+ host_blocks_per_request += required
1092
+ else:
1093
+ num_blocks_per_request += required
1094
+ limits = [kv_cache_config.num_blocks / num_blocks_per_request]
1095
+ if host_blocks_per_request:
1096
+ assert kv_cache_config.hisparse_host_num_blocks is not None
1097
+ limits.append(
1098
+ kv_cache_config.hisparse_host_num_blocks / host_blocks_per_request
1099
+ )
1100
+ return min(limits)
1101
1102
1103
def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int:
@@ -1575,6 +1601,13 @@ def _get_kv_cache_bytes_per_block(
1601
for group in kv_cache_groups
1602
)
1603
assert bytes_per_block > 0
1604
+ hot_page_sizes = [
1605
+ group.kv_cache_spec.page_size_bytes
1606
+ for group in kv_cache_groups
1607
+ if isinstance(group.kv_cache_spec, HiSparseHotSpec)
1608
+ ]
1609
+ if hot_page_sizes:
1610
+ bytes_per_block = round_up(bytes_per_block, math.lcm(*hot_page_sizes))
1611
return bytes_per_block
1612
1613
@@ -1640,6 +1673,12 @@ def get_kv_cache_config_from_groups(
1673
),
1674
)
1675
1676
+ if vllm_config.attention_config.hisparse_config is not None:
1677
+ host_budget = get_hisparse_host_pool_bytes(vllm_config)
1678
+ return get_hisparse_kv_cache_config(
1679
+ vllm_config, kv_cache_groups, available_memory, host_budget
1680
+ )
1681
+
1682
if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
1683
(
1684
attn_group,
@@ -2210,6 +2249,9 @@ def get_kv_cache_groups(
2249
# attention free models.
2250
return []
2251
2252
+ if hisparse_groups := get_hisparse_kv_cache_groups(vllm_config, kv_cache_spec):
2253
+ return hisparse_groups
2254
+
2255
if is_kv_cache_spec_uniform(kv_cache_spec):
2256
# KV cache of all layers are the same, which is true for
2257
# most models. Allocate the same amount of memory for
@@ -2287,6 +2329,10 @@ def generate_scheduler_kv_cache_config(
2329
assert all(
2330
[cfg.num_blocks == kv_cache_configs[0].num_blocks for cfg in kv_cache_configs]
2331
)
2332
+ assert all(
2333
+ cfg.hisparse_host_num_blocks == kv_cache_configs[0].hisparse_host_num_blocks
2334
+ for cfg in kv_cache_configs
2335
+ )
2336
# All workers have the same kv_cache_config except layer names, so use
2337
# an arbitrary one to initialize the scheduler.
2338
cfg = copy.deepcopy(kv_cache_configs[0])
@@ -2347,6 +2393,9 @@ def _max_memory_usage_bytes_from_groups(
2393
if not kv_cache_groups:
2394
return 0
2395
2396
+ if vllm_config.attention_config.hisparse_config is not None:
2397
+ return get_hisparse_gpu_memory_usage(vllm_config, kv_cache_groups)
2398
+
2399
if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
2400
(
2401
attn_group,
@@ -2396,9 +2445,21 @@ def _estimate_max_model_len_from_groups(
2445
Returns 0 if even 1 token doesn't fit.
2446
"""
2447
original_max = vllm_config.model_config.max_model_len
2448
+ hisparse_enabled = (
2449
+ vllm_config.attention_config.hisparse_config is not None
2450
+ and bool(kv_cache_groups)
2451
+ )
2452
2453
def fits(model_len: int) -> bool:
2454
vllm_config.model_config.max_model_len = model_len
2455
+ if hisparse_enabled:
2456
+ try:
2457
+ config = get_kv_cache_config_from_groups(
2458
+ vllm_config, kv_cache_groups, available_memory
2459
+ )
2460
+ except ValueError:
2461
+ return False
2462
+ return get_max_concurrency_for_kv_cache_config(vllm_config, config) >= 1
2463
return (
2464
_max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups)
2465
<= available_memory
@@ -2629,6 +2690,9 @@ def get_kv_cache_configs(
2690
adjusted_memory.append(override * bytes_per_block)
2691
available_memory = adjusted_memory
2692
2693
+ if vllm_config.attention_config.hisparse_config is not None:
2694
+ available_memory = [min(available_memory)] * len(available_memory)
2695
+
2696
# Reserve the null block BlockPool permanently holds back, so auto-fit and
2697
# the capacity check both plan against usable blocks. Allocation below
2698
# still uses the full memory.
vllm/v1/core/sched/output.py
+3
@@ -294,6 +294,9 @@ class SchedulerOutput:
294
# CoW copies to apply after zeroing new blocks and before forward.
295
kv_cache_block_copies: list[KVCacheBlockCopy] | None = None
296
297
+ # Complete block-table rows that replace incrementally appended block IDs.
298
+ block_table_updates: dict[str, tuple[list[int], ...]] | None = None
299
+
300
# Scheduler-local; always None by the time this reaches a worker.
301
kv_connector_block_state: KVConnectorBlockState | None = None
302
vllm/v1/core/sched/scheduler.py
+33
-7
@@ -316,10 +316,9 @@ class Scheduler(SchedulerInterface):
316
self.cache_config.enable_mamba_fine_grained_prefix_cache
317
),
318
)
319
- # Bind GPU block pool to the KV connector. This must happen after
320
- # kv_cache_manager is constructed so block_pool is available.
319
+ # Bind after construction so connectors can access the cache manager.
320
if self.connector is not None:
322
- self.connector.bind_gpu_block_pool(self.kv_cache_manager.block_pool)
321
+ self.connector.bind_kv_cache_manager(self.kv_cache_manager)
322
323
self.use_pp = self.parallel_config.pipeline_parallel_size > 1
324
self.use_v2_model_runner = vllm_config.use_v2_model_runner
@@ -738,6 +737,12 @@ class Scheduler(SchedulerInterface):
737
# The request can be scheduled.
738
break
739
740
+ if (
741
+ self.connector is not None
742
+ and self.connector.has_pending_block_frees()
743
+ ):
744
+ break
745
+
746
# The request cannot be scheduled.
747
# Preempt the lowest-priority request.
748
if self.policy == SchedulingPolicy.PRIORITY:
@@ -1895,14 +1900,20 @@ class Scheduler(SchedulerInterface):
1900
outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
1901
spec_decoding_stats: SpecDecodingStats | None = None
1902
1898
- failed_kv_load_req_ids = None
1903
+ failed_kv_load_req_ids: set[str] = set()
1904
+ if kv_connector_output and kv_connector_output.failed_recving:
1905
+ failed_kv_load_req_ids.update(
1906
+ self._handle_failed_recving(kv_connector_output.failed_recving)
1907
+ )
1908
if kv_connector_output and kv_connector_output.invalid_block_ids:
1909
# These blocks contain externally computed tokens that failed to
1910
# load. Identify affected requests and adjust their computed token
1911
# count to trigger recomputation of the invalid blocks.
1903
- failed_kv_load_req_ids = self._handle_invalid_blocks(
1904
- kv_connector_output.invalid_block_ids,
1905
- num_scheduled_tokens,
1912
+ failed_kv_load_req_ids.update(
1913
+ self._handle_invalid_blocks(
1914
+ kv_connector_output.invalid_block_ids,
1915
+ num_scheduled_tokens,
1916
+ )
1917
)
1918
1919
# Persist per-step routed experts into the scheduler-side slot
@@ -3135,6 +3146,21 @@ class Scheduler(SchedulerInterface):
3146
3147
return affected_req_ids, total_affected_tokens, blocks_to_evict
3148
3149
+ def _handle_failed_recving(self, failed_req_ids: set[str]) -> set[str]:
3150
+ """Fail closed for layouts whose block IDs are not globally unique."""
3151
+ affected_req_ids: set[str] = set()
3152
+ for req_id in failed_req_ids:
3153
+ request = self.requests.get(req_id)
3154
+ if (
3155
+ request is not None
3156
+ and request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
3157
+ ):
3158
+ affected_req_ids.add(req_id)
3159
+ if self.recompute_kv_load_failures:
3160
+ request.num_computed_tokens = 0
3161
+ self.failed_recving_kv_req_ids.add(req_id)
3162
+ return set() if self.recompute_kv_load_failures else affected_req_ids
3163
+
3164
def _handle_invalid_blocks(
3165
self, invalid_block_ids: set[int], num_scheduled_tokens: dict[str, int]
3166
) -> set[str]:
vllm/v1/core/single_type_kv_cache_manager.py
+405
-3
@@ -4,8 +4,9 @@ import itertools
4
from abc import ABC, abstractmethod
5
from collections import defaultdict
6
from collections.abc import Sequence
7
-from typing import ClassVar
7
+from typing import TYPE_CHECKING, ClassVar
8
9
+from vllm.distributed.kv_events import MEDIUM_CPU
10
from vllm.logger import init_logger
11
from vllm.utils.math_utils import cdiv
12
from vllm.v1.core.block_pool import BlockPool
@@ -16,6 +17,7 @@ from vllm.v1.core.kv_cache_utils import (
17
KVCacheBlock,
18
resolve_block_hashes,
19
)
20
+from vllm.v1.hisparse.block_pool import SharedEventQueueBlockPool
21
from vllm.v1.kv_cache_interface import (
22
AttentionSpec,
23
ChunkedLocalAttentionSpec,
@@ -23,7 +25,10 @@ from vllm.v1.kv_cache_interface import (
25
CrossAttentionSpec,
26
FullAttentionSpec,
27
HiddenStateCacheSpec,
28
+ HiSparseHotSpec,
29
+ HiSparseResidentSpec,
30
KpoolTailSpec,
31
+ KVCacheGroupRole,
32
KVCacheSpec,
33
MambaSpec,
34
MLAAttentionSpec,
@@ -37,6 +42,9 @@ from vllm.v1.kv_cache_interface import (
42
from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
43
from vllm.v1.request import Request
44
45
+if TYPE_CHECKING:
46
+ from vllm.v1.hisparse.coordinator import HiSparseCoordinator
47
+
48
logger = init_logger(__name__)
49
50
@@ -48,6 +56,9 @@ class SingleTypeKVCacheManager(ABC):
56
57
supports_fine_grained_hash_lookup: ClassVar[bool] = False
58
59
+ # Keep this group's longer prefix until external cache lookup completes.
60
+ retains_longer_hit: bool = False
61
+
62
@property
63
def has_positionally_stable_blocks(self) -> bool:
64
"""Whether positional offload scans can follow this block table.
@@ -2191,10 +2202,388 @@ class SinkFullAttentionManager(FullAttentionManager):
2202
assert sink_len is not None and sink_len > 0 and sink_len % self.block_size == 0
2203
2204
2205
+class HiSparseSourceManager(FullAttentionManager):
2206
+ """Host-tier manager with a private pool; publishes hashes once durable.
2207
+
2208
+ Host capacity is best effort: a page that cannot get a host block keeps a
2209
+ null host entry, so its GPU copy is never written back and stays pinned.
2210
+ """
2211
+
2212
+ coordinator: "HiSparseCoordinator | None" = None
2213
+ # The host tier keeps prefixes the device groups have already lost.
2214
+ retains_longer_hit = True
2215
+
2216
+ @property
2217
+ def records_new_block_ids(self) -> bool:
2218
+ return False
2219
+
2220
+ def take_new_block_ids(self) -> list[int]:
2221
+ self.new_block_ids = []
2222
+ return []
2223
+
2224
+ def bind_host_pool(self, num_blocks: int) -> None:
2225
+ """Replace the device pool this group was built with by a host one."""
2226
+ device_pool = self.block_pool
2227
+ self.block_pool = SharedEventQueueBlockPool(
2228
+ num_gpu_blocks=num_blocks,
2229
+ enable_caching=self.enable_caching and self.kv_cache_spec.prefix_cacheable,
2230
+ hash_block_size=device_pool.hash_block_size,
2231
+ enable_kv_cache_events=device_pool.enable_kv_cache_events,
2232
+ metrics_collector=device_pool.metrics_collector,
2233
+ medium=MEDIUM_CPU,
2234
+ event_owner=device_pool,
2235
+ )
2236
+ self._null_block = self.block_pool.null_block
2237
+
2238
+ def take_pending_cow_copies(self) -> list[tuple[KVCacheBlock, KVCacheBlock]]:
2239
+ """Host copies never reach the worker's generic block-copy path."""
2240
+ return []
2241
+
2242
+ def take_host_cow_copies(self) -> list[tuple[KVCacheBlock, KVCacheBlock]]:
2243
+ """Drain host copies for the coordinator to hand to the connector."""
2244
+ copies = self._pending_cow_copies
2245
+ self._pending_cow_copies = []
2246
+ return copies
2247
+
2248
+ def get_num_blocks_to_allocate(
2249
+ self,
2250
+ request_id: str,
2251
+ num_tokens: int,
2252
+ new_computed_blocks: Sequence[KVCacheBlock],
2253
+ total_computed_tokens: int,
2254
+ num_local_computed_tokens: int,
2255
+ num_tokens_main_model: int,
2256
+ apply_admission_cap: bool = False,
2257
+ ) -> int:
2258
+ if (
2259
+ total_computed_tokens > num_local_computed_tokens
2260
+ or self._has_partial_local_hit(
2261
+ new_computed_blocks, num_local_computed_tokens
2262
+ )
2263
+ ):
2264
+ # External loads need real destinations; future GPU-computed pages
2265
+ # remain best effort. Use the same admission sentinel as Mamba.
2266
+ required = super().get_num_blocks_to_allocate(
2267
+ request_id,
2268
+ total_computed_tokens,
2269
+ new_computed_blocks,
2270
+ total_computed_tokens,
2271
+ num_local_computed_tokens,
2272
+ total_computed_tokens,
2273
+ )
2274
+ if required > self.block_pool.get_num_free_blocks():
2275
+ assert self.coordinator is not None
2276
+ assert self.coordinator.gpu_pool is not None
2277
+ return self.coordinator.gpu_pool.num_gpu_blocks + 1
2278
+ return 0
2279
+
2280
+ def allocate_new_blocks(
2281
+ self, request_id: str, num_tokens: int, num_tokens_main_model: int
2282
+ ) -> list[KVCacheBlock]:
2283
+ req_blocks = self.req_to_blocks[request_id]
2284
+ num_required = cdiv(num_tokens, self.block_size)
2285
+ num_free = self.block_pool.get_num_free_blocks()
2286
+ if request_id in self._partial_hit_reqs:
2287
+ num_free -= 1
2288
+ fit_blocks = min(num_required, len(req_blocks) + max(0, num_free))
2289
+ new_blocks = super().allocate_new_blocks(
2290
+ request_id,
2291
+ fit_blocks * self.block_size,
2292
+ min(num_tokens_main_model, fit_blocks * self.block_size),
2293
+ )
2294
+ req_blocks.extend([self._null_block] * (num_required - len(req_blocks)))
2295
+ return new_blocks
2296
+
2297
+ def allocate_external_computed_blocks(
2298
+ self,
2299
+ request_id: str,
2300
+ num_local_computed_tokens: int,
2301
+ num_external_computed_tokens: int,
2302
+ ) -> None:
2303
+ if num_external_computed_tokens <= 0:
2304
+ return
2305
+ # The connector writes these pages; only a successful receive makes
2306
+ # them readable, not advancing the request's computed-token count.
2307
+ assert self.coordinator is not None
2308
+ self.coordinator.record_pending_host_import(
2309
+ request_id, num_local_computed_tokens + num_external_computed_tokens
2310
+ )
2311
+ super().allocate_external_computed_blocks(
2312
+ request_id, num_local_computed_tokens, num_external_computed_tokens
2313
+ )
2314
+
2315
+ def cache_blocks(
2316
+ self,
2317
+ request: Request,
2318
+ num_tokens: int,
2319
+ retention_interval: int | None = None,
2320
+ *,
2321
+ replay_boundaries: Sequence[int],
2322
+ ) -> None:
2323
+ assert self.coordinator is not None
2324
+ self.coordinator.publish_when_ready(
2325
+ request,
2326
+ num_tokens,
2327
+ retention_interval,
2328
+ replay_boundaries=replay_boundaries,
2329
+ )
2330
+
2331
+ def publish_blocks(
2332
+ self,
2333
+ request: Request,
2334
+ num_tokens: int,
2335
+ retention_interval: int | None = None,
2336
+ *,
2337
+ replay_boundaries: Sequence[int],
2338
+ ) -> None:
2339
+ super().cache_blocks(
2340
+ request,
2341
+ num_tokens,
2342
+ retention_interval=retention_interval,
2343
+ replay_boundaries=replay_boundaries,
2344
+ )
2345
+
2346
+ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
2347
+ assert self.coordinator is not None
2348
+ self.coordinator.free(request_id)
2349
+ return super().pop_blocks_for_free(request_id)
2350
+
2351
+
2352
+class _HiSparseAuxiliaryManager(SingleTypeKVCacheManager):
2353
+ """Base for ephemeral groups whose host source owns prefix caching."""
2354
+
2355
+ coordinator: "HiSparseCoordinator | None" = None
2356
+
2357
+ def __init__(self, kv_cache_spec: KVCacheSpec, **kwargs) -> None:
2358
+ # Never prefix-cached, but the per-step ``cache_blocks`` hook is where
2359
+ # residency work runs, so stay opted in regardless of prefix caching.
2360
+ kwargs["enable_caching"] = True
2361
+ super().__init__(kv_cache_spec, **kwargs)
2362
+
2363
+ def cache_blocks(
2364
+ self,
2365
+ request: Request,
2366
+ num_tokens: int,
2367
+ retention_interval: int | None = None,
2368
+ *,
2369
+ replay_boundaries: Sequence[int],
2370
+ ) -> None:
2371
+ return None
2372
+
2373
+ def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
2374
+ return 0
2375
+
2376
+ @classmethod
2377
+ def find_longest_cache_hit(
2378
+ cls,
2379
+ block_hashes: BlockHashList,
2380
+ max_length: int,
2381
+ kv_cache_group_ids: list[int],
2382
+ block_pool: BlockPool,
2383
+ kv_cache_spec: KVCacheSpec,
2384
+ drop_eagle_block: bool,
2385
+ alignment_tokens: int,
2386
+ dcp_world_size: int = 1,
2387
+ pcp_world_size: int = 1,
2388
+ ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
2389
+ return tuple([] for _ in kv_cache_group_ids), 0
2390
+
2391
+
2392
+class HiSparseHotManager(_HiSparseAuxiliaryManager):
2393
+ """Allocate a hot region only after a request acquires CPU-only history."""
2394
+
2395
+ def __init__(self, kv_cache_spec: HiSparseHotSpec, **kwargs) -> None:
2396
+ super().__init__(kv_cache_spec, **kwargs)
2397
+ self.blocks_per_request = kv_cache_spec.blocks_per_request
2398
+ self.hot_required: set[str] = set()
2399
+
2400
+ def require_hot(self, request_id: str) -> None:
2401
+ self.hot_required.add(request_id)
2402
+
2403
+ def has_hot(self, request_id: str) -> bool:
2404
+ return len(self.req_to_blocks.get(request_id, ())) == self.blocks_per_request
2405
+
2406
+ def get_num_blocks_to_allocate(
2407
+ self,
2408
+ request_id: str,
2409
+ num_tokens: int,
2410
+ new_computed_blocks: Sequence[KVCacheBlock],
2411
+ total_computed_tokens: int,
2412
+ num_local_computed_tokens: int,
2413
+ num_tokens_main_model: int,
2414
+ apply_admission_cap: bool = False,
2415
+ ) -> int:
2416
+ assert not new_computed_blocks
2417
+ # A hot region is needed to read host-backed history: an external
2418
+ # import, a new request resuming a host prefix, or one already asked
2419
+ # to transition. Running requests keep their earlier answer.
2420
+ host_import = total_computed_tokens > num_local_computed_tokens
2421
+ resumes_host_prefix = (
2422
+ num_local_computed_tokens > 0 and request_id not in self.num_cached_block
2423
+ )
2424
+ if host_import or resumes_host_prefix or request_id in self.hot_required:
2425
+ return self.get_num_required_blocks(request_id)
2426
+ return 0
2427
+
2428
+ def get_num_required_blocks(self, request_id: str) -> int:
2429
+ return max(
2430
+ self.blocks_per_request - len(self.req_to_blocks.get(request_id, ())),
2431
+ 0,
2432
+ )
2433
+
2434
+ def add_local_computed_blocks(
2435
+ self,
2436
+ request_id: str,
2437
+ new_computed_blocks: Sequence[KVCacheBlock],
2438
+ num_local_computed_tokens: int,
2439
+ num_external_computed_tokens: int,
2440
+ ) -> None:
2441
+ assert not new_computed_blocks
2442
+ self.num_cached_block[request_id] = 0
2443
+ assert self.coordinator is not None
2444
+ if num_local_computed_tokens > 0 or not self.coordinator.resident_managers:
2445
+ self.require_hot(request_id)
2446
+
2447
+ def allocate_external_computed_blocks(
2448
+ self,
2449
+ request_id: str,
2450
+ num_local_computed_tokens: int,
2451
+ num_external_computed_tokens: int,
2452
+ ) -> None:
2453
+ self.require_hot(request_id)
2454
+
2455
+ def allocate_new_blocks(
2456
+ self, request_id: str, num_tokens: int, num_tokens_main_model: int
2457
+ ) -> list[KVCacheBlock]:
2458
+ # Cold admissions can bypass add_local_computed_blocks.
2459
+ self.num_cached_block[request_id] = 0
2460
+ if request_id not in self.hot_required:
2461
+ return []
2462
+ req_blocks = self.req_to_blocks[request_id]
2463
+ num_new_blocks = self.blocks_per_request - len(req_blocks)
2464
+ if num_new_blocks <= 0:
2465
+ return []
2466
+ new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
2467
+ req_blocks.extend(new_blocks)
2468
+ return new_blocks
2469
+
2470
+ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
2471
+ self.hot_required.discard(request_id)
2472
+ return super().pop_blocks_for_free(request_id)
2473
+
2474
+
2475
+class HiSparseResidentManager(_HiSparseAuxiliaryManager):
2476
+ """Track GPU-resident pages for otherwise host-backed KV."""
2477
+
2478
+ def get_num_blocks_to_allocate(
2479
+ self,
2480
+ request_id: str,
2481
+ num_tokens: int,
2482
+ new_computed_blocks: Sequence[KVCacheBlock],
2483
+ total_computed_tokens: int,
2484
+ num_local_computed_tokens: int,
2485
+ num_tokens_main_model: int,
2486
+ apply_admission_cap: bool = False,
2487
+ ) -> int:
2488
+ del num_tokens_main_model
2489
+ assert not new_computed_blocks
2490
+ if total_computed_tokens > num_local_computed_tokens:
2491
+ if total_computed_tokens % self.block_size != 0:
2492
+ raise ValueError(
2493
+ "A host-only HiSparse import must end on a cache-block boundary."
2494
+ )
2495
+ imported_pages = total_computed_tokens // self.block_size
2496
+ return max(cdiv(num_tokens, self.block_size) - imported_pages, 0)
2497
+ existing = len(self.req_to_blocks.get(request_id, ()))
2498
+ host_pages = cdiv(num_local_computed_tokens, self.block_size)
2499
+ required = cdiv(num_tokens, self.block_size)
2500
+ if apply_admission_cap:
2501
+ assert self._max_admission_blocks_per_request is not None
2502
+ required = min(required, self._max_admission_blocks_per_request)
2503
+ return max(required - max(existing, host_pages), 0)
2504
+
2505
+ def allocate_external_computed_blocks(
2506
+ self,
2507
+ request_id: str,
2508
+ num_local_computed_tokens: int,
2509
+ num_external_computed_tokens: int,
2510
+ ) -> None:
2511
+ """Represent imported host history with null resident pages."""
2512
+ assert num_external_computed_tokens > 0
2513
+ num_tokens = num_local_computed_tokens + num_external_computed_tokens
2514
+ blocks = self.req_to_blocks[request_id]
2515
+ tail_page = cdiv(num_tokens, self.block_size) - 1
2516
+ if len(blocks) <= tail_page:
2517
+ blocks.extend([self._null_block] * (tail_page + 1 - len(blocks)))
2518
+
2519
+ def add_local_computed_blocks(
2520
+ self,
2521
+ request_id: str,
2522
+ new_computed_blocks: Sequence[KVCacheBlock],
2523
+ num_local_computed_tokens: int,
2524
+ num_external_computed_tokens: int,
2525
+ ) -> None:
2526
+ assert not new_computed_blocks
2527
+ req_blocks = self.req_to_blocks[request_id]
2528
+ assert not req_blocks
2529
+ num_host_pages = cdiv(num_local_computed_tokens, self.block_size)
2530
+ req_blocks.extend([self._null_block] * num_host_pages)
2531
+ self.num_cached_block[request_id] = 0
2532
+ assert self.coordinator is not None
2533
+ self.coordinator.commit_computed_blocks(request_id, num_host_pages)
2534
+
2535
+ def cache_blocks(
2536
+ self,
2537
+ request: Request,
2538
+ num_tokens: int,
2539
+ retention_interval: int | None = None,
2540
+ *,
2541
+ replay_boundaries: Sequence[int],
2542
+ ) -> None:
2543
+ assert self.coordinator is not None
2544
+ self.coordinator.plan_prefix_materialization(request.request_id, num_tokens)
2545
+ self.coordinator.update_residency(request.request_id)
2546
+
2547
+ def allocate_new_blocks(
2548
+ self, request_id: str, num_tokens: int, num_tokens_main_model: int
2549
+ ) -> list[KVCacheBlock]:
2550
+ del num_tokens_main_model
2551
+ req_blocks = self.req_to_blocks[request_id]
2552
+ num_new_blocks = cdiv(num_tokens, self.block_size) - len(req_blocks)
2553
+ if num_new_blocks <= 0:
2554
+ return []
2555
+ new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
2556
+ req_blocks.extend(new_blocks)
2557
+ return new_blocks
2558
+
2559
+ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
2560
+ assert self.coordinator is not None
2561
+ self.coordinator.free(request_id)
2562
+ return super().pop_blocks_for_free(request_id)
2563
+
2564
+ def adopt_resident_page(
2565
+ self, request_id: str, block_idx: int, block: KVCacheBlock
2566
+ ) -> bool:
2567
+ """Point a null prefix page at a pinned GPU copy of its contents."""
2568
+ blocks = self.req_to_blocks.get(request_id)
2569
+ if blocks is None or block_idx >= len(blocks) or not blocks[block_idx].is_null:
2570
+ return False
2571
+ blocks[block_idx] = block
2572
+ return True
2573
+
2574
+ def get_resident_page(self, request_id: str, block_idx: int) -> KVCacheBlock | None:
2575
+ blocks = self.req_to_blocks.get(request_id)
2576
+ if blocks is None or block_idx >= len(blocks):
2577
+ return None
2578
+ block = blocks[block_idx]
2579
+ return None if block.is_null else block
2580
+
2581
+
2582
def get_manager_for_kv_cache_spec(
2583
kv_cache_spec: KVCacheSpec,
2584
max_in_flight_tokens: int,
2585
max_model_len: int,
2586
+ role: str | None = None,
2587
**kwargs,
2588
) -> SingleTypeKVCacheManager:
2589
"""
@@ -2212,7 +2601,7 @@ def get_manager_for_kv_cache_spec(
2601
Returns:
2602
An instance of the appropriate SingleTypeKVCacheManager subclass
2603
"""
2215
- manager_class = KVCacheSpecRegistry.get_manager_class(kv_cache_spec)
2604
+ manager_class = KVCacheSpecRegistry.get_manager_class(kv_cache_spec, role)
2605
assert manager_class is not None, (
2606
f"No manager registered for KVCacheSpec {type(kv_cache_spec)}"
2607
)
@@ -2224,7 +2613,7 @@ def get_manager_for_kv_cache_spec(
2613
# FullAttentionSpec sizing without a separate admission cap.
2614
if isinstance(
2615
kv_cache_spec,
2227
- (SlidingWindowSpec, ChunkedLocalAttentionSpec),
2616
+ (SlidingWindowSpec, ChunkedLocalAttentionSpec, HiSparseResidentSpec),
2617
):
2618
kwargs["max_admission_blocks_per_request"] = (
2619
kv_cache_spec.max_admission_blocks_per_request(
@@ -2238,11 +2627,24 @@ def get_manager_for_kv_cache_spec(
2627
2628
def register_all_kvcache_specs(vllm_config):
2629
"""Built-in spec registration"""
2630
+ KVCacheSpecRegistry.register_role_manager(
2631
+ KVCacheGroupRole.HISPARSE_SOURCE, HiSparseSourceManager
2632
+ )
2633
KVCacheSpecRegistry.register(
2634
FullAttentionSpec,
2635
FullAttentionManager,
2636
uniform_type_base_spec=FullAttentionSpec,
2637
)
2638
+ KVCacheSpecRegistry.register(
2639
+ HiSparseHotSpec,
2640
+ HiSparseHotManager,
2641
+ uniform_type_base_spec=HiSparseHotSpec,
2642
+ )
2643
+ KVCacheSpecRegistry.register(
2644
+ HiSparseResidentSpec,
2645
+ HiSparseResidentManager,
2646
+ uniform_type_base_spec=HiSparseResidentSpec,
2647
+ )
2648
2649
KVCacheSpecRegistry.register(
2650
SlidingWindowSpec,
vllm/v1/hisparse/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/v1/hisparse/binding.py
new
+277
@@ -0,0 +1,277 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Worker-side wiring of HiSparse caches to attention layers."""
4
+
5
+from collections.abc import Mapping
6
+from copy import copy
7
+from typing import TYPE_CHECKING, Any
8
+
9
+import torch
10
+
11
+from vllm.config import VllmConfig
12
+from vllm.v1.hisparse.layout import (
13
+ HISPARSE_HOT_SUFFIX,
14
+ HISPARSE_RESIDENT_SUFFIX,
15
+)
16
+from vllm.v1.hisparse.runtime import (
17
+ HiSparseCacheHandle,
18
+ HiSparseHostPool,
19
+ initialize_hisparse_runtime_buffers,
20
+ release_pinned_state,
21
+)
22
+from vllm.v1.kv_cache_interface import (
23
+ HiSparseHotSpec,
24
+ HiSparseResidentSpec,
25
+ KVCacheConfig,
26
+ KVCacheLayout,
27
+ KVCacheSpec,
28
+ MLAAttentionSpec,
29
+ UniformTypeKVCacheSpecs,
30
+ create_kv_cache_views,
31
+)
32
+from vllm.v1.worker.utils import allocate_kv_cache, select_common_block_size
33
+
34
+if TYPE_CHECKING:
35
+ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
36
+ from vllm.v1.worker.gpu.block_table import BlockTables
37
+
38
+
39
+def resolve_hisparse_block_size(
40
+ vllm_config: VllmConfig,
41
+ kv_cache_spec: dict[str, KVCacheSpec],
42
+ attn_layers: Mapping[str, "AttentionLayerBase"],
43
+) -> None:
44
+ """Resolve a common kernel block size in-place for HiSparse MLA specs."""
45
+ if vllm_config.attention_config.hisparse_config is None:
46
+ return
47
+ mla_specs = {
48
+ name: spec
49
+ for name, spec in kv_cache_spec.items()
50
+ if isinstance(spec, MLAAttentionSpec)
51
+ }
52
+ if not mla_specs:
53
+ return
54
+ block_sizes = {spec.block_size for spec in mla_specs.values()}
55
+ if len(block_sizes) != 1:
56
+ raise ValueError("HiSparse requires one scheduler block size.")
57
+ backends = [attn_layers[name].get_attn_backend() for name in mla_specs]
58
+ try:
59
+ block_size = select_common_block_size(block_sizes.pop(), backends)
60
+ except ValueError as error:
61
+ raise ValueError(
62
+ "HiSparse requires a GPU block size supported by every sparse "
63
+ f"attention and indexer backend: {error}"
64
+ ) from error
65
+ kv_cache_spec.update(
66
+ (name, spec.copy_with_new_block_size(block_size))
67
+ for name, spec in mla_specs.items()
68
+ )
69
+
70
+
71
+def allocate_hisparse_kv_caches(
72
+ kv_cache_config: KVCacheConfig,
73
+ device: torch.device,
74
+ layout: KVCacheLayout,
75
+ kernel_block_sizes: list[int],
76
+ host_pool: HiSparseHostPool,
77
+) -> dict[str, torch.Tensor]:
78
+ """Allocate the host pool separately from the shared device backing."""
79
+ device_config = copy(kv_cache_config)
80
+ device_config.kv_cache_tensors = [
81
+ tensor
82
+ for tensor in kv_cache_config.kv_cache_tensors
83
+ if not tensor.host_resident
84
+ ]
85
+ kv_caches = allocate_kv_cache(device_config, device, layout, kernel_block_sizes)
86
+ host_tensors = [
87
+ tensor for tensor in kv_cache_config.kv_cache_tensors if tensor.host_resident
88
+ ]
89
+ (host_size,) = {tensor.size for tensor in host_tensors}
90
+ backing = host_pool.allocate(host_size)
91
+ (host_group_id,) = kv_cache_config.host_group_ids
92
+ host_spec = kv_cache_config.kv_cache_groups[host_group_id].kv_cache_spec
93
+ for tensor in host_tensors:
94
+ spec = (
95
+ host_spec.kv_cache_specs[tensor.layers[0]]
96
+ if isinstance(host_spec, UniformTypeKVCacheSpecs)
97
+ else host_spec
98
+ )
99
+ kernel_block_size = kernel_block_sizes[host_group_id]
100
+ if isinstance(spec, MLAAttentionSpec) and spec.storage_block_size is not None:
101
+ kernel_block_size = spec.storage_block_size
102
+ views = create_kv_cache_views(
103
+ backing,
104
+ spec,
105
+ kv_cache_config.num_blocks_of(tensor),
106
+ layout,
107
+ tensor,
108
+ kernel_block_size=kernel_block_size,
109
+ )
110
+ kv_caches.update(zip(tensor.layers, views))
111
+ return kv_caches
112
+
113
+
114
+def init_hisparse_kv_cache(
115
+ kv_cache_config: KVCacheConfig,
116
+ device: torch.device,
117
+ kernel_block_sizes: list[int],
118
+ vllm_config: VllmConfig,
119
+ forward_context: dict[str, Any],
120
+ block_tables: "BlockTables",
121
+) -> dict[str, torch.Tensor]:
122
+ """Allocate and bind HiSparse caches within the caller's allocation context."""
123
+ host_pool = HiSparseHostPool()
124
+ kv_caches = allocate_hisparse_kv_caches(
125
+ kv_cache_config,
126
+ device,
127
+ vllm_config.cache_config.get_resolved_kv_cache_layout(),
128
+ kernel_block_sizes,
129
+ host_pool,
130
+ )
131
+ cache_handles = bind_hisparse_kv_caches(
132
+ forward_context=forward_context,
133
+ kv_cache_config=kv_cache_config,
134
+ kv_caches=kv_caches,
135
+ block_tables=block_tables,
136
+ host_pool=host_pool,
137
+ )
138
+ initialize_hisparse_runtime_buffers(
139
+ cache_handles,
140
+ max_num_reqs=vllm_config.scheduler_config.max_num_seqs,
141
+ max_num_batched_tokens=vllm_config.scheduler_config.max_num_batched_tokens,
142
+ )
143
+ return kv_caches
144
+
145
+
146
+def _get_hisparse_cache(
147
+ forward_context: dict[str, Any], layer_name: str
148
+) -> HiSparseCacheHandle:
149
+ attention_layer = forward_context[layer_name]
150
+ hisparse_cache = attention_layer.hisparse_cache
151
+ assert hisparse_cache is not None
152
+ return hisparse_cache
153
+
154
+
155
+def release_hisparse_profiling_cache(forward_context: dict[str, Any]) -> None:
156
+ cache_handles = [
157
+ cache
158
+ for layer in forward_context.values()
159
+ if (cache := getattr(layer, "hisparse_cache", None)) is not None
160
+ ]
161
+ runtimes = {
162
+ id(cache.runtime): cache.runtime
163
+ for cache in cache_handles
164
+ if hasattr(cache.runtime, "_host_cache")
165
+ }
166
+ if not runtimes:
167
+ return
168
+
169
+ registered_pools = list(
170
+ {
171
+ runtime.registered_host_pool.data_ptr(): runtime.registered_host_pool
172
+ for runtime in runtimes.values()
173
+ }.values()
174
+ )
175
+ release_pinned_state(list(runtimes.values()), registered_pools)
176
+ for cache in cache_handles:
177
+ cache.mirror_staging_cache = None
178
+ cache.mirror_staging_slots = None
179
+
180
+
181
+def bind_hisparse_kv_caches(
182
+ *,
183
+ forward_context: dict[str, Any],
184
+ kv_cache_config: KVCacheConfig,
185
+ kv_caches: dict[str, torch.Tensor],
186
+ block_tables: "BlockTables",
187
+ host_pool: HiSparseHostPool,
188
+) -> list[HiSparseCacheHandle]:
189
+ """Bind existing cache storage and block tables; return the bound handles."""
190
+ assert host_pool.registered is not None
191
+ tensor_configs = {
192
+ name: tensor_config
193
+ for tensor_config in kv_cache_config.kv_cache_tensors
194
+ for name in tensor_config.layers
195
+ }
196
+ resident_source_index = 0
197
+ for group_id, group in enumerate(kv_cache_config.kv_cache_groups):
198
+ if not isinstance(group.kv_cache_spec, HiSparseResidentSpec):
199
+ continue
200
+ for cache_name in group.layer_names:
201
+ assert cache_name.endswith(HISPARSE_RESIDENT_SUFFIX)
202
+ layer_name = cache_name[: -len(HISPARSE_RESIDENT_SUFFIX)]
203
+ tensor_config = tensor_configs[cache_name]
204
+ assert not tensor_config.host_resident
205
+ cache_handle = _get_hisparse_cache(forward_context, layer_name)
206
+ cache_handle.bind_cache(
207
+ kv_caches[cache_name],
208
+ byte_offset=tensor_config.offset,
209
+ block_stride=tensor_config.block_stride,
210
+ num_blocks=kv_cache_config.num_blocks,
211
+ block_size=group.kv_cache_spec.block_size,
212
+ block_table=block_tables.input_block_tables[group_id],
213
+ slot_mapping=block_tables.slot_mappings[group_id],
214
+ )
215
+ assert cache_handle.view is not None
216
+ kv_caches[cache_name] = cache_handle.view.cache
217
+ cache_handle.runtime.resident_source_index = resident_source_index
218
+ resident_source_index += 1
219
+
220
+ hot_backing: torch.Tensor | None = None
221
+ cache_handles: list[HiSparseCacheHandle] = []
222
+ for group_id, group in enumerate(kv_cache_config.kv_cache_groups):
223
+ if not isinstance(group.kv_cache_spec, HiSparseHotSpec):
224
+ continue
225
+ for cache_name in group.layer_names:
226
+ assert cache_name.endswith(HISPARSE_HOT_SUFFIX)
227
+ raw_tensor = kv_caches[cache_name]
228
+ if hot_backing is None:
229
+ hot_backing = raw_tensor
230
+ elif hot_backing.untyped_storage().data_ptr() != (
231
+ raw_tensor.untyped_storage().data_ptr()
232
+ ):
233
+ raise RuntimeError("HiSparse hot tensors must share one GPU backing.")
234
+ layer_name = cache_name[: -len(HISPARSE_HOT_SUFFIX)]
235
+ cache_handle = _get_hisparse_cache(forward_context, layer_name)
236
+ tensor_config = tensor_configs[cache_name]
237
+ assert not tensor_config.host_resident
238
+ cache_handle.runtime.bind_hot_cache(
239
+ raw_tensor,
240
+ byte_offset=tensor_config.offset,
241
+ block_stride=tensor_config.block_stride,
242
+ num_blocks=kv_cache_config.num_blocks,
243
+ block_size=group.kv_cache_spec.block_size,
244
+ block_table=block_tables.input_block_tables[group_id],
245
+ )
246
+ resident = cache_handle.view
247
+ hot = cache_handle.runtime.hot
248
+ assert resident is not None
249
+ if (
250
+ resident.cache.untyped_storage().data_ptr()
251
+ != hot.cache.untyped_storage().data_ptr()
252
+ or resident.cache.data_ptr() != hot.cache.data_ptr()
253
+ or resident.cache.stride() != hot.cache.stride()
254
+ ):
255
+ raise RuntimeError("HiSparse resident and hot layouts must match.")
256
+ source_cache = kv_caches[layer_name]
257
+ assert source_cache.untyped_storage().data_ptr() == (
258
+ host_pool.registered.untyped_storage().data_ptr()
259
+ )
260
+ cache_handle.runtime.bind_source_cache(
261
+ source_cache,
262
+ registered_host_pool=host_pool.registered,
263
+ )
264
+ cache_handles.append(cache_handle)
265
+ # The hot slab is raw storage owned by the runtime, not a layer
266
+ # cache: nothing downstream binds or registers it.
267
+ del kv_caches[cache_name]
268
+
269
+ if hot_backing is None or not cache_handles:
270
+ raise RuntimeError("HiSparse found no hot-cache handles.")
271
+ (source_group_id,) = kv_cache_config.host_group_ids
272
+ for cache_handle in cache_handles:
273
+ cache_handle.source_block_table = block_tables.input_block_tables[
274
+ source_group_id
275
+ ]
276
+ cache_handle.mirror_slot_mapping = block_tables.slot_mappings[source_group_id]
277
+ return cache_handles
vllm/v1/hisparse/block_pool.py
new
+27
@@ -0,0 +1,27 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from vllm.distributed.kv_events import KVCacheEvent
5
+from vllm.v1.core.block_pool import BlockPool
6
+
7
+
8
+class SharedEventQueueBlockPool(BlockPool):
9
+ """A pool that publishes into another pool's live KV event queue.
10
+
11
+ The owner rebinds its queue on every drain, so this reads it through the
12
+ owner rather than holding a reference.
13
+ """
14
+
15
+ def __init__(self, *args, event_owner: BlockPool, **kwargs) -> None:
16
+ self._event_owner = event_owner
17
+ super().__init__(*args, **kwargs)
18
+
19
+ @property
20
+ def kv_event_queue(self) -> list[KVCacheEvent]:
21
+ return self._event_owner.kv_event_queue
22
+
23
+ @kv_event_queue.setter
24
+ def kv_event_queue(self, events: list[KVCacheEvent]) -> None:
25
+ # ``BlockPool.__init__`` seeds an empty queue and ``take_events``
26
+ # swaps in a fresh one; both belong to the owner, which drains it.
27
+ assert not events
vllm/v1/hisparse/coordinator.py
new
+752
@@ -0,0 +1,752 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections.abc import Iterable, Mapping, Sequence
5
+from dataclasses import dataclass, field
6
+from typing import TYPE_CHECKING
7
+
8
+from vllm.utils.math_utils import cdiv
9
+from vllm.v1.core.block_pool import BlockPool
10
+from vllm.v1.core.kv_cache_utils import (
11
+ BlockHashWithGroupId,
12
+ KVCacheBlock,
13
+ KVCacheBlockCopy,
14
+)
15
+from vllm.v1.core.single_type_kv_cache_manager import (
16
+ HiSparseHotManager,
17
+ HiSparseResidentManager,
18
+ HiSparseSourceManager,
19
+ SingleTypeKVCacheManager,
20
+)
21
+from vllm.v1.hisparse.types import (
22
+ SparseKVOffloadCommand,
23
+ SparseKVPageTransfer,
24
+ SparseKVRowMirror,
25
+)
26
+from vllm.v1.kv_cache_interface import (
27
+ HiSparseResidentSpec,
28
+ KVCacheConfig,
29
+)
30
+from vllm.v1.request import Request
31
+
32
+if TYPE_CHECKING:
33
+ from vllm.v1.core.kv_cache_manager import KVCacheManager
34
+
35
+# Sealed pages this many positions behind the block-table tail stay pinned so
36
+# a page written by an in-flight step is never handed out under it.
37
+_ACTIVE_TAIL_PAGES = 2
38
+
39
+
40
+@dataclass
41
+class _PendingPublication:
42
+ request: Request
43
+ num_computed_tokens: int
44
+ num_pages: int
45
+ retention_interval: int | None
46
+ replay_boundaries: Sequence[int]
47
+
48
+
49
+@dataclass
50
+class _HiSparseRequestState:
51
+ """Residency bookkeeping for one request.
52
+
53
+ A resident page moves through three states: dirty (only the GPU copy
54
+ exists), clean (``valid_pages``: the host copy is complete) and unpinned
55
+ (``unpinned_pages``: clean, and its allocation reference was released so
56
+ the pool may reuse it). ``pinned_clean`` tracks clean pages still holding
57
+ their reference, which only happens before the request can read from host.
58
+ """
59
+
60
+ valid_pages: set[int] = field(default_factory=set)
61
+ ready_prefix_pages: int = 0
62
+ pending_pages: dict[int, int] = field(default_factory=dict)
63
+ publication: _PendingPublication | None = None
64
+ copies_recorded_blocks: int = 0
65
+ pinned_clean: set[int] = field(default_factory=set)
66
+ unpinned_pages: set[int] = field(default_factory=set)
67
+
68
+
69
+@dataclass
70
+class _PendingSpill:
71
+ transfer_id: int
72
+ page: tuple[str, int]
73
+ request_state: _HiSparseRequestState
74
+ host_block: KVCacheBlock
75
+ resident_blocks: tuple[KVCacheBlock, ...]
76
+ expected_worker_completions: int = 0
77
+ worker_completions: int = 0
78
+ enqueue_applied: bool = False
79
+
80
+
81
+class HiSparseCoordinator:
82
+ """Own HiSparse host allocation, host publication, spills, and GPU residency.
83
+
84
+ GPU-resident pages are a write-back cache of the host tier. Once a page's
85
+ host copy is durable and its owner can read from host (it has, or has
86
+ asked for, a hot region), the page's allocation reference is released with
87
+ ``BlockPool.unpin_blocks``: the block stays in the block table and keeps
88
+ being read, but the pool counts it as free and may hand it out, at which
89
+ point the owner's page is nulled and its block table republished.
90
+ """
91
+
92
+ def __init__(
93
+ self,
94
+ kv_cache_config: KVCacheConfig,
95
+ managers: tuple[SingleTypeKVCacheManager, ...],
96
+ max_model_len: int,
97
+ ) -> None:
98
+ self.managers = managers
99
+ self.max_model_len = max_model_len
100
+ groups = kv_cache_config.kv_cache_groups
101
+
102
+ resident_managers: list[HiSparseResidentManager] = []
103
+ hot_managers: list[HiSparseHotManager] = []
104
+ resident_group_ids: list[int] = []
105
+ for group_id, (group, manager) in enumerate(zip(groups, managers)):
106
+ if isinstance(group.kv_cache_spec, HiSparseResidentSpec):
107
+ if not isinstance(manager, HiSparseResidentManager):
108
+ raise TypeError(
109
+ "HiSparse resident specs require resident cache managers."
110
+ )
111
+ resident_managers.append(manager)
112
+ resident_group_ids.append(group_id)
113
+ assert not group.host_resident
114
+ if isinstance(manager, HiSparseHotManager):
115
+ hot_managers.append(manager)
116
+ assert not group.host_resident
117
+ if isinstance(manager, (HiSparseHotManager, HiSparseResidentManager)):
118
+ manager.coordinator = self
119
+ self.resident_managers = tuple(resident_managers)
120
+ self.hot_managers = tuple(hot_managers)
121
+ self.host_manager: HiSparseSourceManager | None = None
122
+ self.host_group_id: int | None = None
123
+ self.max_spill_pages = 0
124
+ for group_id, manager in enumerate(managers):
125
+ if isinstance(manager, HiSparseSourceManager):
126
+ if self.host_manager is not None:
127
+ raise ValueError("Only one HiSparse host group is supported.")
128
+ num_host_blocks = kv_cache_config.hisparse_host_num_blocks
129
+ if num_host_blocks is None:
130
+ raise ValueError("HiSparse host group needs host capacity.")
131
+ self.host_group_id = group_id
132
+ self.host_manager = manager
133
+ manager.coordinator = self
134
+ manager.bind_host_pool(num_host_blocks)
135
+ self.has_host_cache = self.host_manager is not None
136
+ self.gpu_pool: BlockPool | None = None
137
+ self.transition_watermark = 0
138
+ if self.resident_managers:
139
+ assert self.host_manager is not None
140
+ assert self.host_group_id is not None
141
+ if self.host_group_id > min(resident_group_ids):
142
+ raise ValueError("HiSparse host group must precede resident groups.")
143
+ resident_block_sizes = {
144
+ manager.block_size for manager in self.resident_managers
145
+ }
146
+ if len(resident_block_sizes) != 1:
147
+ raise ValueError(
148
+ "HiSparse resident cache groups must use one block size."
149
+ )
150
+ resident_block_size = resident_block_sizes.pop()
151
+ if self.host_manager.block_size != resident_block_size:
152
+ raise ValueError("HiSparse host and resident block sizes must match.")
153
+ self.max_spill_pages = max_model_len // resident_block_size
154
+ self.gpu_pool = self.resident_managers[0].block_pool
155
+ # Requests start reading from host once the shared pool runs this
156
+ # low, so admissions find evictable pages rather than pinned ones.
157
+ hot_cost = sum(manager.blocks_per_request for manager in hot_managers)
158
+ self.transition_watermark = max(hot_cost, kv_cache_config.num_blocks // 10)
159
+
160
+ # Host copies handed to the worker, by destination block id, until it
161
+ # reports having run them.
162
+ self._retained_copies: dict[int, tuple[KVCacheBlock, KVCacheBlock]] = {}
163
+ # request -> prefix length an external load is still filling in.
164
+ self._pending_imports: dict[str, int] = {}
165
+ self.block_table_updates: set[str] = set()
166
+ self.spills_to_send: list[SparseKVPageTransfer] = []
167
+ self.pending_spills: dict[int, _PendingSpill] = {}
168
+ self.request_states: dict[str, _HiSparseRequestState] = {}
169
+ self.next_spill_id = 0
170
+ # Published host block hash -> GPU copies of that page, readable until
171
+ # the pool reuses them. Lets a host prefix hit come back GPU-resident.
172
+ self.copies: dict[BlockHashWithGroupId, tuple[KVCacheBlock, ...]] = {}
173
+ self._copy_by_block: dict[int, BlockHashWithGroupId] = {}
174
+ # Unpinned GPU block id -> (request, page) pairs still reading it.
175
+ self._owners: dict[int, set[tuple[str, int]]] = {}
176
+
177
+ def _get_request_state(self, request_id: str) -> _HiSparseRequestState:
178
+ state = self.request_states.get(request_id)
179
+ if state is None:
180
+ state = _HiSparseRequestState()
181
+ self.request_states[request_id] = state
182
+ return state
183
+
184
+ def commit_computed_blocks(self, request_id: str, num_host_pages: int) -> None:
185
+ """Account for a prefix hit on host pages the request now references."""
186
+ if not self.resident_managers or self.host_manager is None:
187
+ return
188
+ host_blocks = self.host_manager.req_to_blocks.get(request_id, ())
189
+ num_host_pages = min(num_host_pages, len(host_blocks))
190
+ if num_host_pages == 0:
191
+ return
192
+ state = self._get_request_state(request_id)
193
+ state.valid_pages.update(range(num_host_pages))
194
+ state.ready_prefix_pages = max(num_host_pages, state.ready_prefix_pages)
195
+ self._adopt_copies(request_id, state, host_blocks[:num_host_pages])
196
+ state.copies_recorded_blocks = max(state.copies_recorded_blocks, num_host_pages)
197
+
198
+ def take_host_block_copies(self) -> tuple[KVCacheBlockCopy, ...]:
199
+ """Drain host copy-on-write work, retaining both endpoints.
200
+
201
+ The worker runs the copies in the step that carries them, so the
202
+ endpoints are released when that step's output comes back.
203
+ """
204
+ if self.host_manager is None:
205
+ return ()
206
+ pairs = self.host_manager.take_host_cow_copies()
207
+ for source_block, cow_block in pairs:
208
+ self._retained_copies[cow_block.block_id] = (source_block, cow_block)
209
+ return tuple(
210
+ KVCacheBlockCopy(
211
+ src_block_id=source_block.block_id,
212
+ dst_block_id=cow_block.block_id,
213
+ )
214
+ for source_block, cow_block in pairs
215
+ )
216
+
217
+ def release_completed_host_copies(self, dst_block_ids: Iterable[int]) -> None:
218
+ """Release the endpoints of host copies the worker reports having run."""
219
+ blocks: list[KVCacheBlock] = []
220
+ for dst_block_id in dst_block_ids:
221
+ pair = self._retained_copies.pop(dst_block_id, None)
222
+ if pair is not None:
223
+ blocks.extend(pair)
224
+ if not blocks:
225
+ return
226
+ assert self.host_manager is not None
227
+ self.host_manager.block_pool.free_blocks(reversed(blocks))
228
+
229
+ def get_host_block_pool(self) -> BlockPool | None:
230
+ manager = self.host_manager
231
+ return manager.block_pool if manager is not None else None
232
+
233
+ # ------------------------------------------------------------------
234
+ # GPU copies of published host pages
235
+ # ------------------------------------------------------------------
236
+
237
+ def _adopt_copies(
238
+ self,
239
+ request_id: str,
240
+ state: _HiSparseRequestState,
241
+ host_blocks: Sequence[KVCacheBlock],
242
+ ) -> None:
243
+ """Point null prefix pages at readable GPU copies of their contents.
244
+
245
+ Runs after the hit length is fully reconciled, so copy hits and misses
246
+ can be scattered per page without affecting the prefix hit.
247
+ """
248
+ if not self.copies:
249
+ return
250
+ for host_idx, host_block in enumerate(host_blocks):
251
+ if host_block.is_null or host_block.block_hash is None:
252
+ continue
253
+ blocks = self.copies.get(host_block.block_hash)
254
+ if blocks is None:
255
+ continue
256
+ for manager, block in zip(self.resident_managers, blocks):
257
+ if manager.adopt_resident_page(request_id, host_idx, block):
258
+ manager.block_pool.touch([block])
259
+ state.pinned_clean.add(host_idx)
260
+
261
+ def _record_copies(self, request_id: str, num_computed_tokens: int) -> None:
262
+ """Index the GPU copies of just-published host blocks for later hits."""
263
+ if not self.resident_managers:
264
+ return
265
+ assert self.host_manager is not None
266
+ host_blocks = self.host_manager.req_to_blocks.get(request_id)
267
+ if not host_blocks:
268
+ return
269
+ state = self._get_request_state(request_id)
270
+ num_host_blocks = min(
271
+ num_computed_tokens // self.host_manager.block_size, len(host_blocks)
272
+ )
273
+ for host_idx in range(state.copies_recorded_blocks, num_host_blocks):
274
+ host_block = host_blocks[host_idx]
275
+ if host_block.is_null or host_block.block_hash is None:
276
+ continue
277
+ blocks: list[KVCacheBlock] = []
278
+ for manager in self.resident_managers:
279
+ block = manager.get_resident_page(request_id, host_idx)
280
+ if block is None:
281
+ break
282
+ blocks.append(block)
283
+ if len(blocks) != len(self.resident_managers):
284
+ continue
285
+ self._drop_copy(host_block.block_hash)
286
+ self.copies[host_block.block_hash] = tuple(blocks)
287
+ for block in blocks:
288
+ self._copy_by_block[block.block_id] = host_block.block_hash
289
+ state.copies_recorded_blocks = max(
290
+ state.copies_recorded_blocks, num_host_blocks
291
+ )
292
+
293
+ def _drop_copy(self, host_hash: BlockHashWithGroupId) -> None:
294
+ blocks = self.copies.pop(host_hash, None)
295
+ if blocks is None:
296
+ return
297
+ for block in blocks:
298
+ self._copy_by_block.pop(block.block_id, None)
299
+
300
+ # ------------------------------------------------------------------
301
+ # Residency: pinning, unpinning and losing pages
302
+ # ------------------------------------------------------------------
303
+
304
+ def _can_read_from_host(self, request_id: str) -> bool:
305
+ """Whether resident pages may be dropped under the request."""
306
+ return bool(self.hot_managers) and all(
307
+ manager.has_hot(request_id) for manager in self.hot_managers
308
+ )
309
+
310
+ def _resident_page_blocks(
311
+ self, request_id: str, page_idx: int
312
+ ) -> list[KVCacheBlock] | None:
313
+ blocks: list[KVCacheBlock] = []
314
+ for manager in self.resident_managers:
315
+ req_blocks = manager.req_to_blocks.get(request_id)
316
+ if (
317
+ req_blocks is None
318
+ or page_idx >= len(req_blocks) - _ACTIVE_TAIL_PAGES
319
+ or req_blocks[page_idx].is_null
320
+ ):
321
+ return None
322
+ blocks.append(req_blocks[page_idx])
323
+ return blocks
324
+
325
+ def _unpin_page(
326
+ self, request_id: str, state: _HiSparseRequestState, page_idx: int
327
+ ) -> bool:
328
+ if page_idx in state.pending_pages or page_idx not in state.valid_pages:
329
+ return False
330
+ blocks = self._resident_page_blocks(request_id, page_idx)
331
+ if blocks is None:
332
+ return False
333
+ for manager, block in zip(self.resident_managers, blocks):
334
+ self._owners.setdefault(block.block_id, set()).add((request_id, page_idx))
335
+ manager.block_pool.unpin_blocks([block], self._on_block_reused)
336
+ state.pinned_clean.discard(page_idx)
337
+ state.unpinned_pages.add(page_idx)
338
+ return True
339
+
340
+ def _unpin_clean_pages(self, request_id: str, state: _HiSparseRequestState) -> None:
341
+ for page_idx in sorted(state.pinned_clean):
342
+ self._unpin_page(request_id, state, page_idx)
343
+
344
+ def _page_became_clean(
345
+ self, request_id: str, state: _HiSparseRequestState, page_idx: int
346
+ ) -> None:
347
+ state.pinned_clean.add(page_idx)
348
+ if self._can_read_from_host(request_id):
349
+ self._unpin_page(request_id, state, page_idx)
350
+
351
+ def _on_block_reused(self, block: KVCacheBlock) -> None:
352
+ host_hash = self._copy_by_block.get(block.block_id)
353
+ if host_hash is not None:
354
+ self._drop_copy(host_hash)
355
+ for request_id, page_idx in self._owners.pop(block.block_id, ()):
356
+ self._lose_page(request_id, page_idx)
357
+
358
+ def _lose_page(self, request_id: str, page_idx: int) -> None:
359
+ state = self.request_states.get(request_id)
360
+ if state is None:
361
+ return
362
+ for manager in self.resident_managers:
363
+ blocks = manager.req_to_blocks.get(request_id)
364
+ if blocks is None or page_idx >= len(blocks) or blocks[page_idx].is_null:
365
+ continue
366
+ block = blocks[page_idx]
367
+ blocks[page_idx] = manager._null_block
368
+ owners = self._owners.get(block.block_id)
369
+ if owners is not None:
370
+ owners.discard((request_id, page_idx))
371
+ state.unpinned_pages.discard(page_idx)
372
+ self.block_table_updates.add(request_id)
373
+ for hot_manager in self.hot_managers:
374
+ hot_manager.require_hot(request_id)
375
+
376
+ def update_residency(self, request_id: str) -> None:
377
+ """Per-step residency policy for a scheduled request.
378
+
379
+ A request that can read from host releases every clean sealed page to
380
+ the pool. One that cannot keeps its pages pinned until the shared pool
381
+ runs low, then asks for a hot region. Pages remain pinned until that region
382
+ is allocated on a subsequent scheduling pass.
383
+ """
384
+ if not self.resident_managers:
385
+ return
386
+ state = self._get_request_state(request_id)
387
+ if not self._can_read_from_host(request_id):
388
+ assert self.gpu_pool is not None
389
+ if self.gpu_pool.get_num_free_blocks() >= self.transition_watermark:
390
+ return
391
+ for manager in self.hot_managers:
392
+ manager.require_hot(request_id)
393
+ return
394
+ self._unpin_clean_pages(request_id, state)
395
+
396
+ # ------------------------------------------------------------------
397
+ # Host publication and spills
398
+ # ------------------------------------------------------------------
399
+
400
+ def plan_prefix_materialization(
401
+ self, request_id: str, num_computed_tokens: int
402
+ ) -> None:
403
+ """Eagerly write back every sealed page that has no host copy yet."""
404
+ if not self.resident_managers:
405
+ return
406
+ assert self.host_manager is not None
407
+ host_block_size = self.host_manager.block_size
408
+ num_pages = num_computed_tokens // host_block_size
409
+ state = self._get_request_state(request_id)
410
+ budget = max(self.max_spill_pages - len(self.spills_to_send), 0)
411
+ for page_idx in range(num_pages):
412
+ if budget == 0:
413
+ break
414
+ if page_idx in state.pending_pages:
415
+ continue
416
+ if page_idx in state.valid_pages:
417
+ continue
418
+ if not all(
419
+ manager.get_resident_page(request_id, page_idx) is not None
420
+ for manager in self.resident_managers
421
+ ):
422
+ continue
423
+ if self._plan_spill(request_id, page_idx, after_forward=True):
424
+ budget -= 1
425
+
426
+ def publish_when_ready(
427
+ self,
428
+ request: Request,
429
+ num_computed_tokens: int,
430
+ retention_interval: int | None,
431
+ *,
432
+ replay_boundaries: Sequence[int],
433
+ ) -> None:
434
+ """Publish host-source hashes only after their pages are durable."""
435
+ manager = self.host_manager
436
+ if manager is None:
437
+ return
438
+ num_pages = num_computed_tokens // manager.block_size
439
+ request_id = request.request_id
440
+ state = self._get_request_state(request_id)
441
+ if state.ready_prefix_pages >= num_pages:
442
+ manager.publish_blocks(
443
+ request,
444
+ num_computed_tokens,
445
+ retention_interval=retention_interval,
446
+ replay_boundaries=replay_boundaries,
447
+ )
448
+ self._record_copies(request_id, num_computed_tokens)
449
+ state.publication = None
450
+ return
451
+ state.publication = _PendingPublication(
452
+ request=request,
453
+ num_computed_tokens=num_computed_tokens,
454
+ num_pages=num_pages,
455
+ retention_interval=retention_interval,
456
+ replay_boundaries=replay_boundaries,
457
+ )
458
+
459
+ def _publish_host_blocks_if_ready(self, request_id: str) -> None:
460
+ state = self.request_states.get(request_id)
461
+ if state is None:
462
+ return
463
+ publication = state.publication
464
+ if publication is None or state.ready_prefix_pages < publication.num_pages:
465
+ return
466
+ assert self.host_manager is not None
467
+ self.host_manager.publish_blocks(
468
+ publication.request,
469
+ publication.num_computed_tokens,
470
+ retention_interval=publication.retention_interval,
471
+ replay_boundaries=publication.replay_boundaries,
472
+ )
473
+ self._record_copies(request_id, publication.num_computed_tokens)
474
+ state.publication = None
475
+
476
+ def record_pending_host_import(self, request_id: str, num_tokens: int) -> None:
477
+ """Note a prefix an external load is populating in host pages."""
478
+ self._pending_imports[request_id] = num_tokens
479
+
480
+ def finish_host_import(self, request_id: str, *, failed: bool) -> None:
481
+ num_tokens = self._pending_imports.pop(request_id, None)
482
+ if num_tokens is not None and not failed:
483
+ self.complete_host_import(request_id, num_tokens)
484
+
485
+ def complete_host_import(self, request_id: str, num_computed_tokens: int) -> None:
486
+ """Publish externally populated host pages after connector completion."""
487
+ if not self.resident_managers:
488
+ return
489
+ block_size = self.resident_managers[0].block_size
490
+ num_pages = num_computed_tokens // block_size
491
+ state = self._get_request_state(request_id)
492
+ state.valid_pages = set(range(num_pages))
493
+ state.ready_prefix_pages = num_pages
494
+ self._publish_host_blocks_if_ready(request_id)
495
+
496
+ def _plan_spill(
497
+ self,
498
+ request_id: str,
499
+ page_idx: int,
500
+ *,
501
+ after_forward: bool,
502
+ ) -> bool:
503
+ assert self.host_manager is not None
504
+ state = self._get_request_state(request_id)
505
+ host_blocks = self.host_manager.req_to_blocks.get(request_id)
506
+ if host_blocks is None or page_idx >= len(host_blocks):
507
+ return False
508
+ host_block = host_blocks[page_idx]
509
+ if host_block.is_null:
510
+ return False
511
+
512
+ blocks: list[KVCacheBlock] = []
513
+ for manager in self.resident_managers:
514
+ block = manager.get_resident_page(request_id, page_idx)
515
+ if block is None:
516
+ return False
517
+ blocks.append(block)
518
+
519
+ self.host_manager.block_pool.touch([host_block])
520
+ for manager, block in zip(self.resident_managers, blocks):
521
+ manager.block_pool.touch([block])
522
+ spill_id = self.next_spill_id
523
+ self.next_spill_id += 1
524
+ plan = SparseKVPageTransfer(
525
+ transfer_id=spill_id,
526
+ destination_block_id=host_block.block_id,
527
+ source_block_ids=tuple(block.block_id for block in blocks),
528
+ after_forward=after_forward,
529
+ )
530
+ self.pending_spills[spill_id] = _PendingSpill(
531
+ transfer_id=spill_id,
532
+ page=(request_id, page_idx),
533
+ request_state=state,
534
+ host_block=host_block,
535
+ resident_blocks=tuple(blocks),
536
+ )
537
+ state.pending_pages[page_idx] = spill_id
538
+ self.spills_to_send.append(plan)
539
+ return True
540
+
541
+ # ------------------------------------------------------------------
542
+ # Worker-facing views
543
+ # ------------------------------------------------------------------
544
+
545
+ def build_row_mirrors(
546
+ self,
547
+ requests: Iterable[tuple[str, int, int]],
548
+ ) -> tuple[SparseKVRowMirror, ...]:
549
+ """Map scheduled rows to scheduler-owned resident and host blocks."""
550
+ if not self.resident_managers or self.host_manager is None:
551
+ return ()
552
+ block_size = self.resident_managers[0].block_size
553
+ mirrors: list[SparseKVRowMirror] = []
554
+ for request_id, num_computed_tokens, num_scheduled_tokens in requests:
555
+ host_blocks = self.host_manager.req_to_blocks.get(request_id)
556
+ resident_blocks = [
557
+ manager.req_to_blocks.get(request_id)
558
+ for manager in self.resident_managers
559
+ ]
560
+ if host_blocks is None or any(blocks is None for blocks in resident_blocks):
561
+ continue
562
+ token_position = num_computed_tokens
563
+ end_position = token_position + num_scheduled_tokens
564
+ while token_position < end_position:
565
+ page_idx, row_offset = divmod(token_position, block_size)
566
+ num_rows = min(block_size - row_offset, end_position - token_position)
567
+ # A page without a GPU copy has no mirror source, and a page
568
+ # without a host block has no destination. Skip just that page:
569
+ # its neighbours in the window are still mirrorable, and a
570
+ # window only ever moves forward, so dropping them here would
571
+ # leave their host rows permanently stale.
572
+ source_starts = []
573
+ for blocks in resident_blocks:
574
+ assert blocks is not None
575
+ if page_idx >= len(blocks) or blocks[page_idx].is_null:
576
+ break
577
+ source_starts.append(
578
+ blocks[page_idx].block_id * block_size + row_offset
579
+ )
580
+ if len(source_starts) != len(resident_blocks):
581
+ token_position += num_rows
582
+ continue
583
+ host_block_idx = page_idx
584
+ if host_block_idx >= len(host_blocks):
585
+ token_position += num_rows
586
+ continue
587
+ host_block = host_blocks[host_block_idx]
588
+ if host_block.is_null:
589
+ token_position += num_rows
590
+ continue
591
+ mirrors.append(
592
+ SparseKVRowMirror(
593
+ source_starts=tuple(source_starts),
594
+ destination_start=(
595
+ host_block.block_id * block_size + row_offset
596
+ ),
597
+ num_rows=num_rows,
598
+ )
599
+ )
600
+ token_position += num_rows
601
+ return tuple(mirrors)
602
+
603
+ def all_context_pages_resident(
604
+ self,
605
+ requests: Iterable[tuple[str, int, int]],
606
+ ) -> bool:
607
+ """Return whether every scheduled request can read only resident KV."""
608
+ if not self.resident_managers:
609
+ return False
610
+ block_size = self.resident_managers[0].block_size
611
+ for request_id, num_computed_tokens, num_scheduled_tokens in requests:
612
+ num_pages = cdiv(num_computed_tokens + num_scheduled_tokens, block_size)
613
+ for page_idx in range(num_pages):
614
+ if any(
615
+ manager.get_resident_page(request_id, page_idx) is None
616
+ for manager in self.resident_managers
617
+ ):
618
+ return False
619
+ return True
620
+
621
+ def take_block_table_updates(self) -> dict[str, tuple[list[int], ...]]:
622
+ updates = {
623
+ request_id: tuple(
624
+ [block.block_id for block in manager.req_to_blocks.get(request_id, [])]
625
+ for manager in self.managers
626
+ )
627
+ for request_id in self.block_table_updates
628
+ }
629
+ self.block_table_updates.clear()
630
+ return updates
631
+
632
+ def build_offload_command(self) -> SparseKVOffloadCommand | None:
633
+ """Package residency decisions for the worker connector."""
634
+ if not self.resident_managers:
635
+ return None
636
+ command = SparseKVOffloadCommand(page_transfers=self.spills_to_send)
637
+ self.spills_to_send = []
638
+ return command
639
+
640
+ def has_pending_reclamation(self) -> bool:
641
+ """Whether an in-flight spill will free GPU blocks on completion."""
642
+ return any(
643
+ pending.page[0] in self.request_states
644
+ and self._can_read_from_host(pending.page[0])
645
+ for pending in self.pending_spills.values()
646
+ )
647
+
648
+ def has_pending_work(self) -> bool:
649
+ return bool(self.spills_to_send or self.pending_spills or self._retained_copies)
650
+
651
+ def update_spills(
652
+ self,
653
+ enqueued_counts: Mapping[int, int],
654
+ completed_counts: Mapping[int, int],
655
+ ) -> None:
656
+ for spill_id, count in enqueued_counts.items():
657
+ pending = self.pending_spills.get(spill_id)
658
+ if pending is not None and pending.expected_worker_completions == 0:
659
+ pending.expected_worker_completions = count
660
+ for spill_id, count in completed_counts.items():
661
+ pending = self.pending_spills.get(spill_id)
662
+ if pending is not None:
663
+ pending.worker_completions += count
664
+
665
+ self._apply_enqueued_spills()
666
+ self._complete_host_writes()
667
+
668
+ def _apply_enqueued_spills(self) -> None:
669
+ """Drop the DMA pins once the worker has enqueued a spill."""
670
+ for pending in self.pending_spills.values():
671
+ if (
672
+ pending.enqueue_applied
673
+ or not pending.expected_worker_completions
674
+ or pending.worker_completions < pending.expected_worker_completions
675
+ ):
676
+ continue
677
+ for manager, block in zip(self.resident_managers, pending.resident_blocks):
678
+ manager.block_pool.free_blocks([block])
679
+ pending.enqueue_applied = True
680
+
681
+ def _complete_host_writes(self) -> None:
682
+ completed = [
683
+ pending
684
+ for pending in self.pending_spills.values()
685
+ if pending.enqueue_applied
686
+ and pending.expected_worker_completions
687
+ and pending.worker_completions >= pending.expected_worker_completions
688
+ ]
689
+ completed_request_ids: set[str] = set()
690
+ for pending in completed:
691
+ self.pending_spills.pop(pending.transfer_id, None)
692
+ request_id, page_idx = pending.page
693
+ pending.request_state.pending_pages.pop(page_idx, None)
694
+ if self.request_states.get(request_id) is pending.request_state:
695
+ pending.request_state.valid_pages.add(page_idx)
696
+ self._page_became_clean(request_id, pending.request_state, page_idx)
697
+ completed_request_ids.add(request_id)
698
+ assert self.host_manager is not None
699
+ self.host_manager.block_pool.free_blocks([pending.host_block])
700
+ for request_id in completed_request_ids:
701
+ state = self.request_states[request_id]
702
+ while state.ready_prefix_pages in state.valid_pages:
703
+ state.ready_prefix_pages += 1
704
+ self._publish_host_blocks_if_ready(request_id)
705
+
706
+ def free(self, request_id: str) -> None:
707
+ """Detach the request; its clean pages stay readable copies in the pool."""
708
+ self._pending_imports.pop(request_id, None)
709
+ state = self.request_states.pop(request_id, None)
710
+ if state is None:
711
+ return
712
+ for manager in self.resident_managers:
713
+ blocks = manager.req_to_blocks.get(request_id)
714
+ if not blocks:
715
+ continue
716
+ for page_idx, block in enumerate(blocks):
717
+ if block.is_null:
718
+ continue
719
+ if page_idx in state.unpinned_pages:
720
+ owners = self._owners.get(block.block_id)
721
+ if owners is not None:
722
+ owners.discard((request_id, page_idx))
723
+ elif (
724
+ page_idx in state.valid_pages
725
+ and page_idx not in state.pending_pages
726
+ ):
727
+ manager.block_pool.unpin_blocks([block], self._on_block_reused)
728
+ else:
729
+ continue
730
+ blocks[page_idx] = manager._null_block
731
+
732
+
733
+def get_hisparse_coordinator(
734
+ kv_cache_manager: "KVCacheManager",
735
+) -> HiSparseCoordinator:
736
+ """Return the coordinator shared by a KV cache manager's HiSparse groups.
737
+
738
+ Built on first call and memoised on the HiSparse managers it binds itself
739
+ to. Must run before the first request is admitted.
740
+ """
741
+ managers = tuple(kv_cache_manager.coordinator.single_type_managers)
742
+ for manager in managers:
743
+ coordinator = getattr(manager, "coordinator", None)
744
+ if coordinator is not None:
745
+ assert isinstance(coordinator, HiSparseCoordinator)
746
+ return coordinator
747
+ coordinator = HiSparseCoordinator(
748
+ kv_cache_manager.kv_cache_config, managers, kv_cache_manager.max_model_len
749
+ )
750
+ if coordinator.host_manager is None:
751
+ raise ValueError("No HiSparse cache group is configured.")
752
+ return coordinator
vllm/v1/hisparse/layout.py
new
+345
@@ -0,0 +1,345 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections import defaultdict
5
+from dataclasses import dataclass
6
+from typing import cast
7
+
8
+from vllm.config import VllmConfig
9
+from vllm.config.kv_transfer import hisparse_host_pool_gib
10
+from vllm.logger import init_logger
11
+from vllm.utils.math_utils import cdiv
12
+from vllm.v1.hisparse.runtime import ResolvedHiSparseConfig
13
+from vllm.v1.kv_cache_interface import (
14
+ HiSparseHotSpec,
15
+ HiSparseResidentSpec,
16
+ KVCacheConfig,
17
+ KVCacheGroupRole,
18
+ KVCacheGroupSpec,
19
+ KVCacheSpec,
20
+ KVCacheTensor,
21
+ MLAAttentionSpec,
22
+ SparseCacheRole,
23
+ UniformTypeKVCacheSpecs,
24
+ compute_layout_strides,
25
+)
26
+from vllm.v1.kv_cache_layout import KVCacheLayout
27
+
28
+logger = init_logger(__name__)
29
+
30
+HISPARSE_HOT_SUFFIX = ".hisparse_hot"
31
+HISPARSE_RESIDENT_SUFFIX = ".hisparse_resident"
32
+
33
+
34
+@dataclass(frozen=True)
35
+class HiSparseLayout:
36
+ source_group: KVCacheGroupSpec
37
+ device_groups: list[KVCacheGroupSpec]
38
+ host_num_blocks: int
39
+
40
+
41
+def get_hisparse_kv_cache_groups(
42
+ vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec]
43
+) -> list[KVCacheGroupSpec] | None:
44
+ attention_config = getattr(vllm_config, "attention_config", None)
45
+ if attention_config is None or attention_config.hisparse_config is None:
46
+ return None
47
+
48
+ mla_specs: dict[str, KVCacheSpec] = {
49
+ name: spec
50
+ for name, spec in kv_cache_spec.items()
51
+ if isinstance(spec, MLAAttentionSpec)
52
+ }
53
+ other_specs = {
54
+ name: spec
55
+ for name, spec in kv_cache_spec.items()
56
+ if not isinstance(spec, MLAAttentionSpec)
57
+ }
58
+ if not mla_specs or not other_specs:
59
+ return None
60
+
61
+ from vllm.v1.core.kv_cache_utils import get_kv_cache_groups
62
+
63
+ mla_group_spec = UniformTypeKVCacheSpecs.from_specs(mla_specs)
64
+ assert mla_group_spec is not None
65
+ mla_group = KVCacheGroupSpec(list(mla_specs), mla_group_spec)
66
+ return [mla_group, *get_kv_cache_groups(vllm_config, other_specs)]
67
+
68
+
69
+def get_hisparse_host_pool_bytes(vllm_config: VllmConfig) -> int:
70
+ host_pool_gib = hisparse_host_pool_gib(vllm_config.kv_transfer_config)
71
+ if host_pool_gib is None:
72
+ raise ValueError("HiSparse requires HiSparseConnector with host_pool_gib")
73
+ return int(host_pool_gib * 2**30)
74
+
75
+
76
+def _partition_hisparse_specs(
77
+ groups: list[KVCacheGroupSpec],
78
+) -> tuple[dict[str, MLAAttentionSpec], dict[str, MLAAttentionSpec]]:
79
+ group_spec = groups[0].kv_cache_spec
80
+ if not isinstance(group_spec, UniformTypeKVCacheSpecs):
81
+ raise ValueError("HiSparse requires uniform sparse-MLA cache specs.")
82
+
83
+ specs = group_spec.kv_cache_specs
84
+ if any(
85
+ isinstance(spec, MLAAttentionSpec) and spec.model_version == "deepseek_v4"
86
+ for spec in specs.values()
87
+ ):
88
+ raise ValueError("HiSparse does not support DeepSeek V4.")
89
+ if not all(isinstance(spec, MLAAttentionSpec) for spec in specs.values()):
90
+ raise ValueError("HiSparse requires its first cache group to contain MLA only.")
91
+
92
+ source_specs = {
93
+ name: spec
94
+ for name, spec in specs.items()
95
+ if isinstance(spec, MLAAttentionSpec)
96
+ and spec.cache_role is SparseCacheRole.SPARSE
97
+ }
98
+ indexer_specs = {
99
+ name: spec
100
+ for name, spec in specs.items()
101
+ if isinstance(spec, MLAAttentionSpec)
102
+ and spec.cache_role is SparseCacheRole.INDEXER
103
+ }
104
+ if not source_specs or not indexer_specs:
105
+ raise ValueError("HiSparse requires sparse-MLA and indexer cache specs.")
106
+ return source_specs, indexer_specs
107
+
108
+
109
+def get_hisparse_gpu_memory_usage(
110
+ vllm_config: VllmConfig,
111
+ kv_cache_groups: list[KVCacheGroupSpec],
112
+) -> int:
113
+ _, indexer_specs = _partition_hisparse_specs(kv_cache_groups)
114
+ return sum(
115
+ spec.max_memory_usage_bytes(vllm_config) for spec in indexer_specs.values()
116
+ ) + sum(
117
+ group.kv_cache_spec.max_memory_usage_bytes(vllm_config)
118
+ for group in kv_cache_groups[1:]
119
+ )
120
+
121
+
122
+def create_hisparse_layout(
123
+ vllm_config: VllmConfig,
124
+ groups: list[KVCacheGroupSpec],
125
+ host_budget: int,
126
+) -> HiSparseLayout:
127
+ source_specs, indexer_specs = _partition_hisparse_specs(groups)
128
+ block_sizes = {
129
+ spec.block_size for spec in (*source_specs.values(), *indexer_specs.values())
130
+ }
131
+ if len(block_sizes) != 1:
132
+ raise ValueError("HiSparse requires one resolved GPU block size.")
133
+ gpu_block_size = block_sizes.pop()
134
+
135
+ config = ResolvedHiSparseConfig.from_vllm_config(
136
+ vllm_config,
137
+ vllm_config.model_config.hf_config.index_topk,
138
+ gpu_block_size,
139
+ )
140
+ assert config is not None
141
+ indexer_group_spec = UniformTypeKVCacheSpecs.from_specs(
142
+ cast(dict[str, KVCacheSpec], indexer_specs)
143
+ )
144
+ assert indexer_group_spec is not None
145
+ indexer_group = KVCacheGroupSpec(
146
+ list(indexer_specs),
147
+ indexer_group_spec,
148
+ enable_kv_transfer=True,
149
+ role=KVCacheGroupRole.HISPARSE_INDEXER,
150
+ )
151
+
152
+ indexer_page = sum(spec.page_size_bytes for spec in indexer_specs.values())
153
+ hot_blocks_per_request = cdiv(config.device_buffer_size, gpu_block_size)
154
+ hot_units: list[list[tuple[str, MLAAttentionSpec]]] = []
155
+ for layer_name, layer_spec in source_specs.items():
156
+ if layer_spec.is_index_group_leader or not hot_units:
157
+ hot_units.append([])
158
+ hot_units[-1].append((f"{layer_name}{HISPARSE_HOT_SUFFIX}", layer_spec))
159
+
160
+ resident_groups: list[KVCacheGroupSpec] = []
161
+ hot_groups: list[KVCacheGroupSpec] = []
162
+
163
+ def append_hot_group(layers: list[tuple[str, MLAAttentionSpec]]) -> None:
164
+ page_sizes = {spec.page_size_bytes for _, spec in layers}
165
+ if len(page_sizes) != 1:
166
+ raise ValueError(
167
+ "HiSparse hot-cache groups require one page size, got "
168
+ f"{sorted(page_sizes)}."
169
+ )
170
+ page_size = page_sizes.pop()
171
+ names = [name for name, _ in layers]
172
+ resident_groups.append(
173
+ KVCacheGroupSpec(
174
+ [
175
+ name[: -len(HISPARSE_HOT_SUFFIX)] + HISPARSE_RESIDENT_SUFFIX
176
+ for name in names
177
+ ],
178
+ HiSparseResidentSpec(
179
+ block_size=gpu_block_size,
180
+ page_size=page_size,
181
+ ),
182
+ enable_kv_transfer=False,
183
+ )
184
+ )
185
+ hot_groups.append(
186
+ KVCacheGroupSpec(
187
+ names,
188
+ HiSparseHotSpec(
189
+ block_size=gpu_block_size,
190
+ page_size=page_size,
191
+ blocks_per_request=hot_blocks_per_request,
192
+ ),
193
+ enable_kv_transfer=False,
194
+ )
195
+ )
196
+
197
+ current: list[tuple[str, MLAAttentionSpec]] = []
198
+ current_page = 0
199
+ for unit in hot_units:
200
+ unit_page = sum(spec.page_size_bytes for _, spec in unit)
201
+ if current and current_page + unit_page > indexer_page:
202
+ append_hot_group(current)
203
+ current = []
204
+ current_page = 0
205
+ current.extend(unit)
206
+ current_page += unit_page
207
+ if current:
208
+ append_hot_group(current)
209
+
210
+ source_group_spec = UniformTypeKVCacheSpecs.from_specs(
211
+ cast(dict[str, KVCacheSpec], source_specs)
212
+ )
213
+ assert source_group_spec is not None
214
+ source_group = KVCacheGroupSpec(
215
+ list(source_specs),
216
+ source_group_spec,
217
+ role=KVCacheGroupRole.HISPARSE_SOURCE,
218
+ host_resident=True,
219
+ enable_kv_transfer=True,
220
+ )
221
+ regular_groups = groups[1:]
222
+ gpu_groups = [indexer_group, *resident_groups, *hot_groups, *regular_groups]
223
+
224
+ host_num_blocks = host_budget // sum(
225
+ spec.page_size_bytes for spec in source_specs.values()
226
+ )
227
+ if host_num_blocks <= 0:
228
+ raise ValueError("HiSparse has no allocatable host blocks.")
229
+
230
+ return HiSparseLayout(
231
+ source_group=source_group,
232
+ device_groups=gpu_groups,
233
+ host_num_blocks=host_num_blocks,
234
+ )
235
+
236
+
237
+def _build_hisparse_kv_cache_tensors(
238
+ kv_cache_groups: list[KVCacheGroupSpec],
239
+ num_blocks: int,
240
+ size: int,
241
+ layout: KVCacheLayout,
242
+ bytes_per_block: int,
243
+ *,
244
+ host_resident: bool = False,
245
+) -> list[KVCacheTensor]:
246
+ interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None
247
+ tensors: list[KVCacheTensor] = []
248
+ for group in kv_cache_groups:
249
+ group_spec = group.kv_cache_spec
250
+ layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list)
251
+ if isinstance(group_spec, UniformTypeKVCacheSpecs):
252
+ for layer_name, spec in group_spec.kv_cache_specs.items():
253
+ layers_by_spec[spec].append(layer_name)
254
+ elif group.layer_names:
255
+ layers_by_spec[group_spec].extend(group.layer_names)
256
+
257
+ byte_offset = 0
258
+ for spec, layer_names in layers_by_spec.items():
259
+ if isinstance(spec, (HiSparseHotSpec, HiSparseResidentSpec)):
260
+ if not layout.is_block_outermost:
261
+ raise ValueError("HiSparse requires a block-outermost KV layout.")
262
+ layer_stride = spec.page_size_bytes
263
+ block_stride = bytes_per_block
264
+ else:
265
+ layer_stride, block_stride, _, _, _ = compute_layout_strides(
266
+ spec,
267
+ num_blocks,
268
+ len(layer_names),
269
+ layout,
270
+ fixed_strides=(None, interleaved_block_stride, None, None, None),
271
+ )
272
+ offset = (
273
+ byte_offset
274
+ * max(layer_stride, spec.page_size_bytes)
275
+ // spec.page_size_bytes
276
+ )
277
+ tensors.append(
278
+ KVCacheTensor(
279
+ size=size,
280
+ layers=layer_names,
281
+ layer_stride=layer_stride,
282
+ block_stride=block_stride,
283
+ offset=offset,
284
+ host_resident=host_resident,
285
+ )
286
+ )
287
+ byte_offset += len(layer_names) * spec.page_size_bytes
288
+ return tensors
289
+
290
+
291
+def get_hisparse_kv_cache_config(
292
+ vllm_config: VllmConfig,
293
+ kv_cache_groups: list[KVCacheGroupSpec],
294
+ available_memory: int,
295
+ host_budget: int,
296
+) -> KVCacheConfig:
297
+ from vllm.v1.core.kv_cache_utils import (
298
+ _get_kv_cache_bytes_per_block,
299
+ may_override_num_blocks,
300
+ validate_kv_cache_layout,
301
+ )
302
+
303
+ hisparse_layout = create_hisparse_layout(vllm_config, kv_cache_groups, host_budget)
304
+ device_groups = hisparse_layout.device_groups
305
+ layout = vllm_config.cache_config.get_resolved_kv_cache_layout()
306
+ validate_kv_cache_layout(layout, device_groups)
307
+ bytes_per_block = _get_kv_cache_bytes_per_block(device_groups)
308
+ num_blocks = may_override_num_blocks(
309
+ vllm_config, available_memory // bytes_per_block
310
+ )
311
+ size = bytes_per_block * num_blocks
312
+ kv_cache_tensors = _build_hisparse_kv_cache_tensors(
313
+ device_groups, num_blocks, size, layout, bytes_per_block
314
+ )
315
+
316
+ host_groups = [hisparse_layout.source_group]
317
+ host_layout = KVCacheLayout.LBNHC
318
+ validate_kv_cache_layout(host_layout, host_groups)
319
+ host_bytes_per_block = _get_kv_cache_bytes_per_block(host_groups)
320
+ host_size = host_bytes_per_block * hisparse_layout.host_num_blocks
321
+ kv_cache_tensors[:0] = _build_hisparse_kv_cache_tensors(
322
+ host_groups,
323
+ hisparse_layout.host_num_blocks,
324
+ host_size,
325
+ host_layout,
326
+ host_bytes_per_block,
327
+ host_resident=True,
328
+ )
329
+ logger.info_once(
330
+ "HiSparse HMA: %.1f GiB host source (%d blocks), %.1f GiB shared "
331
+ "GPU indexer/resident/hot pool (%d blocks).",
332
+ host_size / 2**30,
333
+ hisparse_layout.host_num_blocks,
334
+ size / 2**30,
335
+ num_blocks,
336
+ )
337
+ return KVCacheConfig(
338
+ num_blocks=num_blocks,
339
+ kv_cache_tensors=kv_cache_tensors,
340
+ kv_cache_groups=[*host_groups, *device_groups],
341
+ hisparse_host_num_blocks=hisparse_layout.host_num_blocks,
342
+ prefix_cache_retention_interval=(
343
+ vllm_config.cache_config.prefix_cache_retention_interval
344
+ ),
345
+ )
vllm/v1/hisparse/runtime.py
new
+1100
@@ -0,0 +1,1100 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""NVIDIA HiSparse resident, host, and hot-cache data plane."""
4
+
5
+from __future__ import annotations
6
+
7
+import time
8
+from collections.abc import Callable
9
+from dataclasses import dataclass
10
+from typing import TYPE_CHECKING, Any, TypeAlias
11
+
12
+import psutil
13
+import torch
14
+
15
+from vllm.config import CUDAGraphMode, VllmConfig
16
+from vllm.forward_context import get_forward_context
17
+from vllm.logger import init_logger
18
+from vllm.platforms import current_platform
19
+from vllm.utils.math_utils import round_up
20
+from vllm.utils.torch_utils import current_stream
21
+from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor
22
+
23
+logger = init_logger(__name__)
24
+
25
+if TYPE_CHECKING:
26
+ from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
27
+
28
+# fp8_ds_mla KV row: 512 B quantized NoPE + 16 B scales + 128 B RoPE.
29
+FP8_DS_MLA_ROW_BYTES = 656
30
+HiSparseTopKResult: TypeAlias = torch.Tensor | tuple[torch.Tensor, torch.Tensor]
31
+
32
+
33
+def is_hisparse_decode_batch(attn_metadata: Any | None) -> bool:
34
+ return (
35
+ attn_metadata is not None
36
+ and attn_metadata.num_decode_tokens == attn_metadata.num_actual_tokens
37
+ and attn_metadata.max_query_len == 1
38
+ and attn_metadata.num_reqs == attn_metadata.num_actual_tokens
39
+ )
40
+
41
+
42
+def _get_max_decode_query_len(vllm_config: VllmConfig) -> int:
43
+ speculative_config = getattr(vllm_config, "speculative_config", None)
44
+ if (
45
+ speculative_config is not None
46
+ and speculative_config.num_speculative_tokens is not None
47
+ ):
48
+ return 1 + speculative_config.num_speculative_tokens * (
49
+ 2 if speculative_config.parallel_drafting else 1
50
+ )
51
+ return 1
52
+
53
+
54
+def _get_max_swap_rows(vllm_config: VllmConfig) -> int:
55
+ max_query_len = _get_max_decode_query_len(vllm_config)
56
+ scheduler_config = vllm_config.scheduler_config
57
+ return min(
58
+ scheduler_config.max_num_batched_tokens,
59
+ scheduler_config.max_num_seqs * max_query_len,
60
+ )
61
+
62
+
63
+@dataclass(frozen=True)
64
+class PagedCacheView:
65
+ cache: torch.Tensor
66
+ attention_cache: torch.Tensor
67
+ block_size: int
68
+ attention_block_stride: int
69
+
70
+ @classmethod
71
+ def bind(
72
+ cls,
73
+ raw_tensor: torch.Tensor,
74
+ *,
75
+ dtype: torch.dtype,
76
+ row_width: int,
77
+ byte_offset: int,
78
+ block_stride: int,
79
+ num_blocks: int,
80
+ block_size: int,
81
+ ) -> PagedCacheView:
82
+ itemsize = dtype.itemsize
83
+ assert byte_offset % itemsize == 0 and block_stride % itemsize == 0
84
+ cache = torch.as_strided(
85
+ raw_tensor.view(dtype),
86
+ size=(num_blocks, block_size, row_width),
87
+ stride=(block_stride // itemsize, row_width, 1),
88
+ storage_offset=byte_offset // itemsize,
89
+ )
90
+ row_bytes = row_width * itemsize
91
+ if block_stride % (block_size * row_bytes):
92
+ return cls(cache, cache, block_size, block_size)
93
+ attention_cache = (
94
+ raw_tensor[byte_offset:].view(dtype).view(-1, block_size, row_width)
95
+ )
96
+ return cls(cache, attention_cache, block_size, block_stride // row_bytes)
97
+
98
+
99
+@dataclass(frozen=True)
100
+class ResolvedHiSparseConfig:
101
+ top_k: int
102
+ device_buffer_size: int
103
+ eager_host_mirror: bool = True
104
+
105
+ @classmethod
106
+ def from_vllm_config(
107
+ cls,
108
+ vllm_config: VllmConfig,
109
+ model_top_k: int,
110
+ block_size: int | None = None,
111
+ ) -> ResolvedHiSparseConfig | None:
112
+ config = vllm_config.attention_config.hisparse_config
113
+ if config is None:
114
+ return None
115
+
116
+ max_query_len = _get_max_decode_query_len(vllm_config)
117
+ min_device_buffer_size = max_query_len * model_top_k
118
+ # Retain the worst-case union of every speculative step's top-k and
119
+ # leave one additional top-k of LRU slack.
120
+ configured_size = config.device_buffer_size
121
+ if configured_size is None:
122
+ device_buffer_size = (max_query_len + 1) * model_top_k
123
+ else:
124
+ device_buffer_size = configured_size
125
+
126
+ if device_buffer_size < min_device_buffer_size:
127
+ raise ValueError(
128
+ "HiSparse device_buffer_size must cover every decode query's "
129
+ "index_topk rows. Got "
130
+ f"device_buffer_size={device_buffer_size}, "
131
+ f"max_decode_query_len={max_query_len}, "
132
+ f"index_topk={model_top_k}; expected at least "
133
+ f"{min_device_buffer_size}."
134
+ )
135
+ max_device_buffer_size = torch.iinfo(torch.int16).max + 1
136
+ if device_buffer_size > max_device_buffer_size:
137
+ raise ValueError(
138
+ "HiSparse device_buffer_size exceeds the int16 slot-index "
139
+ f"limit: got {device_buffer_size}, maximum is "
140
+ f"{max_device_buffer_size}."
141
+ )
142
+ if configured_size is not None and block_size is not None:
143
+ padding = -device_buffer_size % block_size
144
+ if padding:
145
+ logger.warning(
146
+ "HiSparse device_buffer_size=%d is not aligned to the "
147
+ "%d-token kernel block size and allocates %d unused "
148
+ "padding rows per hot group. Use %d to use all allocated "
149
+ "rows.",
150
+ device_buffer_size,
151
+ block_size,
152
+ padding,
153
+ device_buffer_size + padding,
154
+ )
155
+ return cls(
156
+ top_k=model_top_k,
157
+ device_buffer_size=device_buffer_size,
158
+ eager_host_mirror=config.eager_host_mirror,
159
+ )
160
+
161
+
162
+def check_hisparse_host_memory(pool_bytes: int) -> None:
163
+ """Reject a physical host-pool allocation that cannot fit in RAM."""
164
+ mem = psutil.virtual_memory()
165
+ if pool_bytes > mem.available * 0.95:
166
+ raise ValueError(
167
+ f"HiSparse pinned host pool needs ~{pool_bytes / 2**30:.0f} GiB "
168
+ f"but only {mem.available / 2**30:.0f} GiB of RAM is available. "
169
+ "Lower the HiSparseConnector host_pool_gib or leave headroom "
170
+ "for co-tenants."
171
+ )
172
+
173
+
174
+def allocate_pinned_host_pool(size: int) -> tuple[torch.Tensor, torch.Tensor]:
175
+ """Allocate and deterministically register an exact-size host KV region."""
176
+ page = 4096
177
+ padded_size = round_up(size, page)
178
+ backing = torch.empty(padded_size + page, dtype=torch.int8, device="cpu")
179
+ aligned_offset = (-backing.data_ptr()) % page
180
+ registered = backing[aligned_offset : aligned_offset + padded_size]
181
+ pin_tensor(registered)
182
+ return registered[:size], registered
183
+
184
+
185
+class HiSparseHostPool:
186
+ """Keep the single host backing and its exact CUDA registration range."""
187
+
188
+ def __init__(self) -> None:
189
+ self.backing: torch.Tensor | None = None
190
+ self.registered: torch.Tensor | None = None
191
+
192
+ def allocate(self, size: int) -> torch.Tensor:
193
+ assert self.backing is None, "HiSparse host pool is already allocated."
194
+ check_hisparse_host_memory(size)
195
+ self.backing, self.registered = allocate_pinned_host_pool(size)
196
+ return self.backing
197
+
198
+
199
+def release_pinned_state(
200
+ runtimes: list[HiSparseRuntime], pinned_host_pools: list[torch.Tensor]
201
+) -> None:
202
+ """Synchronize and release registered host KV pools."""
203
+ if pinned_host_pools:
204
+ try:
205
+ torch.accelerator.synchronize()
206
+ except RuntimeError as e:
207
+ logger.warning(
208
+ "HiSparse: CUDA context unusable at teardown (%s); leaving "
209
+ "%d host-pool tensors pinned for kernel exit reclaim.",
210
+ e,
211
+ len(pinned_host_pools),
212
+ )
213
+ return
214
+
215
+ cudart = torch.cuda.cudart()
216
+ release_start = time.perf_counter()
217
+ freed_bytes = 0
218
+ while pinned_host_pools:
219
+ tensor = pinned_host_pools[-1]
220
+ err = cudart.cudaHostUnregister(tensor.data_ptr())
221
+ if err.value != 0:
222
+ logger.warning(
223
+ "HiSparse: cudaHostUnregister failed (code=%d); leaving "
224
+ "%d host-pool tensors pinned for kernel exit reclaim.",
225
+ err.value,
226
+ len(pinned_host_pools),
227
+ )
228
+ cudart.cudaGetLastError()
229
+ break
230
+ freed_bytes += tensor.nbytes
231
+ pinned_host_pools.pop()
232
+ if freed_bytes:
233
+ logger.info(
234
+ "HiSparse: unpinned %.1f GiB of host pool in %.1fs.",
235
+ freed_bytes / 2**30,
236
+ time.perf_counter() - release_start,
237
+ )
238
+
239
+ for runtime in runtimes:
240
+ del runtime._host_cache
241
+ del runtime.registered_host_pool
242
+ del runtime.hot_backing
243
+
244
+
245
+def hisparse_prefill_staging_remap(
246
+ block_table: torch.Tensor, block_size: int
247
+) -> tuple[torch.Tensor, torch.Tensor]:
248
+ """Renumber blocks without a data-dependent CUDA output allocation."""
249
+ block_ids = torch.cat(
250
+ (
251
+ torch.zeros(1, dtype=torch.int32, device=block_table.device),
252
+ block_table.flatten().clamp(min=0).to(torch.int32),
253
+ )
254
+ )
255
+ sorted_ids, permutation = torch.sort(block_ids)
256
+ is_new = torch.cat(
257
+ (
258
+ torch.ones(1, dtype=torch.bool, device=block_table.device),
259
+ sorted_ids[1:] != sorted_ids[:-1],
260
+ )
261
+ )
262
+ compact_sorted = torch.cumsum(is_new, dim=0, dtype=torch.int32) - 1
263
+ inverse = torch.empty_like(compact_sorted)
264
+ inverse.scatter_(0, permutation, compact_sorted)
265
+
266
+ unique_ids = torch.full_like(sorted_ids, -1)
267
+ unique_ids.scatter_(0, compact_sorted, sorted_ids)
268
+ valid = (
269
+ torch.arange(sorted_ids.numel(), device=block_table.device)
270
+ <= compact_sorted[-1]
271
+ )
272
+ unique_ids = torch.where(valid, unique_ids, -1)
273
+ new_bt = inverse[1:].view_as(block_table)
274
+ offsets = torch.arange(block_size, dtype=torch.int32, device=block_table.device)
275
+ row_ids = (
276
+ torch.where(
277
+ valid[:, None],
278
+ unique_ids[:, None] * block_size + offsets,
279
+ -1,
280
+ )
281
+ ).view(1, -1)
282
+ return new_bt, row_ids
283
+
284
+
285
+@dataclass
286
+class HiSparsePrefillStagingPlan:
287
+ block_table: torch.Tensor
288
+ row_ids: torch.Tensor
289
+ dst_rows: torch.Tensor
290
+ miss_mask: torch.Tensor
291
+ block_size: int
292
+ # Host rows with a GPU-resident copy (adopted shadow pages): the flat
293
+ # resident-cache row to read instead of DMAing from host, -1 for misses.
294
+ gpu_row_ids: torch.Tensor | None = None
295
+ gpu_source_key: tuple[int, int] | None = None
296
+
297
+ def ensure_gpu_sources(
298
+ self,
299
+ resident_block_table: torch.Tensor,
300
+ resident_block_size: int,
301
+ ) -> None:
302
+ """Resolve which staged rows can be served from the resident cache.
303
+
304
+ Computed once per plan (the resident block table is shared by every
305
+ layer in the group); non-null resident pages become miss_mask=0 rows
306
+ gathered device-to-device by ``gather_prefill_cache``.
307
+ """
308
+ source_key = (resident_block_table.data_ptr(), resident_block_size)
309
+ if self.gpu_source_key == source_key:
310
+ return
311
+ block_size = self.block_size
312
+ if resident_block_size <= 0 or block_size % resident_block_size != 0:
313
+ return
314
+ device = self.row_ids.device
315
+ num_unique = self.row_ids.shape[1] // block_size
316
+ host_ids = self.row_ids[0].view(num_unique, block_size)[:, 0] // block_size
317
+ new_bt = self.block_table.to(torch.int64)
318
+ num_rows, num_cols = new_bt.shape
319
+ if num_rows == 0 or resident_block_table.shape[0] < num_rows:
320
+ return
321
+ # One representative (row, col) per unique host block: any request
322
+ # referencing the block holds an equivalent (refcounted) resident view.
323
+ flat_pos = torch.arange(num_rows * num_cols, device=device)
324
+ rep = torch.full(
325
+ (num_unique,), num_rows * num_cols, device=device, dtype=torch.int64
326
+ )
327
+ rep.scatter_reduce_(
328
+ 0, new_bt.reshape(-1), flat_pos, reduce="amin", include_self=True
329
+ )
330
+ rep_row = (rep // num_cols).clamp(max=resident_block_table.shape[0] - 1)
331
+ rep_col = rep % num_cols
332
+ pages_per_block = block_size // resident_block_size
333
+ res_pos = rep_col[:, None] * pages_per_block + torch.arange(
334
+ pages_per_block, device=device
335
+ )
336
+ res_cols = resident_block_table.shape[1]
337
+ res_blocks = torch.where(
338
+ res_pos < res_cols,
339
+ torch.gather(
340
+ resident_block_table.to(torch.int64)[rep_row],
341
+ 1,
342
+ res_pos.clamp(max=max(res_cols - 1, 0)),
343
+ ),
344
+ torch.zeros_like(res_pos),
345
+ )
346
+ offsets = torch.arange(block_size, device=device)
347
+ per_off_block = res_blocks[:, offsets // resident_block_size]
348
+ gpu_rows = (
349
+ per_off_block * resident_block_size
350
+ + (offsets % resident_block_size)[None, :]
351
+ )
352
+ # Block id 0 is the null block in both id spaces.
353
+ valid_rows = self.row_ids.view(num_unique, block_size) >= 0
354
+ hit = (per_off_block > 0) & (host_ids[:, None] > 0) & valid_rows
355
+ self.gpu_row_ids = torch.where(hit, gpu_rows, -1).reshape(1, -1).to(torch.int32)
356
+ self.miss_mask = ((self.gpu_row_ids < 0) & valid_rows.view(1, -1)).int()
357
+ self.gpu_source_key = source_key
358
+
359
+
360
+def build_hisparse_prefill_staging_plan(
361
+ block_table: torch.Tensor,
362
+ seq_lens: torch.Tensor,
363
+ block_size: int,
364
+ staging_block_capacity: int,
365
+) -> HiSparsePrefillStagingPlan:
366
+ """Build an asynchronous layer-independent host-cache staging remap."""
367
+ device = block_table.device
368
+ used = (seq_lens.to(torch.int64) + block_size - 1) // block_size
369
+ packed_offsets = torch.arange(staging_block_capacity, device=device)
370
+ row_ends = torch.cumsum(used, dim=0)
371
+ rows = torch.searchsorted(row_ends, packed_offsets, right=True)
372
+ valid = rows < block_table.shape[0]
373
+ rows = rows.clamp(max=block_table.shape[0] - 1)
374
+ row_starts = torch.cat((torch.zeros_like(row_ends[:1]), row_ends[:-1]))
375
+ columns = packed_offsets - row_starts[rows]
376
+ valid &= columns < block_table.shape[1]
377
+ flat_positions = rows * block_table.shape[1] + columns.clamp(
378
+ max=block_table.shape[1] - 1
379
+ )
380
+ referenced = torch.where(valid, block_table.flatten()[flat_positions], 0)
381
+ referenced_bt, row_ids = hisparse_prefill_staging_remap(referenced, block_size)
382
+ scratch_index = block_table.numel()
383
+ scatter_positions = torch.where(valid, flat_positions, scratch_index)
384
+ remapped = torch.where(valid, referenced_bt.flatten(), 0)
385
+ new_bt_storage = torch.zeros(scratch_index + 1, dtype=torch.int32, device=device)
386
+ new_bt_storage.scatter_(0, scatter_positions, remapped)
387
+ new_bt = new_bt_storage[:-1].view_as(block_table)
388
+ valid_rows = row_ids >= 0
389
+ dst_rows = torch.arange(row_ids.shape[1], dtype=torch.int32, device=device).view(
390
+ 1, -1
391
+ )
392
+ return HiSparsePrefillStagingPlan(
393
+ block_table=new_bt,
394
+ row_ids=row_ids,
395
+ dst_rows=dst_rows,
396
+ miss_mask=valid_rows.to(torch.int32),
397
+ block_size=block_size,
398
+ )
399
+
400
+
401
+def _has_hisparse_ops() -> bool:
402
+ if not hasattr(torch.ops, "_C_cache_ops"):
403
+ return False
404
+ return (
405
+ hasattr(torch.ops._C_cache_ops, "hisparse_resolve_residency")
406
+ and hasattr(torch.ops._C_cache_ops, "hisparse_invalidate_written_slots")
407
+ and hasattr(torch.ops._C_cache_ops, "hisparse_gather_plan")
408
+ and hasattr(torch.ops._C_cache_ops, "hisparse_gather_compact")
409
+ )
410
+
411
+
412
+class _SharedTopKState:
413
+ """CUDA-graph-safe residency result shared by index-sharing layers."""
414
+
415
+ def __init__(self, device: torch.device, max_rows: int, top_k: int) -> None:
416
+ self.device_topk_rows = torch.full(
417
+ (max_rows, top_k), -1, dtype=torch.int32, device=device
418
+ )
419
+ self.physical_topk_indices = torch.empty_like(self.device_topk_rows)
420
+ self.swap_host_physical_rows = torch.empty(
421
+ (max_rows, top_k), dtype=torch.int32, device=device
422
+ )
423
+ self.swap_device_physical_rows = torch.empty_like(self.swap_host_physical_rows)
424
+ self.swap_counts = torch.empty(max_rows, dtype=torch.int32, device=device)
425
+ self.valid_topk_counts = torch.empty(max_rows, dtype=torch.int32, device=device)
426
+
427
+
428
+def _create_shared_topk_state(
429
+ device: torch.device, max_rows: int, top_k: int
430
+) -> _SharedTopKState:
431
+ return _SharedTopKState(device, max_rows, top_k)
432
+
433
+
434
+def _create_copy_stream(device: torch.device) -> torch.Stream:
435
+ return torch.Stream(device=device)
436
+
437
+
438
+class HiSparseIndexGroup:
439
+ """GPU index-resolution state shared by one leader and its followers."""
440
+
441
+ leader: HiSparseRuntime
442
+
443
+ def __init__(
444
+ self,
445
+ device: torch.device,
446
+ max_num_reqs: int,
447
+ region_stride: int,
448
+ max_swap_rows: int,
449
+ top_k: int,
450
+ copy_stream: torch.Stream | None = None,
451
+ logical_topk_ready: torch.Event | None = None,
452
+ ) -> None:
453
+ self.device_global_indices = torch.full(
454
+ (max_num_reqs, region_stride),
455
+ -1,
456
+ dtype=torch.int32,
457
+ device=device,
458
+ )
459
+ self.lru_init = torch.arange(region_stride, dtype=torch.int16, device=device)
460
+ self.lru_slots = self.lru_init.repeat(max_num_reqs, 1).contiguous()
461
+ self.shared_topk = _create_shared_topk_state(device, max_swap_rows, top_k)
462
+ self.copy_stream = (
463
+ copy_stream if copy_stream is not None else _create_copy_stream(device)
464
+ )
465
+ self.logical_topk_ready = logical_topk_ready
466
+ self.followers: list[HiSparseRuntime] = []
467
+
468
+
469
+class HiSparseRuntime:
470
+ """Per-cache host/hot data plane and GPU replacement state."""
471
+
472
+ hot: PagedCacheView
473
+ hot_backing: torch.Tensor
474
+ hot_block_table: torch.Tensor
475
+ _host_cache: torch.Tensor
476
+ registered_host_pool: torch.Tensor
477
+
478
+ def __init__(
479
+ self,
480
+ config: ResolvedHiSparseConfig,
481
+ max_num_reqs: int,
482
+ row_width: int,
483
+ kv_dtype: torch.dtype,
484
+ device: torch.device | str,
485
+ max_swap_rows: int | None = None,
486
+ index_group: HiSparseIndexGroup | None = None,
487
+ copy_stream: torch.Stream | None = None,
488
+ logical_topk_ready: torch.Event | None = None,
489
+ ) -> None:
490
+ if not _has_hisparse_ops():
491
+ raise RuntimeError(
492
+ "HiSparse requires its compiled _C_cache_ops CUDA kernels "
493
+ "(host-resident decode has no Python fallback). Rebuild vLLM "
494
+ "from source so "
495
+ "csrc/libtorch_stable/hisparse_kernels.cu is included."
496
+ )
497
+ self.max_num_reqs = max_num_reqs
498
+ self.row_width = row_width
499
+ self.kv_dtype = kv_dtype
500
+ self.device = torch.device(device)
501
+ # Logical slots per request. Physical rows come from its ephemeral HMA
502
+ # block table and need not be contiguous in the shared slab.
503
+ self.region_stride = config.device_buffer_size
504
+
505
+ row_bytes = row_width * kv_dtype.itemsize
506
+ if row_bytes % 16 != 0:
507
+ raise ValueError(
508
+ f"HiSparse requires 16-byte aligned KV rows, got {row_bytes}B."
509
+ )
510
+
511
+ self._layer_ready_event: torch.Event | None = None
512
+ self._swap_staged = False
513
+ self._swap_step = 0
514
+ self.is_group_leader = index_group is None
515
+ if index_group is None:
516
+ index_group = HiSparseIndexGroup(
517
+ self.device,
518
+ max_num_reqs,
519
+ self.region_stride,
520
+ max_swap_rows or max_num_reqs,
521
+ config.top_k,
522
+ copy_stream,
523
+ logical_topk_ready,
524
+ )
525
+ index_group.leader = self
526
+ else:
527
+ index_group.followers.append(self)
528
+ self.index_group = index_group
529
+
530
+ self.eager_host_mirror = config.eager_host_mirror
531
+ self.resident_source_index = -1
532
+ self.request_state_indices: torch.Tensor | None = None
533
+
534
+ @property
535
+ def host_cache(self) -> torch.Tensor:
536
+ return self._host_cache
537
+
538
+ def bind_hot_cache(
539
+ self,
540
+ raw_tensor: torch.Tensor,
541
+ *,
542
+ byte_offset: int,
543
+ block_stride: int,
544
+ num_blocks: int,
545
+ block_size: int,
546
+ block_table: torch.Tensor,
547
+ ) -> None:
548
+ """Bind this runtime's strided view into the shared GPU HMA slab."""
549
+ self.hot = PagedCacheView.bind(
550
+ raw_tensor,
551
+ dtype=self.kv_dtype,
552
+ row_width=self.row_width,
553
+ byte_offset=byte_offset,
554
+ block_stride=block_stride,
555
+ num_blocks=num_blocks,
556
+ block_size=block_size,
557
+ )
558
+ self.hot_backing = raw_tensor
559
+ self.hot_block_table = block_table
560
+
561
+ def bind_source_cache(
562
+ self,
563
+ kv_cache: torch.Tensor,
564
+ *,
565
+ registered_host_pool: torch.Tensor | None = None,
566
+ ) -> None:
567
+ if kv_cache.dtype != self.kv_dtype or kv_cache.shape[-1] != self.row_width:
568
+ raise ValueError(
569
+ "HiSparse runtime bound to a KV cache with mismatched "
570
+ f"layout: expected ({self.row_width}, {self.kv_dtype}), got "
571
+ f"({kv_cache.shape[-1]}, {kv_cache.dtype})."
572
+ )
573
+ # Host-resident pool: the cache itself is the only full-size store;
574
+ # every allocated slot is written before the indexer can select it.
575
+ if kv_cache.device.type != "cpu":
576
+ raise ValueError(
577
+ "HiSparse requires a host-resident KV pool; got a KV cache "
578
+ f"on {kv_cache.device}."
579
+ )
580
+ # The pool is pinned via cudaHostRegister (exact-size, deterministic
581
+ # unpin at shutdown); torch's is_pinned() only recognizes its own
582
+ # caching-host-allocator memory, so also accept ranges the model
583
+ # allocator explicitly registered.
584
+ if not (kv_cache.is_pinned() or registered_host_pool is not None):
585
+ raise ValueError("HiSparse host-resident KV pool must be pinned memory.")
586
+
587
+ self._host_cache = kv_cache.view(-1, kv_cache.shape[-1])
588
+ self.registered_host_pool = (
589
+ registered_host_pool if registered_host_pool is not None else kv_cache
590
+ )
591
+
592
+ def gather_prefill_cache(
593
+ self,
594
+ kv_cache: torch.Tensor,
595
+ plan: HiSparsePrefillStagingPlan,
596
+ resident_cache: torch.Tensor | None = None,
597
+ ) -> torch.Tensor:
598
+ """Gather this runtime's host cache using a shared staging plan."""
599
+ if kv_cache.shape[1] != plan.block_size:
600
+ raise ValueError(
601
+ f"HiSparse staging block size {plan.block_size} does not match "
602
+ f"the KV cache block size {kv_cache.shape[1]}."
603
+ )
604
+ row_width = kv_cache.shape[-1]
605
+
606
+ staged = torch.empty(
607
+ (
608
+ plan.row_ids.shape[1] // plan.block_size,
609
+ plan.block_size,
610
+ row_width,
611
+ ),
612
+ dtype=kv_cache.dtype,
613
+ device=plan.block_table.device,
614
+ )
615
+ if plan.gpu_row_ids is not None and resident_cache is not None:
616
+ gpu_rows = plan.gpu_row_ids[0].to(torch.long)
617
+ src = gpu_rows.clamp_min(0)
618
+ resident_block_size = resident_cache.shape[1]
619
+ rows = resident_cache[src // resident_block_size, src % resident_block_size]
620
+ if rows.dtype != staged.dtype:
621
+ rows = rows.contiguous().view(staged.dtype)
622
+ staged.view(-1, row_width).copy_(rows)
623
+ torch.ops._C_cache_ops.hisparse_gather_plan(
624
+ kv_cache.view(-1, row_width),
625
+ staged,
626
+ plan.row_ids,
627
+ plan.dst_rows,
628
+ plan.miss_mask,
629
+ None,
630
+ None,
631
+ 0,
632
+ )
633
+ return staged
634
+
635
+ def reset_hot_state(self) -> None:
636
+ """Drop all hot-buffer bookkeeping (hits become misses)."""
637
+ group = self.index_group
638
+ compute_stream = current_stream()
639
+ compute_stream.wait_stream(group.copy_stream)
640
+ group.device_global_indices.fill_(-1)
641
+ group.lru_slots.copy_(group.lru_init.expand_as(group.lru_slots))
642
+ group.copy_stream.wait_stream(compute_stream)
643
+
644
+ def invalidate_sorted_slots(
645
+ self,
646
+ sorted_slots: torch.Tensor,
647
+ state_indices: torch.Tensor,
648
+ ) -> None:
649
+ """Drop hot copies using shared, preprocessed invalidation inputs."""
650
+ device_global_indices = self.index_group.device_global_indices
651
+ active_indices = device_global_indices.index_select(0, state_indices)
652
+ positions = torch.searchsorted(sorted_slots, active_indices)
653
+ positions.clamp_(max=sorted_slots.numel() - 1)
654
+ active_indices.masked_fill_(
655
+ sorted_slots[positions] == active_indices,
656
+ -1,
657
+ )
658
+ device_global_indices.index_copy_(0, state_indices, active_indices)
659
+
660
+ def invalidate_written_slots(
661
+ self,
662
+ written_slots: torch.Tensor,
663
+ req_id_per_token: torch.Tensor,
664
+ ) -> None:
665
+ num_tokens = min(written_slots.numel(), req_id_per_token.numel())
666
+ if num_tokens == 0:
667
+ return
668
+ assert self.request_state_indices is not None
669
+ torch.ops._C_cache_ops.hisparse_invalidate_written_slots(
670
+ self.index_group.device_global_indices,
671
+ self.request_state_indices,
672
+ req_id_per_token[:num_tokens],
673
+ written_slots[:num_tokens],
674
+ )
675
+
676
+ def begin_forward(self) -> None:
677
+ self._swap_step = 0
678
+
679
+ def _step_rows(self, num_tokens: int) -> slice:
680
+ start = self._swap_step * num_tokens
681
+ stop = start + num_tokens
682
+ if stop > self.index_group.shared_topk.physical_topk_indices.shape[0]:
683
+ raise ValueError(
684
+ "HiSparse swap rows exceed the configured speculative decode "
685
+ f"capacity: stop={stop}, capacity="
686
+ f"{self.index_group.shared_topk.physical_topk_indices.shape[0]}."
687
+ )
688
+ self._swap_step += 1
689
+ return slice(start, stop)
690
+
691
+ def _resolve_residency(
692
+ self,
693
+ *,
694
+ resident: HiSparseCacheHandle | None = None,
695
+ req_id_per_token: torch.Tensor,
696
+ block_table: torch.Tensor,
697
+ topk_indices: torch.Tensor,
698
+ block_size: int,
699
+ return_valid_counts: bool = False,
700
+ shared_rows: slice,
701
+ attention_indices_out: torch.Tensor | None = None,
702
+ valid_counts_out: torch.Tensor | None = None,
703
+ ) -> None:
704
+ """Resolve logical top-k positions and compact rows requiring swaps."""
705
+ num_tokens = topk_indices.shape[0]
706
+ hot = self.hot
707
+ host_cache = self.host_cache
708
+ assert hot.block_size == block_size
709
+ request_state_indices = self.request_state_indices
710
+ assert request_state_indices is not None
711
+ group = self.index_group
712
+ shared = group.shared_topk
713
+
714
+ device_topk_rows = shared.device_topk_rows[shared_rows]
715
+ valid_topk_counts = None
716
+ if return_valid_counts:
717
+ valid_topk_counts = (
718
+ shared.valid_topk_counts[shared_rows]
719
+ if valid_counts_out is None
720
+ else valid_counts_out
721
+ )
722
+ physical_topk_indices = (
723
+ shared.physical_topk_indices[shared_rows]
724
+ if attention_indices_out is None
725
+ else attention_indices_out
726
+ )
727
+ swap_host_physical_rows = shared.swap_host_physical_rows[shared_rows]
728
+ swap_device_physical_rows = shared.swap_device_physical_rows[shared_rows]
729
+ swap_counts = shared.swap_counts[shared_rows]
730
+
731
+ # Padded rows are skipped by the kernel (request_state_indices) and must
732
+ # come out as -1 so the attention kernel masks them.
733
+ torch.ops._C_cache_ops.hisparse_resolve_residency(
734
+ host_cache,
735
+ hot.cache,
736
+ self.hot_block_table,
737
+ topk_indices,
738
+ device_topk_rows,
739
+ group.device_global_indices,
740
+ group.lru_slots,
741
+ request_state_indices,
742
+ self.region_stride,
743
+ None,
744
+ physical_topk_indices,
745
+ hot.attention_block_stride,
746
+ req_id_per_token[:num_tokens].contiguous(),
747
+ block_table,
748
+ block_size,
749
+ None,
750
+ valid_topk_counts,
751
+ swap_host_physical_rows,
752
+ swap_device_physical_rows,
753
+ swap_counts,
754
+ resident.block_table if resident is not None else None,
755
+ resident.view.block_size
756
+ if resident is not None and resident.view is not None
757
+ else 0,
758
+ 0,
759
+ )
760
+
761
+ def _swap_rows(self, shared_rows: slice) -> None:
762
+ hot = self.hot
763
+ shared = self.index_group.shared_topk
764
+ torch.ops._C_cache_ops.hisparse_gather_compact(
765
+ self.host_cache,
766
+ hot.cache,
767
+ shared.swap_host_physical_rows[shared_rows],
768
+ shared.swap_device_physical_rows[shared_rows],
769
+ shared.swap_counts[shared_rows],
770
+ )
771
+
772
+ def _resolve_and_stage_group(
773
+ self,
774
+ *,
775
+ resident: HiSparseCacheHandle,
776
+ req_id_per_token: torch.Tensor,
777
+ block_table: torch.Tensor,
778
+ logical_topk_indices: torch.Tensor,
779
+ block_size: int,
780
+ return_valid_counts: bool,
781
+ shared_rows: slice,
782
+ attention_indices_out: torch.Tensor | None,
783
+ valid_counts_out: torch.Tensor | None,
784
+ ) -> None:
785
+ group = self.index_group
786
+ compute_stream = current_stream()
787
+ if group.logical_topk_ready is not None:
788
+ group.copy_stream.wait_event(group.logical_topk_ready)
789
+ else:
790
+ group.copy_stream.wait_stream(compute_stream)
791
+ with group.copy_stream:
792
+ self._resolve_residency(
793
+ resident=resident,
794
+ req_id_per_token=req_id_per_token,
795
+ block_table=block_table,
796
+ topk_indices=logical_topk_indices,
797
+ block_size=block_size,
798
+ return_valid_counts=return_valid_counts,
799
+ shared_rows=shared_rows,
800
+ attention_indices_out=attention_indices_out,
801
+ valid_counts_out=valid_counts_out,
802
+ )
803
+ runtimes = [self]
804
+ if resident.decode_batch:
805
+ runtimes.extend(group.followers)
806
+ for runtime in runtimes:
807
+ runtime._swap_rows(shared_rows)
808
+ if runtime._layer_ready_event is None:
809
+ runtime._layer_ready_event = torch.Event()
810
+ runtime._layer_ready_event.record(group.copy_stream)
811
+ runtime._swap_staged = True
812
+
813
+ def _stage_this_layer(self, shared_rows: slice) -> None:
814
+ group = self.index_group
815
+ compute_stream = current_stream()
816
+ group.copy_stream.wait_stream(compute_stream)
817
+ with group.copy_stream:
818
+ self._swap_rows(shared_rows)
819
+ if self._layer_ready_event is None:
820
+ self._layer_ready_event = torch.Event()
821
+ self._layer_ready_event.record(group.copy_stream)
822
+ self._swap_staged = True
823
+
824
+ def wait_for_staged_swap(self) -> None:
825
+ if not self._swap_staged or self._layer_ready_event is None:
826
+ raise RuntimeError("HiSparse layer swap was not prefetched.")
827
+ current_stream().wait_event(self._layer_ready_event)
828
+ self._swap_staged = False
829
+
830
+ def swap_in(
831
+ self,
832
+ *,
833
+ resident: HiSparseCacheHandle,
834
+ req_id_per_token: torch.Tensor,
835
+ block_table: torch.Tensor,
836
+ logical_topk_indices: torch.Tensor,
837
+ block_size: int,
838
+ return_valid_counts: bool = False,
839
+ attention_indices_out: torch.Tensor | None = None,
840
+ valid_counts_out: torch.Tensor | None = None,
841
+ ) -> HiSparseTopKResult:
842
+ """Resolve once per group and ensure this layer's rows are available."""
843
+ num_tokens = logical_topk_indices.shape[0]
844
+ shared_rows = self._step_rows(num_tokens)
845
+ group = self.index_group
846
+ hot = self.hot
847
+ assert hot.block_size == block_size
848
+ assert hot.attention_block_stride == group.leader.hot.attention_block_stride
849
+
850
+ if self.is_group_leader:
851
+ self._resolve_and_stage_group(
852
+ resident=resident,
853
+ req_id_per_token=req_id_per_token,
854
+ block_table=block_table,
855
+ logical_topk_indices=logical_topk_indices,
856
+ block_size=block_size,
857
+ return_valid_counts=return_valid_counts,
858
+ shared_rows=shared_rows,
859
+ attention_indices_out=attention_indices_out,
860
+ valid_counts_out=valid_counts_out,
861
+ )
862
+
863
+ if not self._swap_staged:
864
+ self._stage_this_layer(shared_rows)
865
+ self.wait_for_staged_swap()
866
+
867
+ physical_topk_indices = group.shared_topk.physical_topk_indices[shared_rows]
868
+ if return_valid_counts:
869
+ return (
870
+ physical_topk_indices,
871
+ group.shared_topk.valid_topk_counts[shared_rows],
872
+ )
873
+ return physical_topk_indices
874
+
875
+
876
+class HiSparseCacheHandle:
877
+ """Attention-facing handle for resident KV and sparse offload state."""
878
+
879
+ def __init__(self, runtime: HiSparseRuntime) -> None:
880
+ self.view: PagedCacheView | None = None
881
+ self.block_table: torch.Tensor | None = None
882
+ self.source_block_table: torch.Tensor | None = None
883
+ self.slot_mapping: torch.Tensor | None = None
884
+ self.runtime = runtime
885
+ self.dummy_batch = False
886
+ self.decode_batch = False
887
+ self.all_context_pages_resident = True
888
+ self.num_actual_tokens = 0
889
+ self.num_decode_tokens = 0
890
+ self.req_id_per_token: torch.Tensor | None = None
891
+ self.host_mirror_required = False
892
+ self.mirror_from_resident = False
893
+ self.mirror_slot_mapping: torch.Tensor | None = None
894
+ self.mirror_staging_cache: torch.Tensor | None = None
895
+ self.mirror_staging_slots: torch.Tensor | None = None
896
+ self.submit_layer_mirror: Callable[[], None] | None = None
897
+ self.index_group_caches: list[HiSparseCacheHandle] = [self]
898
+
899
+ def prepare_group_for_batch(self, attn_metadata: Any | None) -> None:
900
+ assert self.runtime.is_group_leader
901
+ for cache in self.index_group_caches:
902
+ cache._prepare_for_batch(attn_metadata)
903
+
904
+ def _prepare_for_batch(self, attn_metadata: Any | None) -> None:
905
+ self.dummy_batch = attn_metadata is None
906
+ self.runtime.begin_forward()
907
+ self.num_actual_tokens = (
908
+ attn_metadata.num_actual_tokens if attn_metadata is not None else 0
909
+ )
910
+ self.num_decode_tokens = (
911
+ attn_metadata.num_decode_tokens if attn_metadata is not None else 0
912
+ )
913
+ self.req_id_per_token = (
914
+ attn_metadata.req_id_per_token if attn_metadata is not None else None
915
+ )
916
+ self.decode_batch = is_hisparse_decode_batch(attn_metadata)
917
+ self.host_mirror_required = attn_metadata is not None and (
918
+ not self.decode_batch or self.runtime.eager_host_mirror
919
+ )
920
+
921
+ def write_target(
922
+ self, num_input_rows: int, num_slot_rows: int
923
+ ) -> tuple[torch.Tensor, torch.Tensor, int]:
924
+ assert self.view is not None and self.slot_mapping is not None
925
+ num_rows = min(num_input_rows, num_slot_rows, self.num_actual_tokens)
926
+ if self.decode_batch:
927
+ num_rows = min(num_rows, self.runtime.max_num_reqs)
928
+ return self.view.cache, self.slot_mapping[:num_rows], num_rows
929
+
930
+ def mirror_write_target(
931
+ self, num_rows: int
932
+ ) -> tuple[torch.Tensor, torch.Tensor] | None:
933
+ if (
934
+ not self.host_mirror_required
935
+ or self.decode_batch
936
+ or self.mirror_from_resident
937
+ ):
938
+ return None
939
+ cache = self.mirror_staging_cache
940
+ slots = self.mirror_staging_slots
941
+ if cache is None or slots is None:
942
+ raise RuntimeError("HiSparse prefill mirror staging is not bound.")
943
+ if num_rows > slots.numel():
944
+ raise RuntimeError(
945
+ "HiSparse prefill mirror exceeds staging capacity: "
946
+ f"{num_rows} > {slots.numel()}."
947
+ )
948
+ return cache, slots[:num_rows]
949
+
950
+ def finish_kv_update(self) -> None:
951
+ if self.dummy_batch:
952
+ return
953
+ if (
954
+ self.submit_layer_mirror is not None
955
+ and not self.decode_batch
956
+ and get_forward_context().cudagraph_runtime_mode != CUDAGraphMode.FULL
957
+ ):
958
+ self.submit_layer_mirror()
959
+
960
+ def bind_cache(
961
+ self,
962
+ raw_tensor: torch.Tensor,
963
+ *,
964
+ byte_offset: int,
965
+ block_stride: int,
966
+ num_blocks: int,
967
+ block_size: int,
968
+ block_table: torch.Tensor,
969
+ slot_mapping: torch.Tensor,
970
+ ) -> None:
971
+ self.view = PagedCacheView.bind(
972
+ raw_tensor,
973
+ dtype=self.runtime.kv_dtype,
974
+ row_width=self.runtime.row_width,
975
+ byte_offset=byte_offset,
976
+ block_stride=block_stride,
977
+ num_blocks=num_blocks,
978
+ block_size=block_size,
979
+ )
980
+ self.block_table = block_table
981
+ self.slot_mapping = slot_mapping
982
+
983
+ def swap_in(
984
+ self,
985
+ req_id_per_token: torch.Tensor,
986
+ block_table: torch.Tensor,
987
+ logical_topk_indices: torch.Tensor,
988
+ *,
989
+ block_size: int,
990
+ return_valid_counts: bool = False,
991
+ attention_indices_out: torch.Tensor | None = None,
992
+ valid_counts_out: torch.Tensor | None = None,
993
+ ) -> HiSparseTopKResult:
994
+ return self.runtime.swap_in(
995
+ resident=self,
996
+ req_id_per_token=req_id_per_token[: logical_topk_indices.shape[0]],
997
+ block_table=block_table,
998
+ logical_topk_indices=logical_topk_indices,
999
+ block_size=block_size,
1000
+ return_valid_counts=return_valid_counts,
1001
+ attention_indices_out=attention_indices_out,
1002
+ valid_counts_out=valid_counts_out,
1003
+ )
1004
+
1005
+
1006
+def initialize_hisparse_runtime_buffers(
1007
+ cache_handles: list[HiSparseCacheHandle],
1008
+ *,
1009
+ max_num_reqs: int,
1010
+ max_num_batched_tokens: int,
1011
+) -> None:
1012
+ """Allocate shared request state and per-layer prefill staging buffers."""
1013
+ assert cache_handles
1014
+ resident = cache_handles[0].view
1015
+ assert resident is not None
1016
+ device = resident.cache.device
1017
+ request_state_indices = torch.full(
1018
+ (max_num_reqs,), -1, dtype=torch.int32, device=device
1019
+ )
1020
+ staging_blocks = (
1021
+ max_num_batched_tokens + resident.block_size - 1
1022
+ ) // resident.block_size
1023
+ mirror_staging_caches = torch.empty(
1024
+ (
1025
+ len(cache_handles),
1026
+ staging_blocks,
1027
+ resident.block_size,
1028
+ resident.cache.shape[-1],
1029
+ ),
1030
+ dtype=resident.cache.dtype,
1031
+ device=device,
1032
+ )
1033
+ mirror_staging_slots = torch.arange(
1034
+ max_num_batched_tokens, dtype=torch.int64, device=device
1035
+ )
1036
+ for layer_index, cache_handle in enumerate(cache_handles):
1037
+ cache_handle.runtime.request_state_indices = request_state_indices
1038
+ cache_handle.mirror_staging_cache = mirror_staging_caches[layer_index]
1039
+ cache_handle.mirror_staging_slots = mirror_staging_slots
1040
+
1041
+
1042
+def create_hisparse_cache_handle(
1043
+ vllm_config: VllmConfig,
1044
+ model_top_k: int,
1045
+ *,
1046
+ is_index_group_leader: bool,
1047
+ row_width: int,
1048
+ kv_dtype: torch.dtype,
1049
+ index_group: HiSparseMLAIndexGroup | None = None,
1050
+ device: torch.device | str | None = None,
1051
+) -> HiSparseCacheHandle | None:
1052
+ config = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k)
1053
+ if config is None:
1054
+ return None
1055
+
1056
+ max_num_reqs = vllm_config.scheduler_config.max_num_seqs
1057
+ max_swap_rows = _get_max_swap_rows(vllm_config)
1058
+ if device is None:
1059
+ device = torch.device(
1060
+ current_platform.device_type, torch.accelerator.current_device_index()
1061
+ )
1062
+
1063
+ hisparse_group = (
1064
+ None
1065
+ if is_index_group_leader or index_group is None
1066
+ else index_group.hisparse_group
1067
+ )
1068
+ runtime = HiSparseRuntime(
1069
+ config=config,
1070
+ max_num_reqs=max_num_reqs,
1071
+ max_swap_rows=max_swap_rows,
1072
+ row_width=row_width,
1073
+ kv_dtype=kv_dtype,
1074
+ device=device,
1075
+ index_group=hisparse_group,
1076
+ copy_stream=index_group.side_stream if index_group is not None else None,
1077
+ logical_topk_ready=(
1078
+ index_group.logical_topk_ready
1079
+ if index_group is not None and index_group.has_indexer
1080
+ else None
1081
+ ),
1082
+ )
1083
+ if is_index_group_leader and index_group is not None:
1084
+ index_group.hisparse_group = runtime.index_group
1085
+ logger.info_once(
1086
+ "Enabled experimental HiSparse HMA hot cache: top_k=%d, "
1087
+ "device_buffer_size=%d (%d LRU rows), max_num_seqs=%d.",
1088
+ config.top_k,
1089
+ config.device_buffer_size,
1090
+ config.device_buffer_size,
1091
+ max_num_reqs,
1092
+ )
1093
+ handle = HiSparseCacheHandle(runtime)
1094
+ speculative_config = vllm_config.speculative_config
1095
+ handle.mirror_from_resident = bool(
1096
+ vllm_config.scheduler_config.async_scheduling
1097
+ and speculative_config is not None
1098
+ and speculative_config.uses_draft_kv_cache()
1099
+ )
1100
+ return handle
vllm/v1/hisparse/types.py
new
+30
@@ -0,0 +1,30 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from dataclasses import dataclass
5
+
6
+
7
+@dataclass(frozen=True)
8
+class SparseKVPageTransfer:
9
+ """Copy one logical KV page between cache-manager and worker-owned tiers."""
10
+
11
+ transfer_id: int
12
+ destination_block_id: int
13
+ source_block_ids: tuple[int, ...]
14
+ after_forward: bool
15
+
16
+
17
+@dataclass(frozen=True)
18
+class SparseKVRowMirror:
19
+ """Mirror one contiguous resident-row span into the host cache."""
20
+
21
+ source_starts: tuple[int, ...]
22
+ destination_start: int
23
+ num_rows: int
24
+
25
+
26
+@dataclass
27
+class SparseKVOffloadCommand:
28
+ """Opaque scheduler-to-worker command for one model step."""
29
+
30
+ page_transfers: list[SparseKVPageTransfer]
vllm/v1/kv_cache_interface.py
+154
-3
@@ -147,6 +147,11 @@ class KVCacheSpecKind(str, Enum):
147
UNKNOWN = "unknown"
148
149
150
+class SparseCacheRole(str, Enum):
151
+ SPARSE = "sparse"
152
+ INDEXER = "indexer"
153
+
154
+
155
@dataclass(frozen=True)
156
class KVCacheSpec:
157
"""
@@ -192,6 +197,10 @@ class KVCacheSpec:
197
return kernel_block_size // self.tokens_per_state
198
return 1
199
200
+ @property
201
+ def block_table_token_alignment(self) -> int | None:
202
+ return 128
203
+
204
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
205
"""
206
The maximum possible memory usage of this KV cache in bytes.
@@ -245,6 +254,25 @@ class KVCacheSpec:
254
isinstance(spec, uniform_type_base_spec) for spec in kv_cache_specs.values()
255
)
256
257
+ @property
258
+ def has_layer_views(self) -> bool:
259
+ """Whether generic allocation creates per-layer cache views.
260
+
261
+ Specs without a per-layer shape keep the raw backing tensor and lay it
262
+ out themselves. They also have no attention module of their own, so
263
+ they take no part in attention-backend or kernel-block selection.
264
+ """
265
+ return True
266
+
267
+ @property
268
+ def uses_slot_mapping(self) -> bool:
269
+ """Whether the worker computes a per-token slot mapping for this spec.
270
+
271
+ Specs that address their pages themselves (raw storage, ring buffers)
272
+ take no slot mapping row.
273
+ """
274
+ return True
275
+
276
277
def group_kernel_blocks(cache: torch.Tensor, num_blocks: int) -> torch.Tensor:
278
"""View a kernel-block-granular layer cache with manager blocks as dim 0.
@@ -383,6 +411,70 @@ def create_kv_cache_views(
411
return views
412
413
414
+@dataclass(frozen=True, kw_only=True)
415
+class HiSparseHotSpec(KVCacheSpec):
416
+ """Ephemeral per-request HiSparse hot-cache allocation."""
417
+
418
+ page_size: int
419
+ blocks_per_request: int
420
+
421
+ @property
422
+ def prefix_cacheable(self) -> bool:
423
+ return False
424
+
425
+ @property
426
+ def page_size_bytes(self) -> int:
427
+ return self.page_size
428
+
429
+ @property
430
+ def block_table_token_alignment(self) -> int | None:
431
+ return None
432
+
433
+ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
434
+ return self.blocks_per_request * self.page_size
435
+
436
+ def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
437
+ return self.blocks_per_request
438
+
439
+ @property
440
+ def has_layer_views(self) -> bool:
441
+ return False
442
+
443
+ @property
444
+ def uses_slot_mapping(self) -> bool:
445
+ return False
446
+
447
+
448
+@dataclass(frozen=True, kw_only=True)
449
+class HiSparseResidentSpec(KVCacheSpec):
450
+ """Reclaimable GPU-resident pages for host-backed HiSparse KV."""
451
+
452
+ page_size: int
453
+
454
+ @property
455
+ def prefix_cacheable(self) -> bool:
456
+ return False
457
+
458
+ @property
459
+ def page_size_bytes(self) -> int:
460
+ return self.page_size
461
+
462
+ def max_admission_blocks_per_request(
463
+ self, max_in_flight_tokens: int, max_model_len: int
464
+ ) -> int:
465
+ num_tokens = min(max_in_flight_tokens, max_model_len)
466
+ return cdiv(num_tokens, self.block_size)
467
+
468
+ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
469
+ return cdiv(vllm_config.model_config.max_model_len, self.block_size) * (
470
+ self.page_size
471
+ )
472
+
473
+ @property
474
+ def has_layer_views(self) -> bool:
475
+ return False
476
+
477
+
478
@dataclass(frozen=True, kw_only=True)
479
class AttentionSpec(KVCacheSpec):
480
num_kv_heads: int
@@ -557,6 +649,8 @@ class MLAAttentionSpec(FullAttentionSpec):
649
# DeepseekV4 only fields. Non-DeepseekV4 MLA models leave these at defaults.
650
alignment: int | None = None # Default to None for no padding.
651
model_version: str | None = None
652
+ cache_role: SparseCacheRole = SparseCacheRole.SPARSE
653
+ is_index_group_leader: bool = False
654
storage_block_size: int | None = None
655
"""Token width used to view storage when it differs from the kernel block."""
656
# Group capability enabled when any member flattens a non-causal query block
@@ -577,16 +671,20 @@ class MLAAttentionSpec(FullAttentionSpec):
671
cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs)
672
tokens_per_state_set = set(spec.tokens_per_state for spec in specs)
673
model_version_set = set(spec.model_version for spec in specs)
674
+ cache_role_set = {spec.cache_role for spec in specs}
675
+ index_group_leader_set = {spec.is_index_group_leader for spec in specs}
676
storage_block_size_set = set(spec.storage_block_size for spec in specs)
677
assert (
678
len(cache_dtype_str_set) == 1
679
and len(tokens_per_state_set) == 1
680
and len(model_version_set) == 1
681
+ and len(cache_role_set) == 1
682
+ and len(index_group_leader_set) == 1
683
and len(storage_block_size_set) == 1
684
), (
685
"All attention layers in the same KV cache group must use the same "
588
- "quantization method, tokens per state, model version, and storage "
589
- "block size."
686
+ "quantization method, tokens per state, model version, cache role, "
687
+ "index-sharing role, and storage block size."
688
)
689
merged_spec = cls(
690
block_size=specs[0].block_size,
@@ -600,6 +698,8 @@ class MLAAttentionSpec(FullAttentionSpec):
698
cache_dtype_str=cache_dtype_str_set.pop(),
699
tokens_per_state=tokens_per_state_set.pop(),
700
model_version=model_version_set.pop(),
701
+ cache_role=cache_role_set.pop(),
702
+ is_index_group_leader=index_group_leader_set.pop(),
703
storage_block_size=storage_block_size_set.pop(),
704
non_causal_multi_token_decode=any(
705
spec.non_causal_multi_token_decode for spec in specs
@@ -772,6 +872,10 @@ class CircularBufferSpec(AttentionSpec):
872
reads the open group's committed keys from the ring.
873
"""
874
875
+ @property
876
+ def block_table_token_alignment(self) -> int | None:
877
+ return None
878
+
879
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
880
# The ring occupies one block per request for its whole lifetime.
881
del vllm_config
@@ -792,6 +896,10 @@ class CircularBufferSpec(AttentionSpec):
896
def prefix_cacheable(self) -> bool:
897
return False
898
899
+ @property
900
+ def uses_slot_mapping(self) -> bool:
901
+ return False
902
+
903
904
@dataclass(frozen=True, kw_only=True)
905
class SlidingWindowMLASpec(SlidingWindowSpec):
@@ -919,6 +1027,10 @@ class MambaSpec(KVCacheSpec):
1027
return self.page_size_padded
1028
return page_size
1029
1030
+ @property
1031
+ def block_table_token_alignment(self) -> int | None:
1032
+ return None
1033
+
1034
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
1035
if vllm_config.cache_config.mamba_cache_mode == "all":
1036
max_model_len = vllm_config.model_config.max_model_len
@@ -1093,6 +1205,10 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
1205
def page_size_bytes(self) -> int:
1206
return sum(spec.page_size_bytes for spec in self.kv_cache_specs.values())
1207
1208
+ @property
1209
+ def block_table_token_alignment(self) -> int | None:
1210
+ return self.first_spec.block_table_token_alignment
1211
+
1212
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
1213
max_num_pages = max(
1214
cdiv(spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes)
@@ -1154,6 +1270,10 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
1270
for spec in self.kv_cache_specs.values()
1271
)
1272
1273
+ @property
1274
+ def uses_slot_mapping(self) -> bool:
1275
+ return self.first_spec.uses_slot_mapping
1276
+
1277
1278
def iter_layer_specs(kv_cache_spec: KVCacheSpec) -> Collection[KVCacheSpec]:
1279
"""The per-layer specs a KV cache group spec covers.
@@ -1259,6 +1379,13 @@ class KVCacheTensor:
1379
layer_stride: int
1380
block_stride: int
1381
offset: int = 0 # byte offset of layers[0]'s block 0
1382
+ host_resident: bool = False
1383
+
1384
+
1385
+class KVCacheGroupRole(str, Enum):
1386
+ DEFAULT = "default"
1387
+ HISPARSE_SOURCE = "hisparse_source"
1388
+ HISPARSE_INDEXER = "hisparse_indexer"
1389
1390
1391
@dataclass
@@ -1274,8 +1401,13 @@ class KVCacheGroupSpec:
1401
kv_cache_spec: KVCacheSpec
1402
# Whether this group contains EAGLE/MTP draft attention layers.
1403
is_eagle_group: bool = False
1404
+ # Host groups use the dedicated HiSparse pool; others share the device pool.
1405
+ host_resident: bool = False
1406
# Whether this group is part of the externally transferable KV state.
1407
+ # Ephemeral accelerator-side replicas are rebuilt from their transferable
1408
+ # source group after a connector load.
1409
enable_kv_transfer: bool = True
1410
+ role: KVCacheGroupRole = KVCacheGroupRole.DEFAULT
1411
1412
1413
@dataclass
@@ -1300,6 +1432,8 @@ class KVCacheConfig:
1432
"""Resolved retention policy for local prefix-cache checkpoints."""
1433
kv_cache_layout: str | None = None
1434
"""The KV cache layout resolved by the engine core, adopted by all workers."""
1435
+ hisparse_host_num_blocks: int | None = None
1436
+ """Capacity of the dedicated HiSparse host-block manager, when enabled."""
1437
1438
@cached_property
1439
def transfer_group_ids(self) -> tuple[int, ...]:
@@ -1337,6 +1471,21 @@ class KVCacheConfig:
1471
)
1472
return tuple(block_ids[group_id] for group_id in self.transfer_group_ids)
1473
1474
+ @property
1475
+ def host_group_ids(self) -> tuple[int, ...]:
1476
+ return tuple(
1477
+ group_id
1478
+ for group_id, group in enumerate(self.kv_cache_groups)
1479
+ if group.host_resident
1480
+ )
1481
+
1482
+ def num_blocks_of(self, tensor: KVCacheTensor) -> int:
1483
+ """Number of blocks addressable by the pool backing ``tensor``."""
1484
+ if not tensor.host_resident:
1485
+ return self.num_blocks
1486
+ assert self.hisparse_host_num_blocks is not None
1487
+ return self.hisparse_host_num_blocks
1488
+
1489
@property
1490
def has_mamba_layers(self) -> bool:
1491
return any(
@@ -1347,9 +1496,11 @@ class KVCacheConfig:
1496
1497
@property
1498
def has_mixed_precision_kv_cache(self) -> bool:
1350
- """Whether attention groups store their KV cache at more than one precision."""
1499
+ """Whether device attention caches use more than one precision."""
1500
kv_cache_precisions: set[tuple[torch.dtype, KVQuantMode]] = set()
1501
for group in self.kv_cache_groups:
1502
+ if group.host_resident:
1503
+ continue
1504
kv_cache_precisions.update(
1505
(spec.dtype, spec.kv_quant_mode)
1506
for spec in iter_layer_specs(group.kv_cache_spec)
vllm/v1/kv_cache_spec_registry.py
+16
-1
@@ -34,6 +34,7 @@ class KVCacheSpecMetadata:
34
35
36
_REGISTRY_KVCACHESPEC_LIST: dict[type["KVCacheSpec"], KVCacheSpecMetadata] = {}
37
+_REGISTRY_ROLE_MANAGERS: dict[str, type["SingleTypeKVCacheManager"]] = {}
38
39
40
class KVCacheSpecRegistry:
@@ -102,20 +103,34 @@ class KVCacheSpecRegistry:
103
uniform_type_base_spec=uniform_type_base_spec,
104
)
105
106
+ @classmethod
107
+ def register_role_manager(
108
+ cls, role: str, manager_class: type["SingleTypeKVCacheManager"]
109
+ ) -> None:
110
+ """Select ``manager_class`` for every group with the given role."""
111
+ registered = _REGISTRY_ROLE_MANAGERS.get(role)
112
+ assert registered is None or registered is manager_class, (
113
+ f"Conflicting manager registration for group role {role}"
114
+ )
115
+ _REGISTRY_ROLE_MANAGERS[role] = manager_class
116
+
117
@classmethod
118
def get_manager_class(
107
- cls, kvcache_spec: "KVCacheSpec"
119
+ cls, kvcache_spec: "KVCacheSpec", role: str | None = None
120
) -> type["SingleTypeKVCacheManager"] | None:
121
"""
122
Get the single type kvcache manager class for a given kvcache spec instance.
123
124
Args:
125
kvcache_spec: A KVCacheSpec instance
126
+ role: The group's role; a role-registered manager takes precedence.
127
128
Returns:
129
The SingleTypeKVCacheManager class to use for this kvcache_spec
130
"""
131
cls._ensure_registered()
132
+ if role is not None and (manager_class := _REGISTRY_ROLE_MANAGERS.get(role)):
133
+ return manager_class
134
kvcache_spec_cls = type(kvcache_spec)
135
136
# Walk up the MRO to find a registered base class
vllm/v1/kv_offload/config.py
+2
@@ -14,6 +14,8 @@ class OffloadingGroupConfig:
14
tokens_per_block: int
15
# Layer names belonging to this group.
16
layer_names: tuple[str, ...]
17
+ # Original KVCacheConfig group index.
18
+ group_id: int
19
20
21
@dataclass(frozen=True)
vllm/v1/outputs.py
+4
@@ -281,6 +281,9 @@ class KVConnectorOutput:
281
# IDs of externally computed KV blocks that failed to load.
282
# Requests referencing these blocks should be rescheduled to recompute them
283
invalid_block_ids: set[int] = field(default_factory=set)
284
+ # Receive failures keyed by request identity. This remains unambiguous for
285
+ # hybrid/multi-pool cache layouts where numeric block IDs overlap.
286
+ failed_recving: set[str] = field(default_factory=set)
287
# Configuration describing how many finished sending/receiving
288
# notifications should be expected for each request. This allows
289
# handshake-based connectors like Nixl to update the KVOutputAggregator.
@@ -295,6 +298,7 @@ class KVConnectorOutput:
298
and not self.kv_connector_stats
299
and not self.kv_cache_events
300
and not self.invalid_block_ids
301
+ and not self.failed_recving
302
and not self.kv_connector_worker_meta
303
)
304
vllm/v1/worker/gpu/attn_utils.py
+36
-11
@@ -18,6 +18,10 @@ from vllm.v1.attention.backend import (
18
CommonAttentionMetadata,
19
)
20
from vllm.v1.attention.backends.utils import create_fast_prefill_custom_backend
21
+from vllm.v1.hisparse.binding import (
22
+ init_hisparse_kv_cache,
23
+ resolve_hisparse_block_size,
24
+)
25
from vllm.v1.kv_cache_interface import (
26
AttentionSpec,
27
KVCacheConfig,
@@ -30,11 +34,12 @@ from vllm.v1.worker.utils import (
34
AttentionGroup,
35
add_kv_sharing_layers_to_kv_cache_groups,
36
allocate_kv_cache,
33
- bind_kv_cache,
37
+ bind_kv_cache_to_layers,
38
prepare_kernel_block_sizes,
39
)
40
41
if TYPE_CHECKING:
42
+ from vllm.v1.worker.gpu.block_table import BlockTables
43
from vllm.v1.worker.gpu.cudagraph_utils import (
44
BatchExecutionDescriptor,
45
CudaGraphManager,
@@ -132,6 +137,7 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
137
if isinstance(spec, AttentionSpec):
138
spec = attn_module.get_attn_backend().customize_spec(spec)
139
kv_cache_spec[layer_name] = spec
140
+ resolve_hisparse_block_size(vllm_config, kv_cache_spec, attn_layers)
141
return kv_cache_spec
142
143
@@ -198,6 +204,9 @@ def init_attn_backend(
204
kv_cache_config.kv_cache_groups
205
):
206
layer_names = kv_cache_group_spec.layer_names
207
+ if not kv_cache_group_spec.kv_cache_spec.has_layer_views:
208
+ attn_groups.append([])
209
+ continue
210
if active_layer_names is not None:
211
layer_names = list(active_layer_names.intersection(layer_names))
212
@@ -316,22 +325,34 @@ def get_query_lens_mismatch_unsupported_backend(
325
326
327
def init_kv_cache(
319
- runner_kv_caches: list[torch.Tensor | list[torch.Tensor]],
328
forward_context: dict[str, Any],
329
kv_cache_config: KVCacheConfig,
330
device: torch.device,
331
kernel_block_sizes: list[int],
332
vllm_config: VllmConfig,
333
kv_cache_allocation_context: AbstractContextManager | None = None,
334
+ *,
335
+ block_tables: "BlockTables | None" = None,
336
) -> dict[str, Any]:
337
allocation_context = kv_cache_allocation_context or nullcontext()
338
with allocation_context:
329
- kv_caches = allocate_kv_cache(
330
- kv_cache_config,
331
- device,
332
- vllm_config.cache_config.get_resolved_kv_cache_layout(),
333
- kernel_block_sizes,
334
- )
339
+ if vllm_config.attention_config.hisparse_config is not None:
340
+ assert block_tables is not None
341
+ kv_caches = init_hisparse_kv_cache(
342
+ kv_cache_config,
343
+ device,
344
+ kernel_block_sizes,
345
+ vllm_config,
346
+ forward_context,
347
+ block_tables,
348
+ )
349
+ else:
350
+ kv_caches = allocate_kv_cache(
351
+ kv_cache_config,
352
+ device,
353
+ vllm_config.cache_config.get_resolved_kv_cache_layout(),
354
+ kernel_block_sizes,
355
+ )
356
for layer_name, target in get_shared_kv_cache_layers(vllm_config).items():
357
kv_caches[layer_name] = kv_caches[target]
358
# Dual-attention models (e.g. LongCat-Flash) put two Attention modules per
@@ -342,10 +363,12 @@ def init_kv_cache(
363
in ("longcat_flash", "longcat_flash_ngram")
364
else 1
365
)
345
- bind_kv_cache(
346
- kv_caches,
366
+ bindable_caches = {
367
+ name: cache for name, cache in kv_caches.items() if name in forward_context
368
+ }
369
+ bind_kv_cache_to_layers(
370
+ bindable_caches,
371
forward_context,
348
- runner_kv_caches,
372
num_attn_module,
373
kv_cache_groups=kv_cache_config.kv_cache_groups,
374
)
@@ -396,6 +419,8 @@ def build_attn_metadata(
419
attn_metadata: dict[str, Any] = {}
420
num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
421
for i in range(num_kv_cache_groups):
422
+ if not attn_groups[i]:
423
+ continue
424
block_table = block_tables[i]
425
slot_mapping = slot_mappings[i]
426
# Per-group causal for hybrid drafters (mixed SWA/full attention).
vllm/v1/worker/gpu/cudagraph_utils.py
+2
@@ -34,6 +34,7 @@ from vllm.platforms import current_platform
34
from vllm.sequence import IntermediateTensors
35
from vllm.utils.math_utils import round_up
36
from vllm.utils.torch_utils import current_stream
37
+from vllm.v1.hisparse.binding import release_hisparse_profiling_cache
38
from vllm.v1.kv_cache_interface import KVCacheConfig
39
from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
40
from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
@@ -931,6 +932,7 @@ def _teardown_profiling_state(runner: "GPUModelRunner") -> None:
932
if (model := getattr(runner, "model", None)) is not None:
933
layers = itertools.chain(layers, model.modules())
934
clear_layer_kv_caches(layers)
935
+ release_hisparse_profiling_cache(runner.compilation_config.static_forward_context)
936
runner.cache_config.num_gpu_blocks = None
937
runner.maybe_remove_all_loras(runner.lora_config)
938
gc.collect()
vllm/v1/worker/gpu/kv_connector.py
+31
-15
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import TYPE_CHECKING
3
+from typing import TYPE_CHECKING, Any
4
5
import torch
6
@@ -29,7 +29,10 @@ if TYPE_CHECKING:
29
class KVConnector:
30
"""KVConnector interface used by GPUModelRunner."""
31
32
- def pre_forward(self, scheduler_output: "SchedulerOutput") -> None:
32
+ def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> None:
33
+ pass
34
+
35
+ def finish_forward(self) -> None:
36
pass
37
38
def post_forward(
@@ -43,6 +46,9 @@ class KVConnector:
46
def set_disabled(self, disabled: bool) -> None:
47
pass
48
49
+ def reset_capture_state(self) -> None:
50
+ pass
51
+
52
53
class ActiveKVConnector(KVConnector):
54
def __init__(
@@ -54,10 +60,10 @@ class ActiveKVConnector(KVConnector):
60
self.kv_connector.register_kv_caches(kv_caches_dict)
61
self.kv_connector.set_host_xfer_buffer_ops(copy_kv_blocks)
62
57
- self._pending_load_start = False
63
+ self._pending_load_kwargs: dict[str, Any] | None = None
64
self._disabled = False
65
60
- def pre_forward(self, scheduler_output: "SchedulerOutput") -> None:
66
+ def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> None:
67
if self._disabled:
68
return
69
@@ -65,23 +71,31 @@ class ActiveKVConnector(KVConnector):
71
assert kv_connector_metadata is not None
72
self.kv_connector.handle_preemptions(kv_connector_metadata)
73
self.kv_connector.bind_connector_metadata(kv_connector_metadata)
74
+ self._pending_load_kwargs = kwargs
75
76
if scheduler_output.has_sync_kv_loads:
77
# Sync loads need to run before this step's forward.
78
self._start_load_kv()
72
- else:
73
- # Start any async loads in post-forward instead, keeping
74
- # their host-side submission cost off the critical path.
75
- self._pending_load_start = True
79
+ # Otherwise start the async load after forward to keep submission
80
+ # off the critical path.
81
82
def _start_load_kv(self) -> None:
78
- self._pending_load_start = False
83
+ load_kwargs = self._pending_load_kwargs
84
+ assert load_kwargs is not None
85
+ self._pending_load_kwargs = None
86
# TODO: sort out KV Connectors' use of forward_context
87
if is_forward_context_available():
81
- self.kv_connector.start_load_kv(get_forward_context())
88
+ self.kv_connector.start_load_kv(get_forward_context(), **load_kwargs)
89
else:
90
with set_forward_context(None, self.vllm_config):
84
- self.kv_connector.start_load_kv(get_forward_context())
91
+ self.kv_connector.start_load_kv(get_forward_context(), **load_kwargs)
92
+
93
+ def finish_forward(self) -> None:
94
+ if not self._disabled:
95
+ self.kv_connector.finish_forward()
96
+
97
+ def reset_capture_state(self) -> None:
98
+ self.kv_connector.reset_capture_state()
99
100
def post_forward(
101
self, finished_req_ids: set[str], wait_for_save: bool = True
@@ -89,15 +103,16 @@ class ActiveKVConnector(KVConnector):
103
if self._disabled:
104
return None
105
92
- if self._pending_load_start:
106
+ if self._pending_load_kwargs is not None:
107
self._start_load_kv()
108
109
output = KVConnectorOutput()
110
if wait_for_save:
111
self.kv_connector.wait_for_save()
98
- output.finished_sending, output.finished_recving = (
99
- self.kv_connector.get_finished(finished_req_ids)
100
- )
112
+ transfer_results = self.kv_connector.get_transfer_results(finished_req_ids)
113
+ output.finished_sending = transfer_results.finished_sending or None
114
+ output.finished_recving = transfer_results.finished_recving or None
115
+ output.failed_recving = transfer_results.failed_recving
116
output.invalid_block_ids = self.kv_connector.get_block_ids_with_load_errors()
117
output.kv_connector_stats = self.kv_connector.get_kv_connector_stats()
118
output.kv_cache_events = self.kv_connector.get_kv_connector_kv_cache_events()
@@ -112,6 +127,7 @@ class ActiveKVConnector(KVConnector):
127
return EMPTY_MODEL_RUNNER_OUTPUT
128
129
self.pre_forward(scheduler_output)
130
+ self.finish_forward()
131
finished_req_ids = scheduler_output.finished_req_ids
132
kv_connector_output = self.post_forward(finished_req_ids, wait_for_save=False)
133
return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
vllm/v1/worker/gpu/model_runner.py
+29
-13
@@ -67,7 +67,6 @@ from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
67
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
68
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
69
from vllm.v1.kv_cache_interface import (
70
- CircularBufferSpec,
70
KVCacheConfig,
71
MambaSpec,
72
UniformTypeKVCacheSpecs,
@@ -590,7 +589,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
589
layer_spec = (
590
spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec
591
)
593
- slot_mapping_enabled.append(not isinstance(layer_spec, CircularBufferSpec))
592
+ slot_mapping_enabled.append(layer_spec.uses_slot_mapping)
593
# Let each cache type account for CP. Attention KV is DCP-sharded,
594
# while Mamba/GDN recurrent state is replicated across DCP ranks.
595
max_num_blocks = spec.max_num_blocks_per_req(
@@ -598,12 +597,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
597
)
598
# Preserve each cache type's alignment requirements after applying
599
# its topology-aware block-table width.
601
- if isinstance(layer_spec, (MambaSpec, CircularBufferSpec)):
602
- max_num_blocks = get_block_table_width(
603
- max_num_blocks, spec.block_size, token_alignment=None
604
- )
605
- else:
606
- max_num_blocks = get_block_table_width(max_num_blocks, spec.block_size)
600
+ max_num_blocks = get_block_table_width(
601
+ max_num_blocks,
602
+ spec.block_size,
603
+ token_alignment=spec.block_table_token_alignment,
604
+ )
605
max_num_blocks_per_group.append(max_num_blocks)
606
607
target_attn_layer_names = None
@@ -722,18 +720,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):
720
# to its own attention support.
721
self.speculator.init_cudagraph_manager(cudagraph_mode)
722
725
- self.kv_caches: list[torch.Tensor] = []
723
# Capture warmup providers that depend on allocated KV-cache strides.
724
with self.jit_warmup_registry.activate():
725
kv_caches_dict = init_kv_cache(
729
- self.kv_caches,
726
self.compilation_config.static_forward_context,
727
self.kv_cache_config,
728
self.device,
729
self.kernel_block_sizes,
730
self.vllm_config,
731
kv_cache_allocation_context=kv_cache_allocation_context,
732
+ block_tables=self.block_tables,
733
)
734
+ self.kv_caches = [
735
+ cache for cache in kv_caches_dict.values() if cache.device == self.device
736
+ ]
737
if is_profiling:
738
self.kv_connector = NO_OP_KV_CONNECTOR
739
else:
@@ -1014,6 +1014,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1014
):
1015
self._dummy_run(**batch)
1016
self.adaptive_verification.set_initial_cost_curves(timings)
1017
+ self.kv_connector.reset_capture_state()
1018
1019
end_free_gpu_memory = torch.accelerator.get_memory_info()[0]
1020
@@ -1141,13 +1142,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1142
def update_requests(self, scheduler_output: SchedulerOutput) -> None:
1143
# Add new blocks and update num_computed_tokens for the existing requests.
1144
reqs = scheduler_output.scheduled_cached_reqs
1145
+ table_updates = scheduler_output.block_table_updates or {}
1146
+ for req_id, block_ids in table_updates.items():
1147
+ if (idx := self.req_states.req_id_to_index.get(req_id)) is not None:
1148
+ self.block_tables.append_block_ids(idx, block_ids, overwrite=True)
1149
num_computed_tokens_np = self.req_states.num_computed_tokens_np
1150
for req_id, num_computed_tokens, req_new_block_ids in zip(
1151
reqs.req_ids, reqs.num_computed_tokens, reqs.new_block_ids
1152
):
1153
req_index = self.req_states.req_id_to_index[req_id]
1154
num_computed_tokens_np[req_index] = num_computed_tokens
1150
- if req_new_block_ids is not None:
1155
+ if req_new_block_ids is not None and req_id not in table_updates:
1156
self.block_tables.append_block_ids(
1157
req_index, req_new_block_ids, overwrite=False
1158
)
@@ -1841,13 +1846,22 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1846
)
1847
self.step_timing.forward_start()
1848
1849
+ connector_kwargs = dict(
1850
+ scheduler_output=scheduler_output,
1851
+ request_state_indices=input_batch.idx_mapping,
1852
+ request_ids=input_batch.req_ids,
1853
+ num_tokens=input_batch.num_tokens,
1854
+ )
1855
+
1856
# Run model.
1857
if batch_desc.cg_mode == CUDAGraphMode.FULL:
1858
# Use explicit cudagraph replay for FULL mode.
1859
# NOTE(woosuk): Here, we don't need to pass the input tensors,
1860
# because they are already copied to the CUDA graph input buffers.
1861
assert self.cudagraph_manager is not None
1850
- self.kv_connector.pre_forward(scheduler_output)
1862
+ self.kv_connector.pre_forward(
1863
+ **connector_kwargs, attn_metadata=attn_metadata
1864
+ )
1865
model_output = self.cudagraph_manager.run_fullgraph(batch_desc)
1866
else:
1867
# For piecewise and eager mode, just call model().
@@ -1871,7 +1885,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1885
skip_compiled=skip_compiled,
1886
is_padding=input_batch.is_padding,
1887
):
1874
- self.kv_connector.pre_forward(scheduler_output)
1888
+ self.kv_connector.pre_forward(**connector_kwargs)
1889
if ubatch_state is not None:
1890
assert self.ubatch_runner is not None
1891
model_output = self.ubatch_runner.run(
@@ -1889,6 +1903,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1903
# Eager (NONE): call the raw model directly.
1904
model_output = self.model(**model_inputs)
1905
1906
+ self.kv_connector.finish_forward()
1907
+
1908
if self.is_last_pp_rank:
1909
if self.use_aux_hidden_state_outputs:
1910
assert isinstance(model_output, tuple)
vllm/v1/worker/gpu/warmup.py
+10
@@ -224,10 +224,20 @@ def warmup_kernels(
224
# fixed draft counts here, then restore the manager for capture and serving.
225
adaptive_verification = model_runner.adaptive_verification
226
model_runner.adaptive_verification = None
227
+ rejection_sampler = model_runner.rejection_sampler
228
+ adaptive_sampling = (
229
+ rejection_sampler is not None and rejection_sampler.enable_adaptive_verification
230
+ )
231
+ if adaptive_sampling:
232
+ assert rejection_sampler is not None
233
+ rejection_sampler.enable_adaptive_verification = False
234
try:
235
_warmup_kernels(model_runner, worker_execute_model, worker_sample_tokens)
236
finally:
237
model_runner.adaptive_verification = adaptive_verification
238
+ if adaptive_sampling:
239
+ assert rejection_sampler is not None
240
+ rejection_sampler.enable_adaptive_verification = True
241
242
243
def _warmup_kernels(
vllm/v1/worker/kv_connector_model_runner_mixin.py
+5
-2
@@ -94,9 +94,12 @@ class KVConnectorModelRunnerMixin:
94
if wait_for_save and not defer_finalize:
95
kv_connector.wait_for_save()
96
97
- output.finished_sending, output.finished_recving = (
98
- kv_connector.get_finished(scheduler_output.finished_req_ids)
97
+ transfer_results = kv_connector.get_transfer_results(
98
+ scheduler_output.finished_req_ids
99
)
100
+ output.finished_sending = transfer_results.finished_sending
101
+ output.finished_recving = transfer_results.finished_recving
102
+ output.failed_recving = transfer_results.failed_recving
103
output.invalid_block_ids = kv_connector.get_block_ids_with_load_errors()
104
105
output.kv_connector_stats = kv_connector.get_kv_connector_stats()
vllm/v1/worker/utils.py
+36
-25
@@ -166,6 +166,8 @@ class KVBlockZeroer:
166
kv = static_forward_context[layer_name].kv_cache
167
if not isinstance(kv, torch.Tensor):
168
continue
169
+ if kv.device.type != self.device.type:
170
+ continue
171
dp = kv.data_ptr()
172
173
assert kv.shape[0] % num_blocks == 0, (
@@ -366,32 +368,22 @@ def select_common_block_size(
368
return False
369
return True
370
369
- # Case 1: if the block_size of kv cache manager is supported by all backends,
370
- # return it directly.
371
if block_size_is_supported(backends, kv_manager_block_size):
372
return kv_manager_block_size
373
374
- # Case 2: otherwise, the block_size must be an `int`-format supported size of
375
- # at least one backend. Iterate over all `int`-format supported sizes in
376
- # descending order and return the first one that is supported by all backends.
377
- # Simple proof:
378
- # If the supported size b is in MultipleOf(x_i) format for all attention
379
- # backends i, and b a factor of kv_manager_block_size, then
380
- # kv_manager_block_size also satisfies MultipleOf(x_i) for all i. We will
381
- # return kv_manager_block_size in case 1.
382
- all_int_supported_sizes = set(
383
- supported_size
374
+ # MultipleOf constraints also accept the manager size if they accept a divisor.
375
+ # Any remaining candidate must therefore be an explicit size from a backend.
376
+ candidates = {
377
+ size
378
for backend in backends
385
- for supported_size in backend.get_supported_kernel_block_sizes()
386
- if isinstance(supported_size, int)
387
- )
379
+ for size in backend.get_supported_kernel_block_sizes()
380
+ if isinstance(size, int) and kv_manager_block_size % size == 0
381
+ }
382
389
- for supported_size in sorted(all_int_supported_sizes, reverse=True):
390
- if kv_manager_block_size % supported_size != 0:
391
- continue
392
- if block_size_is_supported(backends, supported_size):
393
- return supported_size
394
- raise ValueError(f"No common block size for {kv_manager_block_size}. ")
383
+ for size in sorted(candidates, reverse=True):
384
+ if block_size_is_supported(backends, size):
385
+ return size
386
+ raise ValueError(f"No common block size for {kv_manager_block_size}.")
387
388
389
def allocate_kv_cache(
@@ -440,7 +432,11 @@ def allocate_kv_cache(
432
if isinstance(spec, UniformTypeKVCacheSpecs):
433
spec = spec.kv_cache_specs[layer_name]
434
443
- num_blocks = kv_cache_config.num_blocks
435
+ if not spec.has_layer_views:
436
+ kv_caches.update((name, buf) for name in tensor.layers)
437
+ continue
438
+
439
+ num_blocks = kv_cache_config.num_blocks_of(tensor)
440
kernel_block_size = None
441
if kernel_block_sizes is not None and group_id < len(kernel_block_sizes):
442
kernel_block_size = kernel_block_sizes[group_id]
@@ -485,7 +481,9 @@ def prepare_kernel_block_sizes(
481
kv_cache_spec = next(iter(kv_cache_spec.kv_cache_specs.values()))
482
if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec):
483
continue
488
- if isinstance(kv_cache_spec, AttentionSpec):
484
+ if not kv_cache_spec.has_layer_views:
485
+ kernel_block_sizes.append(kv_cache_spec.block_size)
486
+ elif isinstance(kv_cache_spec, AttentionSpec):
487
# This is an attention backend that supports virtual block splitting.
488
kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size
489
group_backends = [g.backend for g in attn_groups[kv_cache_gid]]
@@ -621,7 +619,6 @@ def bind_kv_cache(
619
for layer_name in kv_caches:
620
index2name[extract_layer_index(layer_name, num_attn_module)].append(layer_name)
621
624
- ordered_layer_names: list[str] = []
622
for layer_index in sorted(index2name.keys()):
623
layer_names = index2name[layer_index]
624
if len(layer_names) > 1:
@@ -635,8 +632,19 @@ def bind_kv_cache(
632
current_platform.check_runner_kv_caches_multi_layer()
633
for layer_name in layer_names:
634
runner_kv_caches.append(kv_caches[layer_name])
638
- ordered_layer_names.append(layer_name)
635
636
+ bind_kv_cache_to_layers(
637
+ kv_caches, forward_context, num_attn_module, kv_cache_groups
638
+ )
639
+
640
+
641
+def bind_kv_cache_to_layers(
642
+ kv_caches: dict[str, torch.Tensor],
643
+ forward_context: dict[str, Attention],
644
+ num_attn_module: int = 1,
645
+ kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None,
646
+) -> None:
647
+ """Bind layer caches and share ReplaySSM trackers in model-layer order."""
648
# Bind kv_caches to forward context. Each layer's bind_kv_cache unpacks
649
# its raw allocation into the per-layer view(s) it needs (e.g. Mamba
650
# splits conv/ssm), so the kv_caches dict can hold a single tensor per
@@ -644,6 +652,9 @@ def bind_kv_cache(
652
for layer_name, kv_cache in kv_caches.items():
653
forward_context[layer_name].bind_kv_cache(kv_cache)
654
655
+ ordered_layer_names = sorted(
656
+ kv_caches, key=lambda name: extract_layer_index(name, num_attn_module)
657
+ )
658
share_replayssm_ring_trackers(ordered_layer_names, forward_context, kv_cache_groups)
659
660