@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 2798f66860

[Bugfix][HiSparse] Resolve MTP verification rows with a union residency kernel (#59235)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Matthew Bonanni committed Sep 29, 2026 at 23:01 UTC 2798f668608155bc8a8c74cb97e3ddd0d3053085
13 files changed +694 -431
csrc/libtorch_stable/hisparse_kernels.cu
+396 -201
@@ -30,9 +30,7 @@
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;
33 +// Empty key in the shared open-addressing table of host rows.
34 constexpr int32_t kHashEmpty = -1;
35
36 bool is_pinned_cpu_tensor(const torch::stable::Tensor& tensor) {
@@ -46,10 +44,15 @@ bool is_pinned_cpu_tensor(const torch::stable::Tensor& tensor) {
44 return attributes.type == cudaMemoryTypeHost;
45 }
46
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));
47 +__device__ __forceinline__ int32_t hash_slot(int32_t key, int size) {
48 + // Knuth multiplicative hash scaled onto [0, size) by a multiply-high, so
49 + // the table can take any size without a division.
50 + return static_cast<int32_t>(
51 + __umulhi(static_cast<uint32_t>(key) * 2654435761u, size));
52 +}
53 +
54 +__device__ __forceinline__ int next_slot(int h, int size) {
55 + return h + 1 == size ? 0 : h + 1;
56 }
57
58 // Copy one row of `row_bytes` bytes with a single warp. MLA rows take the
@@ -176,16 +179,117 @@ __device__ __forceinline__ void store_hot_index(
179 }
180 }
181
179 -// One block per batch row.
182 +// A union entry's value holds its first (row, top-k) position until phase 2
183 +// or 3 resolves it to a hot slot tagged with how it was resolved.
184 +constexpr int32_t kEntryHit = 1 << 30;
185 +constexpr int32_t kEntryMiss = 1 << 29;
186 +constexpr int32_t kEntryInvalid = 1 << 28;
187 +constexpr int32_t kEntrySlotMask = 0xFFFF;
188 +constexpr int kResidencyCounters = 5;
189 +// Per-request table entries per referenced top-k position (load <= 1/4
190 +// until the hash_size cap).
191 +constexpr int kHashLoadInverse = 4;
192 +
193 +// Translate a request-relative top-k position to its resident GPU row (-1 when
194 +// the page is not resident) and returns its host row (-1 when not host-backed).
195 +__device__ __forceinline__ int32_t translate_topk_entry(
196 + const int32_t token_index, const int32_t request_row,
197 + const int32_t* __restrict__ source_block_table,
198 + const int32_t* __restrict__ resident_block_table, const int64_t host_rows,
199 + const int64_t source_bt_stride, const int32_t source_num_reqs,
200 + const int32_t source_num_blocks, const int32_t source_block_size,
201 + const int64_t resident_bt_stride, const int32_t resident_num_reqs,
202 + const int32_t resident_num_blocks, const int32_t resident_block_size,
203 + const int32_t resident_null_block, int32_t& resident_row) {
204 + int32_t g = token_index;
205 + resident_row = -1;
206 + if (source_block_table != nullptr) {
207 + const int32_t source_block =
208 + token_index >= 0 ? token_index / source_block_size : -1;
209 + if (request_row >= 0 && request_row < source_num_reqs &&
210 + source_block >= 0 && source_block < source_num_blocks) {
211 + const int32_t physical_block =
212 + source_block_table[static_cast<int64_t>(request_row) *
213 + source_bt_stride +
214 + source_block];
215 + g = physical_block > 0 ? physical_block * source_block_size +
216 + token_index % source_block_size
217 + : -1;
218 + } else {
219 + g = -1;
220 + }
221 + if (resident_block_table != nullptr) {
222 + const int32_t resident_block =
223 + token_index >= 0 ? token_index / resident_block_size : -1;
224 + if (request_row >= 0 && request_row < resident_num_reqs &&
225 + resident_block >= 0 && resident_block < resident_num_blocks) {
226 + const int32_t physical_block =
227 + resident_block_table[static_cast<int64_t>(request_row) *
228 + resident_bt_stride +
229 + resident_block];
230 + if (physical_block != resident_null_block && physical_block >= 0) {
231 + resident_row = physical_block * resident_block_size +
232 + token_index % resident_block_size;
233 + }
234 + }
235 + }
236 + }
237 + return g >= host_rows ? -1 : g;
238 +}
239 +
240 +// Write the hot index of one reference to a resolved union entry and tally it.
241 +__device__ __forceinline__ void write_entry_index(
242 + const int32_t entry, const int row, const int i, const int32_t top_k,
243 + int32_t* __restrict__ hot_indices, int32_t* __restrict__ attention_indices,
244 + const int64_t attention_row_stride, int32_t* __restrict__ miss_mask,
245 + const int32_t* __restrict__ hot_block_table, const int request_row,
246 + const int64_t hot_table_stride, const int32_t hot_block_size,
247 + const int64_t attention_block_stride, const bool first_reference, int& hits,
248 + int& misses) {
249 + int32_t* row_out = hot_indices + static_cast<int64_t>(row) * top_k;
250 + int32_t* row_attention =
251 + attention_indices != nullptr
252 + ? attention_indices + static_cast<int64_t>(row) * attention_row_stride
253 + : nullptr;
254 + if (entry & kEntryInvalid) {
255 + store_hot_index(row_out, row_attention, i, -1, hot_block_size,
256 + attention_block_stride);
257 + return;
258 + }
259 + store_hot_index(row_out, row_attention, i,
260 + static_cast<int32_t>(get_physical_hot_row(
261 + hot_block_table, request_row, hot_table_stride,
262 + hot_block_size, entry & kEntrySlotMask)),
263 + hot_block_size, attention_block_stride);
264 + // Only an entry's first reference loads it; later references reuse the
265 + // slot like a hit.
266 + if ((entry & kEntryMiss) && first_reference) {
267 + if (miss_mask != nullptr) {
268 + miss_mask[static_cast<int64_t>(row) * top_k + i] = 1;
269 + }
270 + ++misses;
271 + } else {
272 + ++hits;
273 + }
274 +}
275 +
276 +// One block per request, launched per batch row. A request's rows (several
277 +// with speculative decoding) are contiguous and share its hot-buffer state, so
278 +// the block of its first row resolves the union of their host misses once and
279 +// maps every row onto it; the other blocks exit. The union is bounded by the
280 +// caller's max_union_rows (every decode query's top-k), which
281 +// device_buffer_size covers; a larger union traps. Misses take eviction slots
282 +// in (row, top-k) order of first reference, so the plan is deterministic.
283 //
284 // 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
285 +// s_hash_keys[hash_size] open addressing: union of the rows' host ids
286 +// s_hash_vals[hash_size] first position, then slot | kEntry* tag
287 +// s_chunk_off[nbc + 1] prefix sums for hit, then miss, compaction
288 // 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
289 +// s_counters[5] distinct hits, union size, per-row valid count,
290 +// misses, per-entry hits
291 +// s_done[nbc] positions phase 3 resolved, one word per chunk
292 // s_lru_out[hot_size] int16, compacted slots: [hits fwd | evict bwd]
188 -// s_hash_vals[hash_size] int16 hash values (top-k index)
293 // Valid global ids must be unique within each row.
294 __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
295 const int32_t* __restrict__ hot_block_table, // [max_rows, hot_blocks]
@@ -218,11 +322,25 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
322 const int64_t attention_row_stride, const int64_t valid_count_stride) {
323 const int NUM_WARPS = blockDim.x / kWarpSize;
324 const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize;
221 - const int num_token_chunks = (top_k + kWarpSize - 1) / kWarpSize;
325
223 - const int batch_row = blockIdx.x;
326 + // The first row of each contiguous run of a request's rows leads it. FULL
327 + // CUDA-graph padding rows map to request 0 after the real rows, forming a
328 + // second run, but carry no valid top-k: that run returns before phase 2
329 + // touches the request's shared state.
330 + const int first_row = blockIdx.x;
331 const int request_row =
225 - request_ids != nullptr ? request_ids[batch_row] : batch_row;
332 + request_ids != nullptr ? request_ids[first_row] : first_row;
333 + if (request_ids != nullptr && first_row > 0 &&
334 + request_ids[first_row - 1] == request_row) {
335 + return;
336 + }
337 + int end_row = first_row + 1;
338 + if (request_ids != nullptr) {
339 + while (end_row < static_cast<int>(gridDim.x) &&
340 + request_ids[end_row] == request_row) {
341 + ++end_row;
342 + }
343 + }
344 const int state_row =
345 request_state_indices != nullptr && request_row >= 0 &&
346 request_row < request_state_count
@@ -230,23 +348,24 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
348 : (request_state_indices == nullptr ? request_row : -1);
349 // V2 publishes -1 for CUDA-graph padding rows.
350 if (state_row < 0) {
233 - for (int i = threadIdx.x; i < top_k; i += blockDim.x) {
234 - const int64_t index = static_cast<int64_t>(batch_row) * top_k + i;
235 - hot_indices[index] = -1;
236 - if (attention_indices != nullptr) {
237 - attention_indices[static_cast<int64_t>(batch_row) *
238 - attention_row_stride +
239 - i] = -1;
351 + for (int row = first_row; row < end_row; ++row) {
352 + for (int i = threadIdx.x; i < top_k; i += blockDim.x) {
353 + const int64_t index = static_cast<int64_t>(row) * top_k + i;
354 + hot_indices[index] = -1;
355 + if (attention_indices != nullptr) {
356 + attention_indices[static_cast<int64_t>(row) * attention_row_stride +
357 + i] = -1;
358 + }
359 + if (resolved_global_indices != nullptr) {
360 + resolved_global_indices[index] = -1;
361 + }
362 }
241 - if (resolved_global_indices != nullptr) {
242 - resolved_global_indices[index] = -1;
363 + if (valid_counts != nullptr && threadIdx.x == 0) {
364 + valid_counts[static_cast<int64_t>(row) * valid_count_stride] = 0;
365 + }
366 + if (swap_counts != nullptr && threadIdx.x == 0) {
367 + swap_counts[row] = 0;
368 }
244 - }
245 - if (valid_counts != nullptr && threadIdx.x == 0) {
246 - valid_counts[static_cast<int64_t>(batch_row) * valid_count_stride] = 0;
247 - }
248 - if (swap_counts != nullptr && threadIdx.x == 0) {
249 - swap_counts[batch_row] = 0;
369 }
370 return;
371 }
@@ -255,136 +374,130 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
374 const int lane_id = tid % kWarpSize;
375 const unsigned int lanes_before = ((unsigned int)1 << lane_id) - 1;
376
258 - const int32_t* row_topk =
259 - global_indices + static_cast<int64_t>(batch_row) * input_row_stride;
260 - int32_t* row_out = hot_indices + static_cast<int64_t>(batch_row) * top_k;
261 - int32_t* row_attention =
262 - attention_indices != nullptr
263 - ? attention_indices +
264 - static_cast<int64_t>(batch_row) * attention_row_stride
265 - : nullptr;
266 - int32_t* row_miss = (miss_mask != nullptr)
267 - ? miss_mask + static_cast<int64_t>(batch_row) * top_k
268 - : nullptr;
377 int32_t* row_dgi =
378 device_global_indices + static_cast<int64_t>(state_row) * region_stride;
379 int16_t* row_lru = lru_slots + static_cast<int64_t>(state_row) * hot_size;
380
381 + // A union within max_union_rows (hash_size - 1) and within the rows' top-k
382 + // leaves an empty slot, so every probe sequence ends. Use a slice sized to
383 + // this request's rows so single-row requests initialize a small table.
384 + const int table_size =
385 + min(hash_size, kHashLoadInverse * (end_row - first_row) * top_k);
386 +
387 extern __shared__ char smem_raw[];
274 - int32_t* s_topk = reinterpret_cast<int32_t*>(smem_raw);
275 - int32_t* s_chunk_off = s_topk + top_k;
388 + int32_t* s_hash_keys = reinterpret_cast<int32_t*>(smem_raw);
389 + int32_t* s_hash_vals = s_hash_keys + hash_size;
390 + int32_t* s_chunk_off = s_hash_vals + hash_size;
391 int32_t* s_evict_off = s_chunk_off + (num_buffer_chunks + 1);
277 - int32_t* s_hash_keys = s_evict_off + (num_buffer_chunks + 1);
278 - int32_t* s_counters = s_hash_keys + hash_size;
279 - int16_t* s_lru_out = reinterpret_cast<int16_t*>(s_counters + 3);
280 - int16_t* s_hash_vals = s_lru_out + hot_size;
392 + int32_t* s_counters = s_evict_off + (num_buffer_chunks + 1);
393 + unsigned int* s_done =
394 + reinterpret_cast<unsigned int*>(s_counters + kResidencyCounters);
395 + int16_t* s_lru_out = reinterpret_cast<int16_t*>(s_done + num_buffer_chunks);
396
282 - if (tid < 3) {
397 + for (int i = tid; i < table_size; i += blockDim.x) {
398 + s_hash_keys[i] = kHashEmpty;
399 + s_hash_vals[i] = INT32_MAX;
400 + }
401 + for (int i = tid; i < num_buffer_chunks + 1; i += blockDim.x) {
402 + s_chunk_off[i] = 0;
403 + s_evict_off[i] = 0;
404 + }
405 + if (tid < kResidencyCounters) {
406 s_counters[tid] = 0;
407 }
408 __syncthreads();
409
287 - // Phase 1: translate request-relative positions and resolve resident rows.
288 - for (int i = tid; i < top_k; i += blockDim.x) {
289 - const int32_t token_index = row_topk[i];
290 - int32_t g = token_index;
291 - int32_t resident_row = -1;
292 - if (source_block_table != nullptr) {
293 - const int32_t request_id = request_ids[batch_row];
294 - const int32_t source_block =
295 - token_index >= 0 ? token_index / source_block_size : -1;
296 - if (request_id >= 0 && request_id < source_num_reqs &&
297 - source_block >= 0 && source_block < source_num_blocks) {
298 - const int32_t physical_block =
299 - source_block_table[static_cast<int64_t>(request_id) *
300 - source_bt_stride +
301 - source_block];
302 - g = physical_block > 0 ? physical_block * source_block_size +
303 - token_index % source_block_size
304 - : -1;
305 - } else {
306 - g = -1;
410 + // Phase 1: translate every row's request-relative positions, resolve
411 + // resident rows directly, and collect the union of host rows.
412 + for (int row = first_row; row < end_row; ++row) {
413 + const int32_t* row_topk =
414 + global_indices + static_cast<int64_t>(row) * input_row_stride;
415 + int32_t* row_out = hot_indices + static_cast<int64_t>(row) * top_k;
416 + int32_t* row_attention =
417 + attention_indices != nullptr
418 + ? attention_indices +
419 + static_cast<int64_t>(row) * attention_row_stride
420 + : nullptr;
421 + int32_t* row_miss = miss_mask != nullptr
422 + ? miss_mask + static_cast<int64_t>(row) * top_k
423 + : nullptr;
424 + for (int i = tid; i < top_k; i += blockDim.x) {
425 + int32_t resident_row;
426 + const int32_t g = translate_topk_entry(
427 + row_topk[i], request_row, source_block_table, resident_block_table,
428 + host_rows, source_bt_stride, source_num_reqs, source_num_blocks,
429 + source_block_size, resident_bt_stride, resident_num_reqs,
430 + resident_num_blocks, resident_block_size, resident_null_block,
431 + resident_row);
432 + if (resolved_global_indices != nullptr) {
433 + resolved_global_indices[static_cast<int64_t>(row) * top_k + i] = g;
434 }
308 - if (resident_block_table != nullptr) {
309 - const int32_t resident_block =
310 - token_index >= 0 ? token_index / resident_block_size : -1;
311 - if (request_id >= 0 && request_id < resident_num_reqs &&
312 - resident_block >= 0 && resident_block < resident_num_blocks) {
313 - const int32_t physical_block =
314 - resident_block_table[static_cast<int64_t>(request_id) *
315 - resident_bt_stride +
316 - resident_block];
317 - if (physical_block != resident_null_block && physical_block >= 0) {
318 - resident_row = physical_block * resident_block_size +
319 - token_index % resident_block_size;
435 + if (resident_row >= 0 || g >= 0) atomicAdd(&s_counters[2], 1);
436 + if (row_miss != nullptr) row_miss[i] = 0;
437 + if (resident_row >= 0) {
438 + store_hot_index(row_out, row_attention, i, resident_row, hot_block_size,
439 + attention_block_stride);
440 + } else if (g < 0) {
441 + store_hot_index(row_out, row_attention, i, -1, hot_block_size,
442 + attention_block_stride);
443 + } else {
444 + int h = hash_slot(g, table_size);
445 + bool inserted = false;
446 + for (int probe = 0; probe < table_size; ++probe) {
447 + const int32_t old = atomicCAS(&s_hash_keys[h], kHashEmpty, g);
448 + if (old == kHashEmpty) {
449 + atomicAdd(&s_counters[1], 1);
450 + inserted = true;
451 + break;
452 + }
453 + if (old == g) {
454 + inserted = true;
455 + break;
456 }
457 + h = next_slot(h, table_size);
458 + }
459 + if (inserted) {
460 + atomicMin(&s_hash_vals[h], (row - first_row) * top_k + i);
461 + } else {
462 + // A full table: push the union count past the cap so the check
463 + // below traps.
464 + atomicAdd(&s_counters[1], hash_size);
465 }
466 }
467 }
324 - if (g >= host_rows) {
325 - g = -1;
326 - }
327 - if (resolved_global_indices != nullptr) {
328 - resolved_global_indices[static_cast<int64_t>(batch_row) * top_k + i] = g;
329 - }
330 - if (resident_row >= 0 || g >= 0) atomicAdd(&s_counters[2], 1);
331 - if (row_miss != nullptr) row_miss[i] = 0;
332 - if (resident_row >= 0) {
333 - store_hot_index(row_out, row_attention, i, resident_row, hot_block_size,
334 - attention_block_stride);
335 - s_topk[i] = kTokenDone;
336 - atomicAdd(&s_counters[1], 1);
337 - } else if (g < 0) {
338 - store_hot_index(row_out, row_attention, i, -1, hot_block_size,
339 - attention_block_stride);
340 - s_topk[i] = kTokenDone;
341 - atomicAdd(&s_counters[1], 1);
342 - } else {
343 - s_topk[i] = g;
468 + __syncthreads();
469 + if (tid == 0) {
470 + if (valid_counts != nullptr) {
471 + valid_counts[static_cast<int64_t>(row) * valid_count_stride] =
472 + s_counters[2];
473 + }
474 + s_counters[2] = 0;
475 }
476 + __syncthreads();
477 }
346 - __syncthreads();
347 - if (valid_counts != nullptr && tid == 0) {
348 - valid_counts[static_cast<int64_t>(batch_row) * valid_count_stride] =
349 - s_counters[2];
478 + // The table and eviction slots are sized for max_union_rows (hash_size - 1);
479 + // a larger union would never reach an empty probe slot or run past s_lru_out.
480 + if (s_counters[1] >= hash_size) {
481 + if (tid == 0) {
482 + printf(
483 + "HiSparse residency: request %d resolves %d host rows, above the "
484 + "configured bound of %d\n",
485 + request_row, s_counters[1], hash_size - 1);
486 + }
487 + __trap();
488 }
489 // Fully resident rows need only request-relative page translation. Avoid
490 // scanning or rewriting the hot LRU when no selected row can consult it.
353 - if (s_counters[1] == top_k) {
354 - if (tid == 0 && swap_counts != nullptr) {
355 - swap_counts[batch_row] = 0;
491 + if (s_counters[1] == 0) {
492 + for (int row = first_row + tid; row < end_row; row += blockDim.x) {
493 + if (swap_counts != nullptr) swap_counts[row] = 0;
494 }
495 return;
496 }
497
360 - for (int i = tid; i < hash_size; i += blockDim.x) {
361 - s_hash_keys[i] = kHashEmpty;
362 - }
363 - for (int i = tid; i < num_buffer_chunks + 1; i += blockDim.x) {
364 - s_chunk_off[i] = 0;
365 - s_evict_off[i] = 0;
366 - }
367 - __syncthreads();
368 -
369 - for (int i = tid; i < top_k; i += blockDim.x) {
370 - const int32_t g = s_topk[i];
371 - if (g == kTokenDone) {
372 - continue;
373 - }
374 - int slot = hash_slot(g, hash_size);
375 - while (true) {
376 - const int32_t old = atomicCAS(&s_hash_keys[slot], kHashEmpty, g);
377 - if (old == kHashEmpty || old == g) {
378 - s_hash_vals[slot] = static_cast<int16_t>(i);
379 - break;
380 - }
381 - slot = (slot + 1) % hash_size;
382 - }
383 - }
384 - __syncthreads();
385 -
498 // Phase 2: walk hot slots in LRU order, classify hit / evictable, and
387 - // compact them (hits forward, evictables backward) into s_lru_out.
499 + // compact them (hits forward, evictables backward) into s_lru_out. A hit
500 + // records its slot on the union entry.
501 const int iters_buffer = (num_buffer_chunks + NUM_WARPS - 1) / NUM_WARPS;
502 int total_hit_count = 0;
503 int total_evict_count = 0;
@@ -397,36 +510,28 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
510 const int16_t slot = has_valid_pos ? row_lru[pos] : int16_t(-1);
511 // Corruption tripwire: lru/dgi are long-lived device state; if some
512 // external writer (e.g. stray RDMA into reused VRAM) corrupts a slot
400 - // out of [0, region_stride), treat it as no-hit so it degrades to a re-miss
401 - // instead of an unbounded dgi read (phases 3/5 bound the write side).
513 + // out of [0, region_stride), treat it as no-hit so it degrades to a
514 + // re-miss instead of an unbounded dgi read (phase 3 bounds the write
515 + // side).
516 const int32_t cached_g =
517 (slot >= 0 && slot < region_stride) ? row_dgi[slot] : -1;
518
405 - int found_topk_idx = -1;
519 + bool is_hit = false;
520 if (cached_g >= 0) {
407 - int h = hash_slot(cached_g, hash_size);
521 + int h = hash_slot(cached_g, table_size);
522 while (true) {
523 const int32_t k = s_hash_keys[h];
524 if (k == cached_g) {
411 - found_topk_idx = static_cast<int32_t>(s_hash_vals[h]);
525 + s_hash_vals[h] = kEntryHit | slot;
526 + is_hit = true;
527 break;
528 }
529 if (k == kHashEmpty) break;
415 - h = (h + 1) % hash_size;
530 + h = next_slot(h, table_size);
531 }
532 }
418 - const bool is_hit = found_topk_idx >= 0;
533 const bool is_evictable = has_valid_pos && !is_hit;
534
421 - if (is_hit) {
422 - s_topk[found_topk_idx] = kTokenDone;
423 - store_hot_index(row_out, row_attention, found_topk_idx,
424 - static_cast<int32_t>(get_physical_hot_row(
425 - hot_block_table, request_row, hot_table_stride,
426 - hot_block_size, slot)),
427 - hot_block_size, attention_block_stride);
428 - }
429 -
535 int local_hit_off = 0;
536 int local_evict_off = 0;
537 if (has_valid_chunk) {
@@ -465,39 +570,59 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
570 }
571 __syncthreads();
572
468 - // Reset prefix sums for the miss compaction (token chunks <= buffer
469 - // chunks because hot_size >= top_k).
470 - for (int i = tid; i < num_token_chunks + 1; i += blockDim.x) {
573 + for (int i = tid; i < num_buffer_chunks + 1; i += blockDim.x) {
574 s_chunk_off[i] = 0;
575 }
576 __syncthreads();
577
475 - // Phase 3: compact misses, assign them eviction slots (oldest first) and
476 - // record the new ownership in device_global_indices.
477 - const int iters_token = (num_token_chunks + NUM_WARPS - 1) / NUM_WARPS;
578 + // Phase 3: walk the rows' positions in (row, top-k) order; the first
579 + // reference to a union entry no slot holds is a miss. Assign misses
580 + // eviction slots (oldest first) and record the new ownership in
581 + // device_global_indices. Swaps fill the rows' swap segments in order; the
582 + // union never exceeds their combined capacity, and position chunks never
583 + // exceed buffer chunks because hot_size covers every row's top-k.
584 + const int num_positions = (end_row - first_row) * top_k;
585 + const int num_position_chunks = (num_positions + kWarpSize - 1) / kWarpSize;
586 + const int iters_position = (num_position_chunks + NUM_WARPS - 1) / NUM_WARPS;
587 int miss_running_total = 0;
479 - for (int iter = 0; iter < iters_token; iter++) {
588 + int entry_hits = 0;
589 + int entry_misses = 0;
590 + for (int iter = 0; iter < iters_position; iter++) {
591 const int chunk_idx = warp_id + iter * NUM_WARPS;
481 - const bool has_valid_chunk = chunk_idx < num_token_chunks;
482 -
483 - const int i = chunk_idx * kWarpSize + lane_id;
484 - const bool has_valid_token = has_valid_chunk && (i < top_k);
592 + const bool has_valid_chunk = chunk_idx < num_position_chunks;
593
486 - int32_t g = 0;
594 + const int position = chunk_idx * kWarpSize + lane_id;
595 + int32_t g = -1;
596 + int h = -1;
597 + int32_t entry = 0;
598 bool is_miss = false;
488 - if (has_valid_token) {
489 - is_miss = s_topk[i] != kTokenDone;
490 - if (is_miss) {
491 - g = s_topk[i];
599 + if (has_valid_chunk && position < num_positions) {
600 + const int row = first_row + position / top_k;
601 + int32_t resident_row;
602 + g = translate_topk_entry(
603 + global_indices[static_cast<int64_t>(row) * input_row_stride +
604 + position % top_k],
605 + request_row, source_block_table, resident_block_table, host_rows,
606 + source_bt_stride, source_num_reqs, source_num_blocks,
607 + source_block_size, resident_bt_stride, resident_num_reqs,
608 + resident_num_blocks, resident_block_size, resident_null_block,
609 + resident_row);
610 + if (resident_row < 0 && g >= 0) {
611 + h = hash_slot(g, table_size);
612 + while (s_hash_keys[h] != g) {
613 + h = next_slot(h, table_size);
614 + }
615 + entry = s_hash_vals[h];
616 + is_miss = entry == position;
617 }
618 }
619
620 int local_miss_off = 0;
621 if (has_valid_chunk) {
497 - const unsigned int miss_mask = __ballot_sync(0xFFFFFFFF, is_miss);
498 - local_miss_off = __popc(miss_mask & lanes_before);
622 + const unsigned int miss_mask_bits = __ballot_sync(0xFFFFFFFF, is_miss);
623 + local_miss_off = __popc(miss_mask_bits & lanes_before);
624 if (lane_id == 0) {
500 - s_chunk_off[chunk_idx + 1] = __popc(miss_mask);
625 + s_chunk_off[chunk_idx + 1] = __popc(miss_mask_bits);
626 }
627 }
628 __syncthreads();
@@ -505,60 +630,112 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
630 if (warp_id == 0) {
631 miss_running_total =
632 warp_inclusive_scan(s_chunk_off, lane_id, chunk_idx + 1,
508 - num_token_chunks + 1, miss_running_total);
633 + num_position_chunks + 1, miss_running_total);
634 + if (tid == 0) {
635 + s_counters[3] = miss_running_total;
636 + }
637 }
638 __syncthreads();
639
640 if (is_miss) {
641 const int m = s_chunk_off[chunk_idx] + local_miss_off;
642 const int16_t evict_slot = s_lru_out[hot_size - 1 - m];
643 + const int64_t compact_index = static_cast<int64_t>(first_row) * top_k + m;
644 if (evict_slot < 0 || evict_slot >= region_stride) {
645 // Corruption tripwire: an out-of-range slot (corrupted lru state,
646 // see phase 2) must not become a hot-cache/dgi write. Resolve the
518 - // entry as invalid (-1, masked by attention, not in the miss set);
519 - // it re-misses on a later step. Phase 5 re-checks the same
520 - // s_lru_out value, so its copy is skipped consistently.
521 - store_hot_index(row_out, row_attention, i, -1, hot_block_size,
522 - attention_block_stride);
647 + // entry as invalid (-1, masked by attention); it re-misses on a later
648 + // step, and its copy is skipped.
649 + s_hash_vals[h] = kEntryInvalid;
650 if (swap_host_physical_rows != nullptr) {
524 - const int64_t compact_index =
525 - static_cast<int64_t>(batch_row) * top_k + m;
651 swap_host_physical_rows[compact_index] = g;
652 swap_device_physical_rows[compact_index] = -1;
653 }
654 } else {
530 - // Reuse s_topk as compacted miss scratch: m < i always holds (done
531 - // entries are skipped), so writes never overrun pending reads.
532 - s_topk[m] = g;
533 - const int32_t physical_row = static_cast<int32_t>(
534 - get_physical_hot_row(hot_block_table, request_row, hot_table_stride,
535 - hot_block_size, evict_slot));
536 - store_hot_index(row_out, row_attention, i, physical_row, hot_block_size,
537 - attention_block_stride);
655 + s_hash_vals[h] = kEntryMiss | evict_slot;
656 if (swap_host_physical_rows != nullptr) {
539 - const int64_t compact_index =
540 - static_cast<int64_t>(batch_row) * top_k + m;
657 swap_host_physical_rows[compact_index] = g;
542 - swap_device_physical_rows[compact_index] = physical_row;
658 + swap_device_physical_rows[compact_index] =
659 + static_cast<int32_t>(get_physical_hot_row(
660 + hot_block_table, request_row, hot_table_stride,
661 + hot_block_size, evict_slot));
662 }
544 - if (row_miss != nullptr) row_miss[i] = 1;
663 row_dgi[evict_slot] = g;
664 }
665 + entry = s_hash_vals[h];
666 + }
667 +
668 + // Hits, first references and references to entries resolved in an
669 + // earlier chunk write their hot index now; later references to an entry
670 + // resolved in this same chunk pass (only across rows) are left to phase 4.
671 + const bool resolved =
672 + h >= 0 && (entry & (kEntryHit | kEntryMiss | kEntryInvalid));
673 + if (resolved) {
674 + write_entry_index(
675 + entry, first_row + position / top_k, position % top_k, top_k,
676 + hot_indices, attention_indices, attention_row_stride, miss_mask,
677 + hot_block_table, request_row, hot_table_stride, hot_block_size,
678 + attention_block_stride, is_miss, entry_hits, entry_misses);
679 + }
680 + if (has_valid_chunk) {
681 + const unsigned int done_bits = __ballot_sync(0xFFFFFFFF, resolved);
682 + if (lane_id == 0) {
683 + s_done[chunk_idx] = done_bits;
684 + }
685 }
686 }
687 __syncthreads();
688
689 const int total_hits = s_counters[0];
552 - const int total_misses = top_k - total_hits - s_counters[1];
553 - if (swap_counts != nullptr && tid == 0) {
554 - swap_counts[batch_row] = total_misses;
690 + const int total_misses = s_counters[3];
691 + if (swap_counts != nullptr) {
692 + for (int row = first_row + tid; row < end_row; row += blockDim.x) {
693 + const int remaining = total_misses - (row - first_row) * top_k;
694 + swap_counts[row] = max(0, min(top_k, remaining));
695 + }
696 }
556 - if (stats != nullptr && tid == 0) {
557 - atomicAdd(&stats[0], static_cast<unsigned long long>(total_hits));
558 - atomicAdd(&stats[1], static_cast<unsigned long long>(total_misses));
697 +
698 + // Phase 4: references phase 3 left unresolved (later references across rows
699 + // to entries loaded in the same chunk pass) map onto the union's slots.
700 + // Single-row requests resolve everything in phase 3.
701 + if (end_row - first_row > 1) {
702 + for (int position = tid; position < num_positions; position += blockDim.x) {
703 + if (s_done[position / kWarpSize] & (1u << (position % kWarpSize))) {
704 + continue;
705 + }
706 + const int row = first_row + position / top_k;
707 + int32_t resident_row;
708 + const int32_t g = translate_topk_entry(
709 + global_indices[static_cast<int64_t>(row) * input_row_stride +
710 + position % top_k],
711 + request_row, source_block_table, resident_block_table, host_rows,
712 + source_bt_stride, source_num_reqs, source_num_blocks,
713 + source_block_size, resident_bt_stride, resident_num_reqs,
714 + resident_num_blocks, resident_block_size, resident_null_block,
715 + resident_row);
716 + if (resident_row >= 0 || g < 0) continue;
717 + int h = hash_slot(g, table_size);
718 + while (s_hash_keys[h] != g) {
719 + h = next_slot(h, table_size);
720 + }
721 + write_entry_index(
722 + s_hash_vals[h], row, position % top_k, top_k, hot_indices,
723 + attention_indices, attention_row_stride, miss_mask, hot_block_table,
724 + request_row, hot_table_stride, hot_block_size, attention_block_stride,
725 + false, entry_hits, entry_misses);
726 + }
727 + }
728 + if (stats != nullptr) {
729 + atomicAdd(&s_counters[4], entry_hits);
730 + atomicAdd(&s_counters[2], entry_misses);
731 + __syncthreads();
732 + if (tid == 0) {
733 + atomicAdd(&stats[0], static_cast<unsigned long long>(s_counters[4]));
734 + atomicAdd(&stats[1], static_cast<unsigned long long>(s_counters[2]));
735 + }
736 }
737
561 - // Phase 4: write back the LRU order: stale evictables at the front,
738 + // Phase 5: write back the LRU order: stale evictables at the front,
739 // freshly loaded misses next, then hits at MRU.
740 const int total_evictable = hot_size - total_hits;
741 const int remaining_evictable = total_evictable - total_misses;
@@ -783,7 +960,7 @@ void hisparse_resolve_residency(
960 torch::stable::Tensor& device_global_indices,
961 torch::stable::Tensor& lru_slots,
962 std::optional<torch::stable::Tensor> const& request_state_indices,
786 - int64_t region_stride,
963 + int64_t region_stride, int64_t max_union_rows,
964 std::optional<torch::stable::Tensor> const& miss_mask,
965 std::optional<torch::stable::Tensor> const& stats,
966 std::optional<torch::stable::Tensor> const& attention_indices,
@@ -1034,14 +1211,32 @@ void hisparse_resolve_residency(
1211 }
1212
1213 constexpr int kBlockSize = 1024;
1037 - const int hash_size = 2 * top_k;
1214 + STD_TORCH_CHECK(max_union_rows > 0 && max_union_rows <= hot_size,
1215 + "max_union_rows must be in [1, hot_size], got ",
1216 + max_union_rows);
1217 + // One spare entry above the union bound keeps every probe sequence finite.
1218 + const int hash_size = static_cast<int>(max_union_rows) + 1;
1219 const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize;
1220 const size_t smem_bytes =
1040 - sizeof(int32_t) * (top_k + 2 * (num_buffer_chunks + 1) + hash_size + 3) +
1041 - sizeof(int16_t) * (hot_size + hash_size);
1221 + sizeof(int32_t) * (2 * hash_size + 2 * (num_buffer_chunks + 1) +
1222 + kResidencyCounters + num_buffer_chunks) +
1223 + sizeof(int16_t) * hot_size;
1224
1225 const torch::stable::accelerator::DeviceGuard device_guard(
1226 hot_cache.get_device_index());
1227 + int max_smem_bytes = 0;
1228 + STD_TORCH_CHECK(cudaDeviceGetAttribute(
1229 + &max_smem_bytes, cudaDevAttrMaxSharedMemoryPerBlockOptin,
1230 + hot_cache.get_device_index()) == cudaSuccess,
1231 + "failed to query shared memory per block");
1232 + STD_TORCH_CHECK(smem_bytes <= static_cast<size_t>(max_smem_bytes),
1233 + "HiSparse device_buffer_size=", hot_size, " needs ",
1234 + smem_bytes,
1235 + " bytes of shared memory to resolve residency, "
1236 + "but the device allows ",
1237 + max_smem_bytes,
1238 + "; lower device_buffer_size or the number of speculative "
1239 + "tokens.");
1240 const cudaStream_t stream = get_current_cuda_stream();
1241 const int64_t attention_row_stride =
1242 attention_indices.has_value() ? attention_indices.value().stride(0) : 0;
csrc/libtorch_stable/ops.h
+1 -1
@@ -680,7 +680,7 @@ void hisparse_resolve_residency(
680 torch::stable::Tensor& device_global_indices,
681 torch::stable::Tensor& lru_slots,
682 std::optional<torch::stable::Tensor> const& request_state_indices,
683 - int64_t region_stride,
683 + int64_t region_stride, int64_t max_union_rows,
684 std::optional<torch::stable::Tensor> const& miss_mask,
685 std::optional<torch::stable::Tensor> const& stats,
686 std::optional<torch::stable::Tensor> const& attention_indices,
csrc/libtorch_stable/torch_bindings.cpp
+1
@@ -995,6 +995,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
995 " Tensor! lru_slots,"
996 " Tensor? request_state_indices,"
997 " int region_stride,"
998 + " int max_union_rows,"
999 " Tensor(a!)? miss_mask=None,"
1000 " Tensor(b!)? stats=None,"
1001 " Tensor(c!)? attention_indices=None,"
docs/design/hisparse.md
+4 -4
@@ -130,10 +130,10 @@ path. The resolver consumes the existing graph-stable request mapping from
130 attention metadata; neither the worker nor individual cache handles keep a
131 duplicate mapping.
132
133 -Speculative decoding resolves and consumes each verification step in order.
134 -Each step receives distinct replayable plan rows while sharing the request's
135 -hot-cache state, so a later step cannot reuse a hot row before an earlier step
136 -has consumed it.
133 +Speculative decoding resolves all verification rows of a request in one pass:
134 +one block resolves the union of the rows' top-k against the request's hot-cache
135 +state, so rows that select the same host row share its hot row and no row
136 +evicts a hot row another row of the step still reads.
137
138 ## P/D import target
139
tests/v1/attention/test_sparse_mla_backends.py
+228 -106
@@ -68,10 +68,10 @@ from vllm.utils.torch_utils import current_stream
68 from vllm.v1.attention.backends.mla import index_group as index_group_module
69 from vllm.v1.attention.backends.mla.flashattn_mla_sparse import (
70 FlashAttnMLASparseImpl,
71 + FlashAttnMLASparseMetadataBuilder,
72 )
73 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
74 FlashInferMLASparseImpl,
74 - FlashInferMLASparseMetadataBuilder,
75 FlashInferMLASparseTRTLLMBackend,
76 )
77 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import (
@@ -1880,11 +1880,13 @@ def _make_hisparse_runtime(
1880 block_size: int = 64,
1881 max_swap_rows: int | None = None,
1882 index_group: hisparse_runtime.HiSparseIndexGroup | None = None,
1883 + max_union_rows: int | None = None,
1884 ) -> HiSparseRuntime:
1885 runtime = HiSparseRuntime(
1886 config=ResolvedHiSparseConfig(
1887 top_k=top_k,
1888 device_buffer_size=device_buffer_size,
1889 + max_union_rows=max_union_rows or device_buffer_size,
1890 ),
1891 max_num_reqs=max_num_reqs,
1892 row_width=row_width,
@@ -2015,6 +2017,167 @@ def test_hisparse_maps_speculative_rows_through_request_state():
2017 assert physical.tolist() == [[block_size]] * 4
2018
2019
2020 +@requires_hisparse_ops
2021 +@pytest.mark.parametrize(
2022 + "num_rows,row_stride,top_k",
2023 + [(8, 128, 128), (4, 32, 128), (8, 2048, 2048)],
2024 + ids=["disjoint", "overlap", "long_mtp"],
2025 +)
2026 +def test_hisparse_speculative_rows_resolve_host_misses_consistently(
2027 + num_rows, row_stride, top_k
2028 +):
2029 + """Verification rows of one request must agree on its hot-buffer slots.
2030 +
2031 + Each row resolved residency in its own thread block against the request's
2032 + shared slot ownership and LRU, so rows missing on different host rows could
2033 + claim the same free slot. One row's hot index then pointed at a slot holding
2034 + another row's KV. Overlapping rows must share one slot per host row, and
2035 + long_mtp's rows span several resolution passes.
2036 + """
2037 + device = torch.device(DEVICE_TYPE)
2038 + block_size, row_width = 64, 64
2039 + num_blocks = 2 * num_rows * top_k // block_size
2040 +
2041 + kv_pool = torch.randn(
2042 + (num_blocks + 1, block_size, row_width), dtype=torch.float32
2043 + ).pin_memory()
2044 + flat_pool = kv_pool.reshape(-1, row_width)
2045 + runtime = _make_hisparse_runtime(
2046 + top_k=top_k,
2047 + device_buffer_size=num_rows * top_k,
2048 + max_num_reqs=1,
2049 + row_width=row_width,
2050 + block_size=block_size,
2051 + max_swap_rows=num_rows,
2052 + )
2053 + runtime.bind_source_cache(kv_pool)
2054 + cache = HiSparseCacheHandle(runtime)
2055 + cache.all_context_pages_resident = False
2056 +
2057 + block_table = torch.arange(
2058 + 1, num_blocks + 1, dtype=torch.int32, device=device
2059 + ).view(1, num_blocks)
2060 + req_ids = torch.zeros(num_rows, dtype=torch.int32, device=device)
2061 + rows = torch.arange(num_rows, dtype=torch.int32, device=device)[:, None]
2062 + base = torch.arange(top_k, dtype=torch.int32, device=device)[None, :]
2063 + for step in range(4):
2064 + runtime.begin_forward()
2065 + # Rows start row_stride apart (disjoint at row_stride == top_k) and
2066 + # shift by half a top-k per step: all misses first, then a mix.
2067 + topk = (rows * row_stride + base + step * (top_k // 2)) % (
2068 + num_blocks * block_size
2069 + )
2070 + hot_indices = cache.swap_in(
2071 + req_id_per_token=req_ids,
2072 + block_table=block_table,
2073 + logical_topk_indices=topk.clone(),
2074 + block_size=block_size,
2075 + )
2076 + torch.accelerator.synchronize()
2077 +
2078 + global_ref = _triton_convert_reference_impl(
2079 + req_ids, block_table, topk, block_size, top_k
2080 + )
2081 + gathered = runtime.hot.attention_cache.reshape(-1, row_width)[
2082 + hot_indices.to(torch.long)
2083 + ].cpu()
2084 + torch.testing.assert_close(gathered, flat_pool[global_ref.cpu().to(torch.long)])
2085 +
2086 +
2087 +@requires_hisparse_ops
2088 +def test_hisparse_speculative_rows_count_a_shared_load_once():
2089 + """A host row two verification rows select loads once; the reuse is a hit."""
2090 + device = torch.device(DEVICE_TYPE)
2091 + block_size, row_width, top_k = 64, 64, 128
2092 + num_blocks = top_k // block_size
2093 + kv_pool = torch.randn(
2094 + (num_blocks + 1, block_size, row_width), dtype=torch.float32
2095 + ).pin_memory()
2096 + runtime = _make_hisparse_runtime(
2097 + top_k=top_k,
2098 + device_buffer_size=2 * top_k,
2099 + max_num_reqs=1,
2100 + row_width=row_width,
2101 + block_size=block_size,
2102 + max_swap_rows=2,
2103 + )
2104 + runtime.bind_source_cache(kv_pool)
2105 + cache = HiSparseCacheHandle(runtime)
2106 + cache.all_context_pages_resident = False
2107 + block_table = torch.arange(
2108 + 1, num_blocks + 1, dtype=torch.int32, device=device
2109 + ).view(1, num_blocks)
2110 + topk = torch.arange(top_k, dtype=torch.int32, device=device).repeat(2, 1)
2111 +
2112 + runtime.begin_forward()
2113 + cache.swap_in(
2114 + req_id_per_token=torch.zeros(2, dtype=torch.int32, device=device),
2115 + block_table=block_table,
2116 + logical_topk_indices=topk,
2117 + block_size=block_size,
2118 + )
2119 + torch.accelerator.synchronize()
2120 +
2121 + assert runtime.index_group.swap_stats.cpu().tolist() == [top_k, top_k]
2122 +
2123 +
2124 +@requires_hisparse_ops
2125 +def test_hisparse_union_table_scales_with_union_bound():
2126 + """Shared memory follows the per-request union bound, not the hot buffer.
2127 +
2128 + Sizing the table from device_buffer_size needs ~340 KB here and fails.
2129 + """
2130 + device = torch.device(DEVICE_TYPE)
2131 + block_size, row_width, top_k = 64, 64, 128
2132 + num_blocks = top_k // block_size
2133 + kv_pool = torch.randn(
2134 + (num_blocks + 1, block_size, row_width), dtype=torch.float32
2135 + ).pin_memory()
2136 + runtime = _make_hisparse_runtime(
2137 + top_k=top_k,
2138 + device_buffer_size=32768,
2139 + max_union_rows=top_k,
2140 + max_num_reqs=1,
2141 + row_width=row_width,
2142 + block_size=block_size,
2143 + max_swap_rows=1,
2144 + )
2145 + runtime.bind_source_cache(kv_pool)
2146 + cache = HiSparseCacheHandle(runtime)
2147 + cache.all_context_pages_resident = False
2148 + req_ids = torch.zeros(1, dtype=torch.int32, device=device)
2149 + block_table = torch.arange(
2150 + 1, num_blocks + 1, dtype=torch.int32, device=device
2151 + ).view(1, num_blocks)
2152 + topk = torch.arange(top_k, dtype=torch.int32, device=device).view(1, top_k)
2153 +
2154 + runtime.begin_forward()
2155 + hot_indices = cache.swap_in(
2156 + req_id_per_token=req_ids,
2157 + block_table=block_table,
2158 + logical_topk_indices=topk,
2159 + block_size=block_size,
2160 + )
2161 + torch.accelerator.synchronize()
2162 +
2163 + global_ref = _triton_convert_reference_impl(
2164 + req_ids, block_table, topk, block_size, top_k
2165 + )
2166 + gathered = runtime.hot.attention_cache.reshape(-1, row_width)[
2167 + hot_indices.to(torch.long)
2168 + ].cpu()
2169 + torch.testing.assert_close(
2170 + gathered, kv_pool.reshape(-1, row_width)[global_ref.cpu().to(torch.long)]
2171 + )
2172 +
2173 +
2174 +@requires_hisparse_ops
2175 +def test_hisparse_rejects_union_bound_beyond_shared_memory():
2176 + """An oversized union bound fails at construction, not on the first forward."""
2177 + with pytest.raises(ValueError, match="shared memory"):
2178 + _make_hisparse_runtime(top_k=2048, device_buffer_size=32768)
2179 +
2180 +
2181 @requires_hisparse_ops
2182 def test_hisparse_resident_rows_bypass_hot_lru():
2183 device = torch.device(DEVICE_TYPE)
@@ -3001,8 +3164,8 @@ def test_hisparse_gather_prefill_cache_prefers_resident_rows():
3164 torch.testing.assert_close(staged_flat[i], expected)
3165
3166
3004 -def test_hisparse_fp8_decode_resolves_steps_then_runs_batched_attention(monkeypatch):
3005 - """FP8 decode must use active dimensions when common metadata is padded."""
3167 +def test_hisparse_fp8_decode_resolves_rows_once_then_runs_batched_attention():
3168 + """Verification rows resolve in one swap_in; FP8 attention sees them batched."""
3169 device = torch.device("cpu")
3170 num_decodes = 2
3171 query_len = 3
@@ -3013,48 +3176,31 @@ def test_hisparse_fp8_decode_resolves_steps_then_runs_batched_attention(monkeypa
3176 .view(-1, 1)
3177 .expand(-1, 4)
3178 )
3016 - steps: list[int] = []
3179 kernel_shapes: list[torch.Size] = []
3180
3181 def swap_in(req_ids, *, logical_topk_indices, **kwargs): # noqa: ARG001
3020 - steps.append(int(logical_topk_indices[0, 0]))
3021 - output = kwargs["attention_indices_out"]
3022 - if output is not None:
3023 - output.copy_(logical_topk_indices + 10)
3024 - return topk[:num_decodes]
3182 + return logical_topk_indices + 10
3183
3184 def run_kernel(self, *, q, **kwargs): # noqa: ARG001
3185 kernel_shapes.append(q.shape)
3186 return q[..., :1], None
3187
3030 - runtime = SimpleNamespace(
3031 - hot=SimpleNamespace(attention_cache=torch.empty(1, device=device))
3032 - )
3033 - leader_cache = SimpleNamespace(
3034 - runtime=runtime,
3035 - source_block_table=torch.empty(
3036 - num_decodes, 1, dtype=torch.int32, device=device
3188 + cache = SimpleNamespace(
3189 + runtime=SimpleNamespace(
3190 + hot=SimpleNamespace(attention_cache=torch.empty(1, device=device))
3191 ),
3038 - swap_in=MagicMock(side_effect=swap_in),
3039 - )
3040 - follower_cache = SimpleNamespace(
3041 - runtime=runtime,
3192 source_block_table=torch.empty(
3193 num_decodes, 1, dtype=torch.int32, device=device
3194 ),
3195 swap_in=MagicMock(side_effect=swap_in),
3196 )
3197 index_group = object.__new__(HiSparseMLAIndexGroup)
3048 - index_group.caches = [leader_cache, follower_cache]
3198 + index_group.caches = [cache]
3199 index_group.physical_topk_indices = torch.empty(
3200 (num_tokens + 1, 4), dtype=torch.int32, device=device
3201 )
3052 - index_group.valid_topk_counts = torch.empty(
3053 - num_tokens + 1, dtype=torch.int32, device=device
3054 - )
3055 - index_group.request_ids = torch.arange(
3056 - num_decodes, dtype=torch.int32, device=device
3057 - )
3202 + index_group.row_indices = torch.arange(num_tokens, dtype=torch.int32, device=device)
3203 + index_group.request_ids = torch.empty_like(index_group.row_indices)
3204 impl = SimpleNamespace(
3205 kv_lora_rank=1,
3206 index_group=index_group,
@@ -3062,84 +3208,83 @@ def test_hisparse_fp8_decode_resolves_steps_then_runs_batched_attention(monkeypa
3208 )
3209 impl._fp8_flash_mla_kernel = MethodType(run_kernel, impl)
3210 metadata = SimpleNamespace(
3065 - num_decodes=num_tokens,
3066 - num_decode_tokens=num_tokens,
3067 - decode_max_query_len=1,
3211 query_start_loc=torch.arange(
3069 - 0,
3070 - num_tokens + 1,
3071 - query_len,
3072 - dtype=torch.int32,
3073 - device=device,
3212 + 0, num_tokens + 1, query_len, dtype=torch.int32, device=device
3213 ),
3075 - block_table=torch.empty(num_decodes, 1, dtype=torch.int32, device=device),
3214 + req_id_per_token=torch.arange(
3215 + num_decodes, dtype=torch.int32, device=device
3216 + ).repeat_interleave(query_len),
3217 block_size=64,
3218 )
3219
3220 output = FlashMLASparseImpl._host_backed_fp8_decode(
3080 - impl,
3081 - q,
3082 - topk,
3083 - metadata,
3084 - SimpleNamespace(),
3085 - num_decodes,
3086 - query_len,
3221 + impl, q, topk, metadata, SimpleNamespace(), num_decodes
3222 )
3223
3089 - assert steps == [0, 1, 2]
3090 - assert kernel_shapes == [(num_decodes, query_len, 2, 4)]
3091 - assert output.shape == (num_tokens, 2, 1)
3224 + cache.swap_in.assert_called_once()
3225 torch.testing.assert_close(
3093 - index_group.physical_topk_indices[:num_tokens], topk + 10
3226 + cache.swap_in.call_args.args[0], metadata.req_id_per_token
3227 )
3228 + torch.testing.assert_close(
3229 + cache.swap_in.call_args.kwargs["logical_topk_indices"], topk
3230 + )
3231 + assert kernel_shapes == [(num_decodes, query_len, 2, 4)]
3232 + assert output.shape == (num_tokens, 2, 1)
3233
3096 - follower_indices = index_group.convert_decode_logical_to_physical_topk(
3097 - 1,
3098 - topk,
3099 - metadata,
3100 - return_valid_counts=False,
3101 - num_decodes=num_decodes,
3102 - decode_query_len=query_len,
3234 +
3235 +def test_hisparse_keeps_speculative_rows_as_decodes():
3236 + """HiSparse limits decodes to verification requests, not short extends.
3237 +
3238 + FlashAttn sparse MLA used to classify MTP verification as prefill under
3239 + HiSparse, whose residency-dependent route a FULL CUDA graph freezes at
3240 + capture.
3241 + """
3242 + builder = object.__new__(FlashAttnMLASparseMetadataBuilder)
3243 + builder.vllm_config = SimpleNamespace(
3244 + attention_config=SimpleNamespace(hisparse_config=object()),
3245 + speculative_config=SimpleNamespace(
3246 + num_speculative_tokens=3, parallel_drafting=False
3247 + ),
3248 + parallel_config=SimpleNamespace(decode_context_parallel_size=1),
3249 + use_v2_model_runner=True,
3250 )
3104 - assert steps == [0, 1, 2, 0, 1, 2]
3105 - torch.testing.assert_close(follower_indices, topk + 10)
3251
3252 + builder._init_reorder_batch_threshold(256, supports_spec_as_decode=True)
3253
3108 -def test_hisparse_varlen_decode_uses_canonical_request_rows():
3109 - """Graph metadata's query upper bound must preserve packed token rows."""
3110 - num_decodes = 2
3111 - topk = torch.arange(18, dtype=torch.int32).view(9, 2)
3254 + assert builder.reorder_batch_threshold == 4
3255
3113 - def swap_in(_request_ids, *, logical_topk_indices, **_kwargs):
3114 - return logical_topk_indices + 10
3256
3257 +def test_hisparse_decode_skips_padding_rows():
3258 + """Padding rows are masked, and the ids stay in the group's persistent buffer.
3259 +
3260 + The resolver reads them on the HiSparse copy stream, where a compute-stream
3261 + temporary can already be recycled (CUDA graphs do), scrambling rows' requests.
3262 + """
3263 + topk = torch.arange(20, dtype=torch.int32).view(10, 2)
3264 cache = SimpleNamespace(
3117 - source_block_table=torch.zeros((num_decodes, 1), dtype=torch.int32),
3118 - swap_in=MagicMock(side_effect=swap_in),
3265 + source_block_table=torch.zeros((2, 1), dtype=torch.int32),
3266 + swap_in=MagicMock(return_value=topk),
3267 )
3268 index_group = object.__new__(HiSparseMLAIndexGroup)
3269 index_group.caches = [cache]
3122 - index_group.physical_topk_indices = torch.empty((10, 2), dtype=torch.int32)
3123 - index_group.valid_topk_counts = torch.empty(10, dtype=torch.int32)
3124 - index_group.request_ids = torch.arange(num_decodes, dtype=torch.int32)
3270 + index_group.physical_topk_indices = torch.empty((11, 2), dtype=torch.int32)
3271 + index_group.row_indices = torch.arange(10, dtype=torch.int32)
3272 + index_group.request_ids = torch.empty_like(index_group.row_indices)
3273 metadata = SimpleNamespace(
3126 - num_decodes=num_decodes,
3127 - num_decode_tokens=topk.shape[0],
3128 - decode_max_query_len=15,
3274 query_start_loc=torch.tensor([0, 8, 9], dtype=torch.int32),
3275 + req_id_per_token=torch.tensor([0] * 8 + [1, 0], dtype=torch.int32),
3276 block_size=64,
3277 )
3278
3133 - physical_topk = index_group.convert_decode_logical_to_physical_topk(
3134 - 0,
3135 - topk,
3136 - metadata,
3137 - return_valid_counts=False,
3279 + index_group.convert_logical_to_physical_topk(
3280 + 0, topk, metadata, block_stride_rows=None, return_valid_counts=False
3281 )
3282
3140 - torch.testing.assert_close(physical_topk, topk + 10)
3141 - request_rows = cache.swap_in.call_args_list[0].args[0]
3142 - torch.testing.assert_close(request_rows, index_group.request_ids)
3283 + request_ids = cache.swap_in.call_args.args[0]
3284 + assert request_ids.data_ptr() == index_group.request_ids.data_ptr()
3285 + torch.testing.assert_close(
3286 + request_ids, torch.tensor([0] * 8 + [1, -1], dtype=torch.int32)
3287 + )
3288
3289
3290 def test_flashinfer_hisparse_decode_runs_batched_attention():
@@ -3168,7 +3313,7 @@ def test_flashinfer_hisparse_decode_runs_batched_attention():
3313 )
3314 index_group = object.__new__(HiSparseMLAIndexGroup)
3315 index_group.caches = [cache_handle]
3171 - index_group.convert_decode_logical_to_physical_topk = MethodType(
3316 + index_group.convert_logical_to_physical_topk = MethodType(
3317 convert_decode, index_group
3318 )
3319 impl = object.__new__(FlashInferMLASparseImpl)
@@ -3200,7 +3345,7 @@ def test_flashattn_hisparse_decode_uses_index_group():
3345 physical = topk.clone()
3346 counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE)
3347 index_group = object.__new__(HiSparseMLAIndexGroup)
3203 - index_group.convert_decode_logical_to_physical_topk = MagicMock(
3348 + index_group.convert_logical_to_physical_topk = MagicMock(
3349 return_value=(physical, counts)
3350 )
3351 index_group.physical_kv_cache = MagicMock(
@@ -3223,7 +3368,7 @@ def test_flashattn_hisparse_decode_uses_index_group():
3368
3369 assert output.shape == (num_tokens, 2, 1)
3370 assert lse is None
3226 - index_group.convert_decode_logical_to_physical_topk.assert_called_once()
3371 + index_group.convert_logical_to_physical_topk.assert_called_once()
3372 impl._run_mqa_kernel.assert_called_once()
3373
3374
@@ -3233,9 +3378,7 @@ def test_flashinfer_sm120_hisparse_decode_uses_index_group():
3378 topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=DEVICE_TYPE)
3379 physical = topk.clone()
3380 index_group = object.__new__(HiSparseMLAIndexGroup)
3236 - index_group.convert_decode_logical_to_physical_topk = MagicMock(
3237 - return_value=physical
3238 - )
3381 + index_group.convert_logical_to_physical_topk = MagicMock(return_value=physical)
3382 index_group.physical_kv_cache = MagicMock(
3383 return_value=torch.empty(1, device=DEVICE_TYPE)
3384 )
@@ -3259,7 +3402,7 @@ def test_flashinfer_sm120_hisparse_decode_uses_index_group():
3402
3403 assert output.shape == (num_tokens, 2, 1)
3404 assert lse is None
3262 - index_group.convert_decode_logical_to_physical_topk.assert_called_once()
3405 + index_group.convert_logical_to_physical_topk.assert_called_once()
3406 impl._run_mqa_kernel.assert_called_once()
3407
3408
@@ -3291,27 +3434,6 @@ def test_hisparse_resident_prefill_uses_attention_block_stride():
3434 assert index_group._convert_once.call_args.kwargs["block_stride_rows"] == 832
3435
3436
3294 -def test_hisparse_shared_sparse_builder_routes_multi_token_chunks_to_prefill():
3295 - """Default sparse backends must route multi-token chunks to prefill."""
3296 - builder = object.__new__(FlashInferMLASparseMetadataBuilder)
3297 - builder.vllm_config = SimpleNamespace(
3298 - attention_config=SimpleNamespace(hisparse_config=object()),
3299 - speculative_config=SimpleNamespace(
3300 - num_speculative_tokens=4,
3301 - parallel_drafting=False,
3302 - ),
3303 - parallel_config=SimpleNamespace(decode_context_parallel_size=1),
3304 - )
3305 -
3306 - builder._init_reorder_batch_threshold(
3307 - 1024,
3308 - supports_spec_as_decode=True,
3309 - supports_dcp_with_varlen=True,
3310 - )
3311 -
3312 - assert builder.reorder_batch_threshold == 1
3313 -
3314 -
3437 def test_hisparse_flashmla_reorders_full_speculative_window_as_decode():
3438 builder = object.__new__(FlashMLASparseMetadataBuilder)
3439 builder.vllm_config = SimpleNamespace(
tests/v1/worker/test_utils.py
+1
@@ -1156,6 +1156,7 @@ def test_hisparse_cache_handles_join_index_groups_during_construction(monkeypatc
1156 resolved = hisparse_runtime_module.ResolvedHiSparseConfig(
1157 top_k=4,
1158 device_buffer_size=8,
1159 + max_union_rows=8,
1160 )
1161 monkeypatch.setattr(hisparse_runtime_module, "_has_hisparse_ops", lambda: True)
1162 monkeypatch.setattr(
vllm/model_executor/layers/attention/sparse_mla_attention.py
-3
@@ -165,7 +165,6 @@ T = TypeVar("T", bound=SparseMLACommonMetadata)
165 class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
166 metadata_cls: type[T]
167 require_uniform_decodes: ClassVar[bool] = False
168 - hisparse_supports_multi_token_decode: ClassVar[bool] = False
168
169 def __init__(
170 self,
@@ -259,8 +258,6 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
258 ) -> None:
259 if self.vllm_config.attention_config.hisparse_config is not None:
260 reorder_batch_threshold = 1
262 - if not self.hisparse_supports_multi_token_decode:
263 - supports_spec_as_decode = False
261 super()._init_reorder_batch_threshold(
262 reorder_batch_threshold,
263 supports_spec_as_decode,
vllm/v1/attention/backends/mla/flashattn_mla_sparse.py
+2 -1
@@ -207,10 +207,11 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
207 outputs = []
208 if num_decode_tokens:
209 physical_topk, valid_counts = (
210 - index_group.convert_decode_logical_to_physical_topk(
210 + index_group.convert_logical_to_physical_topk(
211 self.index_group_index,
212 topk_indices[:num_decode_tokens],
213 attn_metadata,
214 + block_stride_rows=None,
215 return_valid_counts=True,
216 )
217 )
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+2 -2
@@ -273,7 +273,6 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde
273 """Metadata builder for the SM100 TRT-LLM sparse MLA kernel."""
274
275 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
276 - hisparse_supports_multi_token_decode: ClassVar[bool] = True
276
277 def __init__(
278 self,
@@ -502,10 +501,11 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
501 decode_lse: torch.Tensor | None = None
502 if num_decode_tokens > 0:
503 physical_topk, valid_counts = (
505 - index_group.convert_decode_logical_to_physical_topk(
504 + index_group.convert_logical_to_physical_topk(
505 self.index_group_index,
506 topk_indices[:num_decode_tokens],
507 attn_metadata,
508 + block_stride_rows=None,
509 return_valid_counts=True,
510 )
511 )
vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py
+2 -1
@@ -125,10 +125,11 @@ class FlashInferMLASparseSM120Impl(SparseMLACommonImpl[FlashInferMLASparseMetada
125 if num_decode_tokens:
126 topk_indices_physical = cast(
127 torch.Tensor,
128 - index_group.convert_decode_logical_to_physical_topk(
128 + index_group.convert_logical_to_physical_topk(
129 self.index_group_index,
130 topk_indices[:num_decode_tokens],
131 attn_metadata,
132 + block_stride_rows=None,
133 return_valid_counts=False,
134 ),
135 )
vllm/v1/attention/backends/mla/flashmla_sparse.py
+4 -7
@@ -317,7 +317,6 @@ class FlashMLASparseMetadataBuilder(
317 ):
318 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
319 require_uniform_decodes: ClassVar[bool] = True
320 - hisparse_supports_multi_token_decode: ClassVar[bool] = True
320 metadata_cls = FlashMLASparseMetadata
321
322 def __init__(
@@ -818,10 +817,11 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
817 num_decode_tokens = attn_metadata.num_decode_tokens
818 if num_decode_tokens > 0:
819 decode_topk, decode_lengths = (
821 - index_group.convert_decode_logical_to_physical_topk(
820 + index_group.convert_logical_to_physical_topk(
821 self.index_group_index,
822 topk_indices[:num_decode_tokens],
823 attn_metadata,
824 + block_stride_rows=None,
825 return_valid_counts=True,
826 )
827 )
@@ -1027,7 +1027,6 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1027 attn_metadata,
1028 fp8_metadata.decode.kernel_metadata,
1029 num_decodes,
1030 - fp8_metadata.decode.decode_query_len,
1030 )
1031 # Reshape q: (num_decode_tokens, num_heads, head_dim)
1032 # -> (num_decodes, seq_len, num_heads, head_dim)
@@ -1267,16 +1266,14 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1266 attn_metadata: FlashMLASparseMetadata,
1267 kernel_metadata: FlashMLASparseMetadata.FP8KernelMetadata,
1268 num_decodes: int,
1270 - decode_query_len: int,
1269 ) -> torch.Tensor:
1270 assert isinstance(self.index_group, HiSparseMLAIndexGroup)
1273 - physical_topk = self.index_group.convert_decode_logical_to_physical_topk(
1271 + physical_topk = self.index_group.convert_logical_to_physical_topk(
1272 self.index_group_index,
1273 topk_indices,
1274 attn_metadata,
1275 + block_stride_rows=None,
1276 return_valid_counts=False,
1278 - num_decodes=num_decodes,
1279 - decode_query_len=decode_query_len,
1277 )
1278 assert isinstance(physical_topk, torch.Tensor)
1279 q = reshape_query_for_spec_decode(q, num_decodes)
vllm/v1/attention/backends/mla/index_group.py
+16 -104
@@ -39,6 +39,7 @@ class SparseMLAIndexGroup:
39 logical_topk_indices: torch.Tensor
40 physical_topk_indices: torch.Tensor
41 valid_topk_counts: torch.Tensor
42 + row_indices: torch.Tensor
43 request_ids: torch.Tensor
44 side_stream: torch.Stream
45 logical_topk_ready: torch.Event
@@ -230,12 +231,10 @@ class HiSparseMLAIndexGroup(SparseMLAIndexGroup):
231 *,
232 block_stride_rows: int | None,
233 return_valid_counts: bool,
233 - req_id_per_token: torch.Tensor | None = None,
234 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
235 cache = self.cache(layer_index)
236 num_tokens = logical_topk_indices.shape[0]
237 - if req_id_per_token is None:
238 - req_id_per_token = attn_metadata.req_id_per_token[:num_tokens]
237 + req_id_per_token = attn_metadata.req_id_per_token[:num_tokens]
238 if num_tokens > self.physical_topk_indices.shape[0]:
239 # Prefill-sized batches do not fit the decode residency workspace.
240 # Non-resident prefills are staged before reaching this path.
@@ -253,113 +252,21 @@ class HiSparseMLAIndexGroup(SparseMLAIndexGroup):
252 )
253 source_block_table = cache.source_block_table
254 assert source_block_table is not None
255 + # CUDA-graph padding rows past the batch's tokens map to request 0;
256 + # mark them -1 so residency resolution skips them.
257 + request_ids = self.request_ids[:num_tokens]
258 + request_ids.copy_(req_id_per_token)
259 + request_ids.masked_fill_(
260 + self.row_indices[:num_tokens] >= attn_metadata.query_start_loc[-1], -1
261 + )
262 return cache.swap_in(
257 - req_id_per_token,
263 + request_ids,
264 block_table=source_block_table,
265 logical_topk_indices=logical_topk_indices,
266 block_size=attn_metadata.block_size,
267 return_valid_counts=return_valid_counts,
268 )
269
264 - def convert_decode_logical_to_physical_topk(
265 - self,
266 - layer_index: int,
267 - logical_topk_indices: torch.Tensor,
268 - attn_metadata: Any,
269 - *,
270 - return_valid_counts: bool,
271 - num_decodes: int | None = None,
272 - decode_query_len: int | None = None,
273 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
274 - num_tokens = logical_topk_indices.shape[0]
275 - if num_decodes is None:
276 - num_decodes = attn_metadata.num_decodes
277 - if decode_query_len is None:
278 - decode_query_len = attn_metadata.decode_max_query_len
279 - if decode_query_len == 1:
280 - return self.convert_logical_to_physical_topk(
281 - layer_index,
282 - logical_topk_indices,
283 - attn_metadata,
284 - block_stride_rows=None,
285 - return_valid_counts=return_valid_counts,
286 - req_id_per_token=self.request_ids[:num_decodes],
287 - )
288 -
289 - if num_tokens != num_decodes * decode_query_len:
290 - query_start_loc = attn_metadata.query_start_loc[: num_decodes + 1]
291 - physical_topk_indices = self.physical_topk_indices[: num_tokens + 1]
292 - valid_topk_counts = self.valid_topk_counts[: num_tokens + 1]
293 - request_ids = self.request_ids[:num_decodes]
294 - cache = self.cache(layer_index)
295 - source_block_table = cache.source_block_table
296 - assert source_block_table is not None
297 - for step in range(decode_query_len):
298 - token_indices = query_start_loc[:-1] + step
299 - active = (token_indices < query_start_loc[1:]) & (
300 - token_indices < num_tokens
301 - )
302 - output_indices = torch.where(
303 - active, token_indices, torch.full_like(token_indices, num_tokens)
304 - ).long()
305 - token_indices = token_indices.clamp(0, num_tokens - 1).long()
306 - step_topk = logical_topk_indices.index_select(
307 - 0, token_indices
308 - ).masked_fill(~active.unsqueeze(1), -1)
309 - step_result = cache.swap_in(
310 - request_ids,
311 - block_table=source_block_table,
312 - logical_topk_indices=step_topk,
313 - block_size=attn_metadata.block_size,
314 - return_valid_counts=return_valid_counts,
315 - )
316 - if layer_index == 0:
317 - if return_valid_counts:
318 - step_indices, step_counts = step_result
319 - valid_topk_counts.index_copy_(0, output_indices, step_counts)
320 - else:
321 - step_indices = step_result
322 - physical_topk_indices.index_copy_(0, output_indices, step_indices)
323 - physical_topk_indices = physical_topk_indices[:num_tokens]
324 - if return_valid_counts:
325 - return physical_topk_indices, valid_topk_counts[:num_tokens]
326 - return physical_topk_indices
327 -
328 - assert num_tokens == num_decodes * decode_query_len
329 - logical_topk_by_request = logical_topk_indices.view(
330 - num_decodes, decode_query_len, -1
331 - )
332 - physical_topk_by_request = self.physical_topk_indices[:num_tokens].view(
333 - num_decodes, decode_query_len, -1
334 - )
335 - valid_topk_by_request = self.valid_topk_counts[:num_tokens].view(
336 - num_decodes, decode_query_len
337 - )
338 - request_ids = self.request_ids[:num_decodes]
339 - for step in range(decode_query_len):
340 - cache = self.cache(layer_index)
341 - source_block_table = cache.source_block_table
342 - assert source_block_table is not None
343 - cache.swap_in(
344 - request_ids,
345 - block_table=source_block_table,
346 - logical_topk_indices=logical_topk_by_request[:, step],
347 - block_size=attn_metadata.block_size,
348 - return_valid_counts=return_valid_counts,
349 - attention_indices_out=(
350 - physical_topk_by_request[:, step] if layer_index == 0 else None
351 - ),
352 - valid_counts_out=(
353 - valid_topk_by_request[:, step]
354 - if layer_index == 0 and return_valid_counts
355 - else None
356 - ),
357 - )
358 - physical_topk_indices = physical_topk_by_request.view(num_tokens, -1)
359 - if return_valid_counts:
360 - return physical_topk_indices, valid_topk_by_request.view(num_tokens)
361 - return physical_topk_indices
362 -
270 def stage_prefill_rows(
271 self, layer_index: int, kv_cache: torch.Tensor, attn_metadata: Any
272 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -477,7 +384,12 @@ class SparseMLAIndexGroupBuilder:
384 dtype=torch.int32,
385 device=self.logical_topk_indices.device,
386 ),
480 - request_ids=torch.arange(
387 + row_indices=torch.arange(
388 + workspace_rows,
389 + dtype=torch.int32,
390 + device=self.logical_topk_indices.device,
391 + ),
392 + request_ids=torch.empty(
393 workspace_rows,
394 dtype=torch.int32,
395 device=self.logical_topk_indices.device,
vllm/v1/hisparse/runtime.py
+37 -1
@@ -19,7 +19,8 @@ from vllm.distributed import get_tp_group
19 from vllm.forward_context import get_forward_context, in_piecewise_cudagraph
20 from vllm.logger import init_logger
21 from vllm.platforms import current_platform
22 -from vllm.utils.math_utils import round_up
22 +from vllm.utils.math_utils import cdiv, round_up
23 +from vllm.utils.mem_utils import get_max_shared_memory_bytes
24 from vllm.utils.torch_utils import current_stream
25 from vllm.v1.attention.backend import max_decode_query_len
26 from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
@@ -96,6 +97,9 @@ class PagedCacheView:
97 class ResolvedHiSparseConfig:
98 top_k: int
99 device_buffer_size: int
100 + # Most host rows one request's decode rows can select in a step: every
101 + # verification row's top-k.
102 + max_union_rows: int
103 eager_host_mirror: bool = True
104
105 @classmethod
@@ -151,10 +155,38 @@ class ResolvedHiSparseConfig:
155 return cls(
156 top_k=model_top_k,
157 device_buffer_size=device_buffer_size,
158 + max_union_rows=min_device_buffer_size,
159 eager_host_mirror=config.eager_host_mirror,
160 )
161
162
163 +def _check_residency_shared_memory(
164 + device: torch.device, config: ResolvedHiSparseConfig
165 +) -> None:
166 + """Reject a config whose residency resolver exceeds per-block shared memory."""
167 + # Mirrors hisparse_resolve_residency's layout: union hash keys and values
168 + # (plus an empty slot), per-chunk offsets, five counters, done bits, and
169 + # the int16 compacted LRU.
170 + hot_size = config.device_buffer_size
171 + num_chunks = cdiv(hot_size, 32)
172 + required = 4 * (2 * (config.max_union_rows + 1) + 3 * num_chunks + 7) + 2 * hot_size
173 + device_index = (
174 + device.index
175 + if device.index is not None
176 + else torch.accelerator.current_device_index()
177 + )
178 + available = get_max_shared_memory_bytes(device_index)
179 + if required > available:
180 + raise ValueError(
181 + "HiSparse residency resolution needs "
182 + f"{required} B of shared memory per block for "
183 + f"device_buffer_size={config.device_buffer_size} and "
184 + f"{config.max_union_rows} rows per request "
185 + "(decode query length x index_topk), but the device allows "
186 + f"{available} B. Lower device_buffer_size or num_speculative_tokens."
187 + )
188 +
189 +
190 def check_hisparse_host_memory(pool_bytes: int) -> None:
191 """Reject a physical host-pool allocation that cannot fit in RAM."""
192 mem = psutil.virtual_memory()
@@ -677,6 +709,9 @@ class HiSparseRuntime:
709 # Logical slots per request. Physical rows come from its ephemeral HMA
710 # block table and need not be contiguous in the shared slab.
711 self.region_stride = config.device_buffer_size
712 + self.max_union_rows = config.max_union_rows
713 + if self.device.type == "cuda":
714 + _check_residency_shared_memory(self.device, config)
715
716 row_bytes = row_width * kv_dtype.itemsize
717 if row_bytes % 16 != 0:
@@ -919,6 +954,7 @@ class HiSparseRuntime:
954 group.lru_slots,
955 request_state_indices,
956 self.region_stride,
957 + self.max_union_rows,
958 None,
959 group.swap_stats,
960 physical_topk_indices,