iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Matthew Bonanni
2026-09-30 03:01:45 +00:00
committed by GitHub
co-authored by Claude Opus 5.5
parent bec710fccf
commit 2798f66860
13 changed files with 705 additions and 442 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/hisparse_kernels.cu
+406 -211
View File
@@ -30,9 +30,7 @@
namespace {
constexpr int kWarpSize = 32;
// Sentinel in the shared top-k scratch: entry already resolved (hit /
// newest / invalid), no miss handling needed.
constexpr int32_t kTokenDone = -1;
// Empty key in the shared open-addressing table of host rows.
constexpr int32_t kHashEmpty = -1;
bool is_pinned_cpu_tensor(const torch::stable::Tensor& tensor) {
@@ -46,10 +44,15 @@ bool is_pinned_cpu_tensor(const torch::stable::Tensor& tensor) {
return attributes.type == cudaMemoryTypeHost;
}
__device__ __forceinline__ int32_t hash_slot(int32_t key, int32_t hash_size) {
// Knuth multiplicative hash for the open-addressing table.
return static_cast<int32_t>((static_cast<uint32_t>(key) * 2654435761u) %
static_cast<uint32_t>(hash_size));
__device__ __forceinline__ int32_t hash_slot(int32_t key, int size) {
// Knuth multiplicative hash scaled onto [0, size) by a multiply-high, so
// the table can take any size without a division.
return static_cast<int32_t>(
__umulhi(static_cast<uint32_t>(key) * 2654435761u, size));
}
__device__ __forceinline__ int next_slot(int h, int size) {
return h + 1 == size ? 0 : h + 1;
}
// 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(
}
}
// One block per batch row.
// A union entry's value holds its first (row, top-k) position until phase 2
// or 3 resolves it to a hot slot tagged with how it was resolved.
constexpr int32_t kEntryHit = 1 << 30;
constexpr int32_t kEntryMiss = 1 << 29;
constexpr int32_t kEntryInvalid = 1 << 28;
constexpr int32_t kEntrySlotMask = 0xFFFF;
constexpr int kResidencyCounters = 5;
// Per-request table entries per referenced top-k position (load <= 1/4
// until the hash_size cap).
constexpr int kHashLoadInverse = 4;
// Translate a request-relative top-k position to its resident GPU row (-1 when
// the page is not resident) and returns its host row (-1 when not host-backed).
__device__ __forceinline__ int32_t translate_topk_entry(
const int32_t token_index, const int32_t request_row,
const int32_t* __restrict__ source_block_table,
const int32_t* __restrict__ resident_block_table, const int64_t host_rows,
const int64_t source_bt_stride, const int32_t source_num_reqs,
const int32_t source_num_blocks, const int32_t source_block_size,
const int64_t resident_bt_stride, const int32_t resident_num_reqs,
const int32_t resident_num_blocks, const int32_t resident_block_size,
const int32_t resident_null_block, int32_t& resident_row) {
int32_t g = token_index;
resident_row = -1;
if (source_block_table != nullptr) {
const int32_t source_block =
token_index >= 0 ? token_index / source_block_size : -1;
if (request_row >= 0 && request_row < source_num_reqs &&
source_block >= 0 && source_block < source_num_blocks) {
const int32_t physical_block =
source_block_table[static_cast<int64_t>(request_row) *
source_bt_stride +
source_block];
g = physical_block > 0 ? physical_block * source_block_size +
token_index % source_block_size
: -1;
} else {
g = -1;
}
if (resident_block_table != nullptr) {
const int32_t resident_block =
token_index >= 0 ? token_index / resident_block_size : -1;
if (request_row >= 0 && request_row < resident_num_reqs &&
resident_block >= 0 && resident_block < resident_num_blocks) {
const int32_t physical_block =
resident_block_table[static_cast<int64_t>(request_row) *
resident_bt_stride +
resident_block];
if (physical_block != resident_null_block && physical_block >= 0) {
resident_row = physical_block * resident_block_size +
token_index % resident_block_size;
}
}
}
}
return g >= host_rows ? -1 : g;
}
// Write the hot index of one reference to a resolved union entry and tally it.
__device__ __forceinline__ void write_entry_index(
const int32_t entry, const int row, const int i, const int32_t top_k,
int32_t* __restrict__ hot_indices, int32_t* __restrict__ attention_indices,
const int64_t attention_row_stride, int32_t* __restrict__ miss_mask,
const int32_t* __restrict__ hot_block_table, const int request_row,
const int64_t hot_table_stride, const int32_t hot_block_size,
const int64_t attention_block_stride, const bool first_reference, int& hits,
int& misses) {
int32_t* row_out = hot_indices + static_cast<int64_t>(row) * top_k;
int32_t* row_attention =
attention_indices != nullptr
? attention_indices + static_cast<int64_t>(row) * attention_row_stride
: nullptr;
if (entry & kEntryInvalid) {
store_hot_index(row_out, row_attention, i, -1, hot_block_size,
attention_block_stride);
return;
}
store_hot_index(row_out, row_attention, i,
static_cast<int32_t>(get_physical_hot_row(
hot_block_table, request_row, hot_table_stride,
hot_block_size, entry & kEntrySlotMask)),
hot_block_size, attention_block_stride);
// Only an entry's first reference loads it; later references reuse the
// slot like a hit.
if ((entry & kEntryMiss) && first_reference) {
if (miss_mask != nullptr) {
miss_mask[static_cast<int64_t>(row) * top_k + i] = 1;
}
++misses;
} else {
++hits;
}
}
// One block per request, launched per batch row. A request's rows (several
// with speculative decoding) are contiguous and share its hot-buffer state, so
// the block of its first row resolves the union of their host misses once and
// maps every row onto it; the other blocks exit. The union is bounded by the
// caller's max_union_rows (every decode query's top-k), which
// device_buffer_size covers; a larger union traps. Misses take eviction slots
// in (row, top-k) order of first reference, so the plan is deterministic.
//
// Shared memory layout (int32 region followed by int16 region):
// s_topk[top_k] top-k global ids; reused as miss scratch
// s_chunk_off[nbc + 1] prefix sums for hit (then miss) compaction
// s_hash_keys[hash_size] open addressing: union of the rows' host ids
// s_hash_vals[hash_size] first position, then slot | kEntry* tag
// s_chunk_off[nbc + 1] prefix sums for hit, then miss, compaction
// s_evict_off[nbc + 1] prefix sums for evictable compaction
// s_hash_keys[hash_size] open addressing: global id -> top-k index
// s_counters[3] hits, phase-1 resolved, valid count
// s_counters[5] distinct hits, union size, per-row valid count,
// misses, per-entry hits
// s_done[nbc] positions phase 3 resolved, one word per chunk
// s_lru_out[hot_size] int16, compacted slots: [hits fwd | evict bwd]
// s_hash_vals[hash_size] int16 hash values (top-k index)
// Valid global ids must be unique within each row.
__global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
const int32_t* __restrict__ hot_block_table, // [max_rows, hot_blocks]
@@ -218,11 +322,25 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
const int64_t attention_row_stride, const int64_t valid_count_stride) {
const int NUM_WARPS = blockDim.x / kWarpSize;
const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize;
const int num_token_chunks = (top_k + kWarpSize - 1) / kWarpSize;
const int batch_row = blockIdx.x;
// The first row of each contiguous run of a request's rows leads it. FULL
// CUDA-graph padding rows map to request 0 after the real rows, forming a
// second run, but carry no valid top-k: that run returns before phase 2
// touches the request's shared state.
const int first_row = blockIdx.x;
const int request_row =
request_ids != nullptr ? request_ids[batch_row] : batch_row;
request_ids != nullptr ? request_ids[first_row] : first_row;
if (request_ids != nullptr && first_row > 0 &&
request_ids[first_row - 1] == request_row) {
return;
}
int end_row = first_row + 1;
if (request_ids != nullptr) {
while (end_row < static_cast<int>(gridDim.x) &&
request_ids[end_row] == request_row) {
++end_row;
}
}
const int state_row =
request_state_indices != nullptr && request_row >= 0 &&
request_row < request_state_count
@@ -230,23 +348,24 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
: (request_state_indices == nullptr ? request_row : -1);
// V2 publishes -1 for CUDA-graph padding rows.
if (state_row < 0) {
for (int i = threadIdx.x; i < top_k; i += blockDim.x) {
const int64_t index = static_cast<int64_t>(batch_row) * top_k + i;
hot_indices[index] = -1;
if (attention_indices != nullptr) {
attention_indices[static_cast<int64_t>(batch_row) *
attention_row_stride +
i] = -1;
for (int row = first_row; row < end_row; ++row) {
for (int i = threadIdx.x; i < top_k; i += blockDim.x) {
const int64_t index = static_cast<int64_t>(row) * top_k + i;
hot_indices[index] = -1;
if (attention_indices != nullptr) {
attention_indices[static_cast<int64_t>(row) * attention_row_stride +
i] = -1;
}
if (resolved_global_indices != nullptr) {
resolved_global_indices[index] = -1;
}
}
if (resolved_global_indices != nullptr) {
resolved_global_indices[index] = -1;
if (valid_counts != nullptr && threadIdx.x == 0) {
valid_counts[static_cast<int64_t>(row) * valid_count_stride] = 0;
}
if (swap_counts != nullptr && threadIdx.x == 0) {
swap_counts[row] = 0;
}
}
if (valid_counts != nullptr && threadIdx.x == 0) {
valid_counts[static_cast<int64_t>(batch_row) * valid_count_stride] = 0;
}
if (swap_counts != nullptr && threadIdx.x == 0) {
swap_counts[batch_row] = 0;
}
return;
}
@@ -255,136 +374,130 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
const int lane_id = tid % kWarpSize;
const unsigned int lanes_before = ((unsigned int)1 << lane_id) - 1;
const int32_t* row_topk =
global_indices + static_cast<int64_t>(batch_row) * input_row_stride;
int32_t* row_out = hot_indices + static_cast<int64_t>(batch_row) * top_k;
int32_t* row_attention =
attention_indices != nullptr
? attention_indices +
static_cast<int64_t>(batch_row) * attention_row_stride
: nullptr;
int32_t* row_miss = (miss_mask != nullptr)
? miss_mask + static_cast<int64_t>(batch_row) * top_k
: nullptr;
int32_t* row_dgi =
device_global_indices + static_cast<int64_t>(state_row) * region_stride;
int16_t* row_lru = lru_slots + static_cast<int64_t>(state_row) * hot_size;
// A union within max_union_rows (hash_size - 1) and within the rows' top-k
// leaves an empty slot, so every probe sequence ends. Use a slice sized to
// this request's rows so single-row requests initialize a small table.
const int table_size =
min(hash_size, kHashLoadInverse * (end_row - first_row) * top_k);
extern __shared__ char smem_raw[];
int32_t* s_topk = reinterpret_cast<int32_t*>(smem_raw);
int32_t* s_chunk_off = s_topk + top_k;
int32_t* s_hash_keys = reinterpret_cast<int32_t*>(smem_raw);
int32_t* s_hash_vals = s_hash_keys + hash_size;
int32_t* s_chunk_off = s_hash_vals + hash_size;
int32_t* s_evict_off = s_chunk_off + (num_buffer_chunks + 1);
int32_t* s_hash_keys = s_evict_off + (num_buffer_chunks + 1);
int32_t* s_counters = s_hash_keys + hash_size;
int16_t* s_lru_out = reinterpret_cast<int16_t*>(s_counters + 3);
int16_t* s_hash_vals = s_lru_out + hot_size;
int32_t* s_counters = s_evict_off + (num_buffer_chunks + 1);
unsigned int* s_done =
reinterpret_cast<unsigned int*>(s_counters + kResidencyCounters);
int16_t* s_lru_out = reinterpret_cast<int16_t*>(s_done + num_buffer_chunks);
if (tid < 3) {
s_counters[tid] = 0;
}
__syncthreads();
// Phase 1: translate request-relative positions and resolve resident rows.
for (int i = tid; i < top_k; i += blockDim.x) {
const int32_t token_index = row_topk[i];
int32_t g = token_index;
int32_t resident_row = -1;
if (source_block_table != nullptr) {
const int32_t request_id = request_ids[batch_row];
const int32_t source_block =
token_index >= 0 ? token_index / source_block_size : -1;
if (request_id >= 0 && request_id < source_num_reqs &&
source_block >= 0 && source_block < source_num_blocks) {
const int32_t physical_block =
source_block_table[static_cast<int64_t>(request_id) *
source_bt_stride +
source_block];
g = physical_block > 0 ? physical_block * source_block_size +
token_index % source_block_size
: -1;
} else {
g = -1;
}
if (resident_block_table != nullptr) {
const int32_t resident_block =
token_index >= 0 ? token_index / resident_block_size : -1;
if (request_id >= 0 && request_id < resident_num_reqs &&
resident_block >= 0 && resident_block < resident_num_blocks) {
const int32_t physical_block =
resident_block_table[static_cast<int64_t>(request_id) *
resident_bt_stride +
resident_block];
if (physical_block != resident_null_block && physical_block >= 0) {
resident_row = physical_block * resident_block_size +
token_index % resident_block_size;
}
}
}
}
if (g >= host_rows) {
g = -1;
}
if (resolved_global_indices != nullptr) {
resolved_global_indices[static_cast<int64_t>(batch_row) * top_k + i] = g;
}
if (resident_row >= 0 || g >= 0) atomicAdd(&s_counters[2], 1);
if (row_miss != nullptr) row_miss[i] = 0;
if (resident_row >= 0) {
store_hot_index(row_out, row_attention, i, resident_row, hot_block_size,
attention_block_stride);
s_topk[i] = kTokenDone;
atomicAdd(&s_counters[1], 1);
} else if (g < 0) {
store_hot_index(row_out, row_attention, i, -1, hot_block_size,
attention_block_stride);
s_topk[i] = kTokenDone;
atomicAdd(&s_counters[1], 1);
} else {
s_topk[i] = g;
}
}
__syncthreads();
if (valid_counts != nullptr && tid == 0) {
valid_counts[static_cast<int64_t>(batch_row) * valid_count_stride] =
s_counters[2];
}
// Fully resident rows need only request-relative page translation. Avoid
// scanning or rewriting the hot LRU when no selected row can consult it.
if (s_counters[1] == top_k) {
if (tid == 0 && swap_counts != nullptr) {
swap_counts[batch_row] = 0;
}
return;
}
for (int i = tid; i < hash_size; i += blockDim.x) {
for (int i = tid; i < table_size; i += blockDim.x) {
s_hash_keys[i] = kHashEmpty;
s_hash_vals[i] = INT32_MAX;
}
for (int i = tid; i < num_buffer_chunks + 1; i += blockDim.x) {
s_chunk_off[i] = 0;
s_evict_off[i] = 0;
}
__syncthreads();
for (int i = tid; i < top_k; i += blockDim.x) {
const int32_t g = s_topk[i];
if (g == kTokenDone) {
continue;
}
int slot = hash_slot(g, hash_size);
while (true) {
const int32_t old = atomicCAS(&s_hash_keys[slot], kHashEmpty, g);
if (old == kHashEmpty || old == g) {
s_hash_vals[slot] = static_cast<int16_t>(i);
break;
}
slot = (slot + 1) % hash_size;
}
if (tid < kResidencyCounters) {
s_counters[tid] = 0;
}
__syncthreads();
// Phase 1: translate every row's request-relative positions, resolve
// resident rows directly, and collect the union of host rows.
for (int row = first_row; row < end_row; ++row) {
const int32_t* row_topk =
global_indices + static_cast<int64_t>(row) * input_row_stride;
int32_t* row_out = hot_indices + static_cast<int64_t>(row) * top_k;
int32_t* row_attention =
attention_indices != nullptr
? attention_indices +
static_cast<int64_t>(row) * attention_row_stride
: nullptr;
int32_t* row_miss = miss_mask != nullptr
? miss_mask + static_cast<int64_t>(row) * top_k
: nullptr;
for (int i = tid; i < top_k; i += blockDim.x) {
int32_t resident_row;
const int32_t g = translate_topk_entry(
row_topk[i], request_row, source_block_table, resident_block_table,
host_rows, source_bt_stride, source_num_reqs, source_num_blocks,
source_block_size, resident_bt_stride, resident_num_reqs,
resident_num_blocks, resident_block_size, resident_null_block,
resident_row);
if (resolved_global_indices != nullptr) {
resolved_global_indices[static_cast<int64_t>(row) * top_k + i] = g;
}
if (resident_row >= 0 || g >= 0) atomicAdd(&s_counters[2], 1);
if (row_miss != nullptr) row_miss[i] = 0;
if (resident_row >= 0) {
store_hot_index(row_out, row_attention, i, resident_row, hot_block_size,
attention_block_stride);
} else if (g < 0) {
store_hot_index(row_out, row_attention, i, -1, hot_block_size,
attention_block_stride);
} else {
int h = hash_slot(g, table_size);
bool inserted = false;
for (int probe = 0; probe < table_size; ++probe) {
const int32_t old = atomicCAS(&s_hash_keys[h], kHashEmpty, g);
if (old == kHashEmpty) {
atomicAdd(&s_counters[1], 1);
inserted = true;
break;
}
if (old == g) {
inserted = true;
break;
}
h = next_slot(h, table_size);
}
if (inserted) {
atomicMin(&s_hash_vals[h], (row - first_row) * top_k + i);
} else {
// A full table: push the union count past the cap so the check
// below traps.
atomicAdd(&s_counters[1], hash_size);
}
}
}
__syncthreads();
if (tid == 0) {
if (valid_counts != nullptr) {
valid_counts[static_cast<int64_t>(row) * valid_count_stride] =
s_counters[2];
}
s_counters[2] = 0;
}
__syncthreads();
}
// The table and eviction slots are sized for max_union_rows (hash_size - 1);
// a larger union would never reach an empty probe slot or run past s_lru_out.
if (s_counters[1] >= hash_size) {
if (tid == 0) {
printf(
"HiSparse residency: request %d resolves %d host rows, above the "
"configured bound of %d\n",
request_row, s_counters[1], hash_size - 1);
}
__trap();
}
// Fully resident rows need only request-relative page translation. Avoid
// scanning or rewriting the hot LRU when no selected row can consult it.
if (s_counters[1] == 0) {
for (int row = first_row + tid; row < end_row; row += blockDim.x) {
if (swap_counts != nullptr) swap_counts[row] = 0;
}
return;
}
// Phase 2: walk hot slots in LRU order, classify hit / evictable, and
// compact them (hits forward, evictables backward) into s_lru_out.
// compact them (hits forward, evictables backward) into s_lru_out. A hit
// records its slot on the union entry.
const int iters_buffer = (num_buffer_chunks + NUM_WARPS - 1) / NUM_WARPS;
int total_hit_count = 0;
int total_evict_count = 0;
@@ -397,36 +510,28 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
const int16_t slot = has_valid_pos ? row_lru[pos] : int16_t(-1);
// Corruption tripwire: lru/dgi are long-lived device state; if some
// external writer (e.g. stray RDMA into reused VRAM) corrupts a slot
// out of [0, region_stride), treat it as no-hit so it degrades to a re-miss
// instead of an unbounded dgi read (phases 3/5 bound the write side).
// out of [0, region_stride), treat it as no-hit so it degrades to a
// re-miss instead of an unbounded dgi read (phase 3 bounds the write
// side).
const int32_t cached_g =
(slot >= 0 && slot < region_stride) ? row_dgi[slot] : -1;
int found_topk_idx = -1;
bool is_hit = false;
if (cached_g >= 0) {
int h = hash_slot(cached_g, hash_size);
int h = hash_slot(cached_g, table_size);
while (true) {
const int32_t k = s_hash_keys[h];
if (k == cached_g) {
found_topk_idx = static_cast<int32_t>(s_hash_vals[h]);
s_hash_vals[h] = kEntryHit | slot;
is_hit = true;
break;
}
if (k == kHashEmpty) break;
h = (h + 1) % hash_size;
h = next_slot(h, table_size);
}
}
const bool is_hit = found_topk_idx >= 0;
const bool is_evictable = has_valid_pos && !is_hit;
if (is_hit) {
s_topk[found_topk_idx] = kTokenDone;
store_hot_index(row_out, row_attention, found_topk_idx,
static_cast<int32_t>(get_physical_hot_row(
hot_block_table, request_row, hot_table_stride,
hot_block_size, slot)),
hot_block_size, attention_block_stride);
}
int local_hit_off = 0;
int local_evict_off = 0;
if (has_valid_chunk) {
@@ -465,39 +570,59 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
}
__syncthreads();
// Reset prefix sums for the miss compaction (token chunks <= buffer
// chunks because hot_size >= top_k).
for (int i = tid; i < num_token_chunks + 1; i += blockDim.x) {
for (int i = tid; i < num_buffer_chunks + 1; i += blockDim.x) {
s_chunk_off[i] = 0;
}
__syncthreads();
// Phase 3: compact misses, assign them eviction slots (oldest first) and
// record the new ownership in device_global_indices.
const int iters_token = (num_token_chunks + NUM_WARPS - 1) / NUM_WARPS;
// Phase 3: walk the rows' positions in (row, top-k) order; the first
// reference to a union entry no slot holds is a miss. Assign misses
// eviction slots (oldest first) and record the new ownership in
// device_global_indices. Swaps fill the rows' swap segments in order; the
// union never exceeds their combined capacity, and position chunks never
// exceed buffer chunks because hot_size covers every row's top-k.
const int num_positions = (end_row - first_row) * top_k;
const int num_position_chunks = (num_positions + kWarpSize - 1) / kWarpSize;
const int iters_position = (num_position_chunks + NUM_WARPS - 1) / NUM_WARPS;
int miss_running_total = 0;
for (int iter = 0; iter < iters_token; iter++) {
int entry_hits = 0;
int entry_misses = 0;
for (int iter = 0; iter < iters_position; iter++) {
const int chunk_idx = warp_id + iter * NUM_WARPS;
const bool has_valid_chunk = chunk_idx < num_token_chunks;
const bool has_valid_chunk = chunk_idx < num_position_chunks;
const int i = chunk_idx * kWarpSize + lane_id;
const bool has_valid_token = has_valid_chunk && (i < top_k);
int32_t g = 0;
const int position = chunk_idx * kWarpSize + lane_id;
int32_t g = -1;
int h = -1;
int32_t entry = 0;
bool is_miss = false;
if (has_valid_token) {
is_miss = s_topk[i] != kTokenDone;
if (is_miss) {
g = s_topk[i];
if (has_valid_chunk && position < num_positions) {
const int row = first_row + position / top_k;
int32_t resident_row;
g = translate_topk_entry(
global_indices[static_cast<int64_t>(row) * input_row_stride +
position % top_k],
request_row, source_block_table, resident_block_table, host_rows,
source_bt_stride, source_num_reqs, source_num_blocks,
source_block_size, resident_bt_stride, resident_num_reqs,
resident_num_blocks, resident_block_size, resident_null_block,
resident_row);
if (resident_row < 0 && g >= 0) {
h = hash_slot(g, table_size);
while (s_hash_keys[h] != g) {
h = next_slot(h, table_size);
}
entry = s_hash_vals[h];
is_miss = entry == position;
}
}
int local_miss_off = 0;
if (has_valid_chunk) {
const unsigned int miss_mask = __ballot_sync(0xFFFFFFFF, is_miss);
local_miss_off = __popc(miss_mask & lanes_before);
const unsigned int miss_mask_bits = __ballot_sync(0xFFFFFFFF, is_miss);
local_miss_off = __popc(miss_mask_bits & lanes_before);
if (lane_id == 0) {
s_chunk_off[chunk_idx + 1] = __popc(miss_mask);
s_chunk_off[chunk_idx + 1] = __popc(miss_mask_bits);
}
}
__syncthreads();
@@ -505,60 +630,112 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
if (warp_id == 0) {
miss_running_total =
warp_inclusive_scan(s_chunk_off, lane_id, chunk_idx + 1,
num_token_chunks + 1, miss_running_total);
num_position_chunks + 1, miss_running_total);
if (tid == 0) {
s_counters[3] = miss_running_total;
}
}
__syncthreads();
if (is_miss) {
const int m = s_chunk_off[chunk_idx] + local_miss_off;
const int16_t evict_slot = s_lru_out[hot_size - 1 - m];
const int64_t compact_index = static_cast<int64_t>(first_row) * top_k + m;
if (evict_slot < 0 || evict_slot >= region_stride) {
// Corruption tripwire: an out-of-range slot (corrupted lru state,
// see phase 2) must not become a hot-cache/dgi write. Resolve the
// entry as invalid (-1, masked by attention, not in the miss set);
// it re-misses on a later step. Phase 5 re-checks the same
// s_lru_out value, so its copy is skipped consistently.
store_hot_index(row_out, row_attention, i, -1, hot_block_size,
attention_block_stride);
// entry as invalid (-1, masked by attention); it re-misses on a later
// step, and its copy is skipped.
s_hash_vals[h] = kEntryInvalid;
if (swap_host_physical_rows != nullptr) {
const int64_t compact_index =
static_cast<int64_t>(batch_row) * top_k + m;
swap_host_physical_rows[compact_index] = g;
swap_device_physical_rows[compact_index] = -1;
}
} else {
// Reuse s_topk as compacted miss scratch: m < i always holds (done
// entries are skipped), so writes never overrun pending reads.
s_topk[m] = g;
const int32_t physical_row = static_cast<int32_t>(
get_physical_hot_row(hot_block_table, request_row, hot_table_stride,
hot_block_size, evict_slot));
store_hot_index(row_out, row_attention, i, physical_row, hot_block_size,
attention_block_stride);
s_hash_vals[h] = kEntryMiss | evict_slot;
if (swap_host_physical_rows != nullptr) {
const int64_t compact_index =
static_cast<int64_t>(batch_row) * top_k + m;
swap_host_physical_rows[compact_index] = g;
swap_device_physical_rows[compact_index] = physical_row;
swap_device_physical_rows[compact_index] =
static_cast<int32_t>(get_physical_hot_row(
hot_block_table, request_row, hot_table_stride,
hot_block_size, evict_slot));
}
if (row_miss != nullptr) row_miss[i] = 1;
row_dgi[evict_slot] = g;
}
entry = s_hash_vals[h];
}
// Hits, first references and references to entries resolved in an
// earlier chunk write their hot index now; later references to an entry
// resolved in this same chunk pass (only across rows) are left to phase 4.
const bool resolved =
h >= 0 && (entry & (kEntryHit | kEntryMiss | kEntryInvalid));
if (resolved) {
write_entry_index(
entry, first_row + position / top_k, position % top_k, top_k,
hot_indices, attention_indices, attention_row_stride, miss_mask,
hot_block_table, request_row, hot_table_stride, hot_block_size,
attention_block_stride, is_miss, entry_hits, entry_misses);
}
if (has_valid_chunk) {
const unsigned int done_bits = __ballot_sync(0xFFFFFFFF, resolved);
if (lane_id == 0) {
s_done[chunk_idx] = done_bits;
}
}
}
__syncthreads();
const int total_hits = s_counters[0];
const int total_misses = top_k - total_hits - s_counters[1];
if (swap_counts != nullptr && tid == 0) {
swap_counts[batch_row] = total_misses;
}
if (stats != nullptr && tid == 0) {
atomicAdd(&stats[0], static_cast<unsigned long long>(total_hits));
atomicAdd(&stats[1], static_cast<unsigned long long>(total_misses));
const int total_misses = s_counters[3];
if (swap_counts != nullptr) {
for (int row = first_row + tid; row < end_row; row += blockDim.x) {
const int remaining = total_misses - (row - first_row) * top_k;
swap_counts[row] = max(0, min(top_k, remaining));
}
}
// Phase 4: write back the LRU order: stale evictables at the front,
// Phase 4: references phase 3 left unresolved (later references across rows
// to entries loaded in the same chunk pass) map onto the union's slots.
// Single-row requests resolve everything in phase 3.
if (end_row - first_row > 1) {
for (int position = tid; position < num_positions; position += blockDim.x) {
if (s_done[position / kWarpSize] & (1u << (position % kWarpSize))) {
continue;
}
const int row = first_row + position / top_k;
int32_t resident_row;
const int32_t g = translate_topk_entry(
global_indices[static_cast<int64_t>(row) * input_row_stride +
position % top_k],
request_row, source_block_table, resident_block_table, host_rows,
source_bt_stride, source_num_reqs, source_num_blocks,
source_block_size, resident_bt_stride, resident_num_reqs,
resident_num_blocks, resident_block_size, resident_null_block,
resident_row);
if (resident_row >= 0 || g < 0) continue;
int h = hash_slot(g, table_size);
while (s_hash_keys[h] != g) {
h = next_slot(h, table_size);
}
write_entry_index(
s_hash_vals[h], row, position % top_k, top_k, hot_indices,
attention_indices, attention_row_stride, miss_mask, hot_block_table,
request_row, hot_table_stride, hot_block_size, attention_block_stride,
false, entry_hits, entry_misses);
}
}
if (stats != nullptr) {
atomicAdd(&s_counters[4], entry_hits);
atomicAdd(&s_counters[2], entry_misses);
__syncthreads();
if (tid == 0) {
atomicAdd(&stats[0], static_cast<unsigned long long>(s_counters[4]));
atomicAdd(&stats[1], static_cast<unsigned long long>(s_counters[2]));
}
}
// Phase 5: write back the LRU order: stale evictables at the front,
// freshly loaded misses next, then hits at MRU.
const int total_evictable = hot_size - total_hits;
const int remaining_evictable = total_evictable - total_misses;
@@ -783,7 +960,7 @@ void hisparse_resolve_residency(
torch::stable::Tensor& device_global_indices,
torch::stable::Tensor& lru_slots,
std::optional<torch::stable::Tensor> const& request_state_indices,
int64_t region_stride,
int64_t region_stride, int64_t max_union_rows,
std::optional<torch::stable::Tensor> const& miss_mask,
std::optional<torch::stable::Tensor> const& stats,
std::optional<torch::stable::Tensor> const& attention_indices,
@@ -1034,14 +1211,32 @@ void hisparse_resolve_residency(
}
constexpr int kBlockSize = 1024;
const int hash_size = 2 * top_k;
STD_TORCH_CHECK(max_union_rows > 0 && max_union_rows <= hot_size,
"max_union_rows must be in [1, hot_size], got ",
max_union_rows);
// One spare entry above the union bound keeps every probe sequence finite.
const int hash_size = static_cast<int>(max_union_rows) + 1;
const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize;
const size_t smem_bytes =
sizeof(int32_t) * (top_k + 2 * (num_buffer_chunks + 1) + hash_size + 3) +
sizeof(int16_t) * (hot_size + hash_size);
sizeof(int32_t) * (2 * hash_size + 2 * (num_buffer_chunks + 1) +
kResidencyCounters + num_buffer_chunks) +
sizeof(int16_t) * hot_size;
const torch::stable::accelerator::DeviceGuard device_guard(
hot_cache.get_device_index());
int max_smem_bytes = 0;
STD_TORCH_CHECK(cudaDeviceGetAttribute(
&max_smem_bytes, cudaDevAttrMaxSharedMemoryPerBlockOptin,
hot_cache.get_device_index()) == cudaSuccess,
"failed to query shared memory per block");
STD_TORCH_CHECK(smem_bytes <= static_cast<size_t>(max_smem_bytes),
"HiSparse device_buffer_size=", hot_size, " needs ",
smem_bytes,
" bytes of shared memory to resolve residency, "
"but the device allows ",
max_smem_bytes,
"; lower device_buffer_size or the number of speculative "
"tokens.");
const cudaStream_t stream = get_current_cuda_stream();
const int64_t attention_row_stride =
attention_indices.has_value() ? attention_indices.value().stride(0) : 0;
csrc/libtorch_stable/ops.h
+1 -1
View File
@@ -680,7 +680,7 @@ void hisparse_resolve_residency(
torch::stable::Tensor& device_global_indices,
torch::stable::Tensor& lru_slots,
std::optional<torch::stable::Tensor> const& request_state_indices,
int64_t region_stride,
int64_t region_stride, int64_t max_union_rows,
std::optional<torch::stable::Tensor> const& miss_mask,
std::optional<torch::stable::Tensor> const& stats,
std::optional<torch::stable::Tensor> const& attention_indices,
csrc/libtorch_stable/torch_bindings.cpp
+1
View File
@@ -995,6 +995,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
" Tensor! lru_slots,"
" Tensor? request_state_indices,"
" int region_stride,"
" int max_union_rows,"
" Tensor(a!)? miss_mask=None,"
" Tensor(b!)? stats=None,"
" Tensor(c!)? attention_indices=None,"
docs/design/hisparse.md
+4 -4
View File
@@ -130,10 +130,10 @@ path. The resolver consumes the existing graph-stable request mapping from
attention metadata; neither the worker nor individual cache handles keep a
duplicate mapping.
Speculative decoding resolves and consumes each verification step in order.
Each step receives distinct replayable plan rows while sharing the request's
hot-cache state, so a later step cannot reuse a hot row before an earlier step
has consumed it.
Speculative decoding resolves all verification rows of a request in one pass:
one block resolves the union of the rows' top-k against the request's hot-cache
state, so rows that select the same host row share its hot row and no row
evicts a hot row another row of the step still reads.
## P/D import target
tests/v1/attention/test_sparse_mla_backends.py
+229 -107
View File
@@ -68,10 +68,10 @@ from vllm.utils.torch_utils import current_stream
from vllm.v1.attention.backends.mla import index_group as index_group_module
from vllm.v1.attention.backends.mla.flashattn_mla_sparse import (
FlashAttnMLASparseImpl,
FlashAttnMLASparseMetadataBuilder,
)
from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
FlashInferMLASparseImpl,
FlashInferMLASparseMetadataBuilder,
FlashInferMLASparseTRTLLMBackend,
)
from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import (
@@ -1880,11 +1880,13 @@ def _make_hisparse_runtime(
block_size: int = 64,
max_swap_rows: int | None = None,
index_group: hisparse_runtime.HiSparseIndexGroup | None = None,
max_union_rows: int | None = None,
) -> HiSparseRuntime:
runtime = HiSparseRuntime(
config=ResolvedHiSparseConfig(
top_k=top_k,
device_buffer_size=device_buffer_size,
max_union_rows=max_union_rows or device_buffer_size,
),
max_num_reqs=max_num_reqs,
row_width=row_width,
@@ -2015,6 +2017,167 @@ def test_hisparse_maps_speculative_rows_through_request_state():
assert physical.tolist() == [[block_size]] * 4
@requires_hisparse_ops
@pytest.mark.parametrize(
"num_rows,row_stride,top_k",
[(8, 128, 128), (4, 32, 128), (8, 2048, 2048)],
ids=["disjoint", "overlap", "long_mtp"],
)
def test_hisparse_speculative_rows_resolve_host_misses_consistently(
num_rows, row_stride, top_k
):
"""Verification rows of one request must agree on its hot-buffer slots.
Each row resolved residency in its own thread block against the request's
shared slot ownership and LRU, so rows missing on different host rows could
claim the same free slot. One row's hot index then pointed at a slot holding
another row's KV. Overlapping rows must share one slot per host row, and
long_mtp's rows span several resolution passes.
"""
device = torch.device(DEVICE_TYPE)
block_size, row_width = 64, 64
num_blocks = 2 * num_rows * top_k // block_size
kv_pool = torch.randn(
(num_blocks + 1, block_size, row_width), dtype=torch.float32
).pin_memory()
flat_pool = kv_pool.reshape(-1, row_width)
runtime = _make_hisparse_runtime(
top_k=top_k,
device_buffer_size=num_rows * top_k,
max_num_reqs=1,
row_width=row_width,
block_size=block_size,
max_swap_rows=num_rows,
)
runtime.bind_source_cache(kv_pool)
cache = HiSparseCacheHandle(runtime)
cache.all_context_pages_resident = False
block_table = torch.arange(
1, num_blocks + 1, dtype=torch.int32, device=device
).view(1, num_blocks)
req_ids = torch.zeros(num_rows, dtype=torch.int32, device=device)
rows = torch.arange(num_rows, dtype=torch.int32, device=device)[:, None]
base = torch.arange(top_k, dtype=torch.int32, device=device)[None, :]
for step in range(4):
runtime.begin_forward()
# Rows start row_stride apart (disjoint at row_stride == top_k) and
# shift by half a top-k per step: all misses first, then a mix.
topk = (rows * row_stride + base + step * (top_k // 2)) % (
num_blocks * block_size
)
hot_indices = cache.swap_in(
req_id_per_token=req_ids,
block_table=block_table,
logical_topk_indices=topk.clone(),
block_size=block_size,
)
torch.accelerator.synchronize()
global_ref = _triton_convert_reference_impl(
req_ids, block_table, topk, block_size, top_k
)
gathered = runtime.hot.attention_cache.reshape(-1, row_width)[
hot_indices.to(torch.long)
].cpu()
torch.testing.assert_close(gathered, flat_pool[global_ref.cpu().to(torch.long)])
@requires_hisparse_ops
def test_hisparse_speculative_rows_count_a_shared_load_once():
"""A host row two verification rows select loads once; the reuse is a hit."""
device = torch.device(DEVICE_TYPE)
block_size, row_width, top_k = 64, 64, 128
num_blocks = top_k // block_size
kv_pool = torch.randn(
(num_blocks + 1, block_size, row_width), dtype=torch.float32
).pin_memory()
runtime = _make_hisparse_runtime(
top_k=top_k,
device_buffer_size=2 * top_k,
max_num_reqs=1,
row_width=row_width,
block_size=block_size,
max_swap_rows=2,
)
runtime.bind_source_cache(kv_pool)
cache = HiSparseCacheHandle(runtime)
cache.all_context_pages_resident = False
block_table = torch.arange(
1, num_blocks + 1, dtype=torch.int32, device=device
).view(1, num_blocks)
topk = torch.arange(top_k, dtype=torch.int32, device=device).repeat(2, 1)
runtime.begin_forward()
cache.swap_in(
req_id_per_token=torch.zeros(2, dtype=torch.int32, device=device),
block_table=block_table,
logical_topk_indices=topk,
block_size=block_size,
)
torch.accelerator.synchronize()
assert runtime.index_group.swap_stats.cpu().tolist() == [top_k, top_k]
@requires_hisparse_ops
def test_hisparse_union_table_scales_with_union_bound():
"""Shared memory follows the per-request union bound, not the hot buffer.
Sizing the table from device_buffer_size needs ~340 KB here and fails.
"""
device = torch.device(DEVICE_TYPE)
block_size, row_width, top_k = 64, 64, 128
num_blocks = top_k // block_size
kv_pool = torch.randn(
(num_blocks + 1, block_size, row_width), dtype=torch.float32
).pin_memory()
runtime = _make_hisparse_runtime(
top_k=top_k,
device_buffer_size=32768,
max_union_rows=top_k,
max_num_reqs=1,
row_width=row_width,
block_size=block_size,
max_swap_rows=1,
)
runtime.bind_source_cache(kv_pool)
cache = HiSparseCacheHandle(runtime)
cache.all_context_pages_resident = False
req_ids = torch.zeros(1, dtype=torch.int32, device=device)
block_table = torch.arange(
1, num_blocks + 1, dtype=torch.int32, device=device
).view(1, num_blocks)
topk = torch.arange(top_k, dtype=torch.int32, device=device).view(1, top_k)
runtime.begin_forward()
hot_indices = cache.swap_in(
req_id_per_token=req_ids,
block_table=block_table,
logical_topk_indices=topk,
block_size=block_size,
)
torch.accelerator.synchronize()
global_ref = _triton_convert_reference_impl(
req_ids, block_table, topk, block_size, top_k
)
gathered = runtime.hot.attention_cache.reshape(-1, row_width)[
hot_indices.to(torch.long)
].cpu()
torch.testing.assert_close(
gathered, kv_pool.reshape(-1, row_width)[global_ref.cpu().to(torch.long)]
)
@requires_hisparse_ops
def test_hisparse_rejects_union_bound_beyond_shared_memory():
"""An oversized union bound fails at construction, not on the first forward."""
with pytest.raises(ValueError, match="shared memory"):
_make_hisparse_runtime(top_k=2048, device_buffer_size=32768)
@requires_hisparse_ops
def test_hisparse_resident_rows_bypass_hot_lru():
device = torch.device(DEVICE_TYPE)
@@ -3001,8 +3164,8 @@ def test_hisparse_gather_prefill_cache_prefers_resident_rows():
torch.testing.assert_close(staged_flat[i], expected)
def test_hisparse_fp8_decode_resolves_steps_then_runs_batched_attention(monkeypatch):
"""FP8 decode must use active dimensions when common metadata is padded."""
def test_hisparse_fp8_decode_resolves_rows_once_then_runs_batched_attention():
"""Verification rows resolve in one swap_in; FP8 attention sees them batched."""
device = torch.device("cpu")
num_decodes = 2
query_len = 3
@@ -3013,48 +3176,31 @@ def test_hisparse_fp8_decode_resolves_steps_then_runs_batched_attention(monkeypa
.view(-1, 1)
.expand(-1, 4)
)
steps: list[int] = []
kernel_shapes: list[torch.Size] = []
def swap_in(req_ids, *, logical_topk_indices, **kwargs): # noqa: ARG001
steps.append(int(logical_topk_indices[0, 0]))
output = kwargs["attention_indices_out"]
if output is not None:
output.copy_(logical_topk_indices + 10)
return topk[:num_decodes]
return logical_topk_indices + 10
def run_kernel(self, *, q, **kwargs): # noqa: ARG001
kernel_shapes.append(q.shape)
return q[..., :1], None
runtime = SimpleNamespace(
hot=SimpleNamespace(attention_cache=torch.empty(1, device=device))
)
leader_cache = SimpleNamespace(
runtime=runtime,
source_block_table=torch.empty(
num_decodes, 1, dtype=torch.int32, device=device
cache = SimpleNamespace(
runtime=SimpleNamespace(
hot=SimpleNamespace(attention_cache=torch.empty(1, device=device))
),
swap_in=MagicMock(side_effect=swap_in),
)
follower_cache = SimpleNamespace(
runtime=runtime,
source_block_table=torch.empty(
num_decodes, 1, dtype=torch.int32, device=device
),
swap_in=MagicMock(side_effect=swap_in),
)
index_group = object.__new__(HiSparseMLAIndexGroup)
index_group.caches = [leader_cache, follower_cache]
index_group.caches = [cache]
index_group.physical_topk_indices = torch.empty(
(num_tokens + 1, 4), dtype=torch.int32, device=device
)
index_group.valid_topk_counts = torch.empty(
num_tokens + 1, dtype=torch.int32, device=device
)
index_group.request_ids = torch.arange(
num_decodes, dtype=torch.int32, device=device
)
index_group.row_indices = torch.arange(num_tokens, dtype=torch.int32, device=device)
index_group.request_ids = torch.empty_like(index_group.row_indices)
impl = SimpleNamespace(
kv_lora_rank=1,
index_group=index_group,
@@ -3062,84 +3208,83 @@ def test_hisparse_fp8_decode_resolves_steps_then_runs_batched_attention(monkeypa
)
impl._fp8_flash_mla_kernel = MethodType(run_kernel, impl)
metadata = SimpleNamespace(
num_decodes=num_tokens,
num_decode_tokens=num_tokens,
decode_max_query_len=1,
query_start_loc=torch.arange(
0,
num_tokens + 1,
query_len,
dtype=torch.int32,
device=device,
0, num_tokens + 1, query_len, dtype=torch.int32, device=device
),
block_table=torch.empty(num_decodes, 1, dtype=torch.int32, device=device),
req_id_per_token=torch.arange(
num_decodes, dtype=torch.int32, device=device
).repeat_interleave(query_len),
block_size=64,
)
output = FlashMLASparseImpl._host_backed_fp8_decode(
impl,
q,
topk,
metadata,
SimpleNamespace(),
num_decodes,
query_len,
impl, q, topk, metadata, SimpleNamespace(), num_decodes
)
assert steps == [0, 1, 2]
cache.swap_in.assert_called_once()
torch.testing.assert_close(
cache.swap_in.call_args.args[0], metadata.req_id_per_token
)
torch.testing.assert_close(
cache.swap_in.call_args.kwargs["logical_topk_indices"], topk
)
assert kernel_shapes == [(num_decodes, query_len, 2, 4)]
assert output.shape == (num_tokens, 2, 1)
torch.testing.assert_close(
index_group.physical_topk_indices[:num_tokens], topk + 10
def test_hisparse_keeps_speculative_rows_as_decodes():
"""HiSparse limits decodes to verification requests, not short extends.
FlashAttn sparse MLA used to classify MTP verification as prefill under
HiSparse, whose residency-dependent route a FULL CUDA graph freezes at
capture.
"""
builder = object.__new__(FlashAttnMLASparseMetadataBuilder)
builder.vllm_config = SimpleNamespace(
attention_config=SimpleNamespace(hisparse_config=object()),
speculative_config=SimpleNamespace(
num_speculative_tokens=3, parallel_drafting=False
),
parallel_config=SimpleNamespace(decode_context_parallel_size=1),
use_v2_model_runner=True,
)
follower_indices = index_group.convert_decode_logical_to_physical_topk(
1,
topk,
metadata,
return_valid_counts=False,
num_decodes=num_decodes,
decode_query_len=query_len,
)
assert steps == [0, 1, 2, 0, 1, 2]
torch.testing.assert_close(follower_indices, topk + 10)
builder._init_reorder_batch_threshold(256, supports_spec_as_decode=True)
assert builder.reorder_batch_threshold == 4
def test_hisparse_varlen_decode_uses_canonical_request_rows():
"""Graph metadata's query upper bound must preserve packed token rows."""
num_decodes = 2
topk = torch.arange(18, dtype=torch.int32).view(9, 2)
def swap_in(_request_ids, *, logical_topk_indices, **_kwargs):
return logical_topk_indices + 10
def test_hisparse_decode_skips_padding_rows():
"""Padding rows are masked, and the ids stay in the group's persistent buffer.
The resolver reads them on the HiSparse copy stream, where a compute-stream
temporary can already be recycled (CUDA graphs do), scrambling rows' requests.
"""
topk = torch.arange(20, dtype=torch.int32).view(10, 2)
cache = SimpleNamespace(
source_block_table=torch.zeros((num_decodes, 1), dtype=torch.int32),
swap_in=MagicMock(side_effect=swap_in),
source_block_table=torch.zeros((2, 1), dtype=torch.int32),
swap_in=MagicMock(return_value=topk),
)
index_group = object.__new__(HiSparseMLAIndexGroup)
index_group.caches = [cache]
index_group.physical_topk_indices = torch.empty((10, 2), dtype=torch.int32)
index_group.valid_topk_counts = torch.empty(10, dtype=torch.int32)
index_group.request_ids = torch.arange(num_decodes, dtype=torch.int32)
index_group.physical_topk_indices = torch.empty((11, 2), dtype=torch.int32)
index_group.row_indices = torch.arange(10, dtype=torch.int32)
index_group.request_ids = torch.empty_like(index_group.row_indices)
metadata = SimpleNamespace(
num_decodes=num_decodes,
num_decode_tokens=topk.shape[0],
decode_max_query_len=15,
query_start_loc=torch.tensor([0, 8, 9], dtype=torch.int32),
req_id_per_token=torch.tensor([0] * 8 + [1, 0], dtype=torch.int32),
block_size=64,
)
physical_topk = index_group.convert_decode_logical_to_physical_topk(
0,
topk,
metadata,
return_valid_counts=False,
index_group.convert_logical_to_physical_topk(
0, topk, metadata, block_stride_rows=None, return_valid_counts=False
)
torch.testing.assert_close(physical_topk, topk + 10)
request_rows = cache.swap_in.call_args_list[0].args[0]
torch.testing.assert_close(request_rows, index_group.request_ids)
request_ids = cache.swap_in.call_args.args[0]
assert request_ids.data_ptr() == index_group.request_ids.data_ptr()
torch.testing.assert_close(
request_ids, torch.tensor([0] * 8 + [1, -1], dtype=torch.int32)
)
def test_flashinfer_hisparse_decode_runs_batched_attention():
@@ -3168,7 +3313,7 @@ def test_flashinfer_hisparse_decode_runs_batched_attention():
)
index_group = object.__new__(HiSparseMLAIndexGroup)
index_group.caches = [cache_handle]
index_group.convert_decode_logical_to_physical_topk = MethodType(
index_group.convert_logical_to_physical_topk = MethodType(
convert_decode, index_group
)
impl = object.__new__(FlashInferMLASparseImpl)
@@ -3200,7 +3345,7 @@ def test_flashattn_hisparse_decode_uses_index_group():
physical = topk.clone()
counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE)
index_group = object.__new__(HiSparseMLAIndexGroup)
index_group.convert_decode_logical_to_physical_topk = MagicMock(
index_group.convert_logical_to_physical_topk = MagicMock(
return_value=(physical, counts)
)
index_group.physical_kv_cache = MagicMock(
@@ -3223,7 +3368,7 @@ def test_flashattn_hisparse_decode_uses_index_group():
assert output.shape == (num_tokens, 2, 1)
assert lse is None
index_group.convert_decode_logical_to_physical_topk.assert_called_once()
index_group.convert_logical_to_physical_topk.assert_called_once()
impl._run_mqa_kernel.assert_called_once()
@@ -3233,9 +3378,7 @@ def test_flashinfer_sm120_hisparse_decode_uses_index_group():
topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=DEVICE_TYPE)
physical = topk.clone()
index_group = object.__new__(HiSparseMLAIndexGroup)
index_group.convert_decode_logical_to_physical_topk = MagicMock(
return_value=physical
)
index_group.convert_logical_to_physical_topk = MagicMock(return_value=physical)
index_group.physical_kv_cache = MagicMock(
return_value=torch.empty(1, device=DEVICE_TYPE)
)
@@ -3259,7 +3402,7 @@ def test_flashinfer_sm120_hisparse_decode_uses_index_group():
assert output.shape == (num_tokens, 2, 1)
assert lse is None
index_group.convert_decode_logical_to_physical_topk.assert_called_once()
index_group.convert_logical_to_physical_topk.assert_called_once()
impl._run_mqa_kernel.assert_called_once()
@@ -3291,27 +3434,6 @@ def test_hisparse_resident_prefill_uses_attention_block_stride():
assert index_group._convert_once.call_args.kwargs["block_stride_rows"] == 832
def test_hisparse_shared_sparse_builder_routes_multi_token_chunks_to_prefill():
"""Default sparse backends must route multi-token chunks to prefill."""
builder = object.__new__(FlashInferMLASparseMetadataBuilder)
builder.vllm_config = SimpleNamespace(
attention_config=SimpleNamespace(hisparse_config=object()),
speculative_config=SimpleNamespace(
num_speculative_tokens=4,
parallel_drafting=False,
),
parallel_config=SimpleNamespace(decode_context_parallel_size=1),
)
builder._init_reorder_batch_threshold(
1024,
supports_spec_as_decode=True,
supports_dcp_with_varlen=True,
)
assert builder.reorder_batch_threshold == 1
def test_hisparse_flashmla_reorders_full_speculative_window_as_decode():
builder = object.__new__(FlashMLASparseMetadataBuilder)
builder.vllm_config = SimpleNamespace(
tests/v1/worker/test_utils.py
+1
View File
@@ -1156,6 +1156,7 @@ def test_hisparse_cache_handles_join_index_groups_during_construction(monkeypatc
resolved = hisparse_runtime_module.ResolvedHiSparseConfig(
top_k=4,
device_buffer_size=8,
max_union_rows=8,
)
monkeypatch.setattr(hisparse_runtime_module, "_has_hisparse_ops", lambda: True)
monkeypatch.setattr(
vllm/model_executor/layers/attention/sparse_mla_attention.py
-3
View File
@@ -165,7 +165,6 @@ T = TypeVar("T", bound=SparseMLACommonMetadata)
class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
metadata_cls: type[T]
require_uniform_decodes: ClassVar[bool] = False
hisparse_supports_multi_token_decode: ClassVar[bool] = False
def __init__(
self,
@@ -259,8 +258,6 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
) -> None:
if self.vllm_config.attention_config.hisparse_config is not None:
reorder_batch_threshold = 1
if not self.hisparse_supports_multi_token_decode:
supports_spec_as_decode = False
super()._init_reorder_batch_threshold(
reorder_batch_threshold,
supports_spec_as_decode,
vllm/v1/attention/backends/mla/flashattn_mla_sparse.py
+2 -1
View File
@@ -207,10 +207,11 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
outputs = []
if num_decode_tokens:
physical_topk, valid_counts = (
index_group.convert_decode_logical_to_physical_topk(
index_group.convert_logical_to_physical_topk(
self.index_group_index,
topk_indices[:num_decode_tokens],
attn_metadata,
block_stride_rows=None,
return_valid_counts=True,
)
)
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+2 -2
View File
@@ -273,7 +273,6 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde
"""Metadata builder for the SM100 TRT-LLM sparse MLA kernel."""
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
hisparse_supports_multi_token_decode: ClassVar[bool] = True
def __init__(
self,
@@ -502,10 +501,11 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
decode_lse: torch.Tensor | None = None
if num_decode_tokens > 0:
physical_topk, valid_counts = (
index_group.convert_decode_logical_to_physical_topk(
index_group.convert_logical_to_physical_topk(
self.index_group_index,
topk_indices[:num_decode_tokens],
attn_metadata,
block_stride_rows=None,
return_valid_counts=True,
)
)
vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py
+2 -1
View File
@@ -125,10 +125,11 @@ class FlashInferMLASparseSM120Impl(SparseMLACommonImpl[FlashInferMLASparseMetada
if num_decode_tokens:
topk_indices_physical = cast(
torch.Tensor,
index_group.convert_decode_logical_to_physical_topk(
index_group.convert_logical_to_physical_topk(
self.index_group_index,
topk_indices[:num_decode_tokens],
attn_metadata,
block_stride_rows=None,
return_valid_counts=False,
),
)
vllm/v1/attention/backends/mla/flashmla_sparse.py
+4 -7
View File
@@ -317,7 +317,6 @@ class FlashMLASparseMetadataBuilder(
):
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
require_uniform_decodes: ClassVar[bool] = True
hisparse_supports_multi_token_decode: ClassVar[bool] = True
metadata_cls = FlashMLASparseMetadata
def __init__(
@@ -818,10 +817,11 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
num_decode_tokens = attn_metadata.num_decode_tokens
if num_decode_tokens > 0:
decode_topk, decode_lengths = (
index_group.convert_decode_logical_to_physical_topk(
index_group.convert_logical_to_physical_topk(
self.index_group_index,
topk_indices[:num_decode_tokens],
attn_metadata,
block_stride_rows=None,
return_valid_counts=True,
)
)
@@ -1027,7 +1027,6 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
attn_metadata,
fp8_metadata.decode.kernel_metadata,
num_decodes,
fp8_metadata.decode.decode_query_len,
)
# Reshape q: (num_decode_tokens, num_heads, head_dim)
# -> (num_decodes, seq_len, num_heads, head_dim)
@@ -1267,16 +1266,14 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
attn_metadata: FlashMLASparseMetadata,
kernel_metadata: FlashMLASparseMetadata.FP8KernelMetadata,
num_decodes: int,
decode_query_len: int,
) -> torch.Tensor:
assert isinstance(self.index_group, HiSparseMLAIndexGroup)
physical_topk = self.index_group.convert_decode_logical_to_physical_topk(
physical_topk = self.index_group.convert_logical_to_physical_topk(
self.index_group_index,
topk_indices,
attn_metadata,
block_stride_rows=None,
return_valid_counts=False,
num_decodes=num_decodes,
decode_query_len=decode_query_len,
)
assert isinstance(physical_topk, torch.Tensor)
q = reshape_query_for_spec_decode(q, num_decodes)
vllm/v1/attention/backends/mla/index_group.py
+16 -104
View File
@@ -39,6 +39,7 @@ class SparseMLAIndexGroup:
logical_topk_indices: torch.Tensor
physical_topk_indices: torch.Tensor
valid_topk_counts: torch.Tensor
row_indices: torch.Tensor
request_ids: torch.Tensor
side_stream: torch.Stream
logical_topk_ready: torch.Event
@@ -230,12 +231,10 @@ class HiSparseMLAIndexGroup(SparseMLAIndexGroup):
*,
block_stride_rows: int | None,
return_valid_counts: bool,
req_id_per_token: torch.Tensor | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
cache = self.cache(layer_index)
num_tokens = logical_topk_indices.shape[0]
if req_id_per_token is None:
req_id_per_token = attn_metadata.req_id_per_token[:num_tokens]
req_id_per_token = attn_metadata.req_id_per_token[:num_tokens]
if num_tokens > self.physical_topk_indices.shape[0]:
# Prefill-sized batches do not fit the decode residency workspace.
# Non-resident prefills are staged before reaching this path.
@@ -253,113 +252,21 @@ class HiSparseMLAIndexGroup(SparseMLAIndexGroup):
)
source_block_table = cache.source_block_table
assert source_block_table is not None
# CUDA-graph padding rows past the batch's tokens map to request 0;
# mark them -1 so residency resolution skips them.
request_ids = self.request_ids[:num_tokens]
request_ids.copy_(req_id_per_token)
request_ids.masked_fill_(
self.row_indices[:num_tokens] >= attn_metadata.query_start_loc[-1], -1
)
return cache.swap_in(
req_id_per_token,
request_ids,
block_table=source_block_table,
logical_topk_indices=logical_topk_indices,
block_size=attn_metadata.block_size,
return_valid_counts=return_valid_counts,
)
def convert_decode_logical_to_physical_topk(
self,
layer_index: int,
logical_topk_indices: torch.Tensor,
attn_metadata: Any,
*,
return_valid_counts: bool,
num_decodes: int | None = None,
decode_query_len: int | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
num_tokens = logical_topk_indices.shape[0]
if num_decodes is None:
num_decodes = attn_metadata.num_decodes
if decode_query_len is None:
decode_query_len = attn_metadata.decode_max_query_len
if decode_query_len == 1:
return self.convert_logical_to_physical_topk(
layer_index,
logical_topk_indices,
attn_metadata,
block_stride_rows=None,
return_valid_counts=return_valid_counts,
req_id_per_token=self.request_ids[:num_decodes],
)
if num_tokens != num_decodes * decode_query_len:
query_start_loc = attn_metadata.query_start_loc[: num_decodes + 1]
physical_topk_indices = self.physical_topk_indices[: num_tokens + 1]
valid_topk_counts = self.valid_topk_counts[: num_tokens + 1]
request_ids = self.request_ids[:num_decodes]
cache = self.cache(layer_index)
source_block_table = cache.source_block_table
assert source_block_table is not None
for step in range(decode_query_len):
token_indices = query_start_loc[:-1] + step
active = (token_indices < query_start_loc[1:]) & (
token_indices < num_tokens
)
output_indices = torch.where(
active, token_indices, torch.full_like(token_indices, num_tokens)
).long()
token_indices = token_indices.clamp(0, num_tokens - 1).long()
step_topk = logical_topk_indices.index_select(
0, token_indices
).masked_fill(~active.unsqueeze(1), -1)
step_result = cache.swap_in(
request_ids,
block_table=source_block_table,
logical_topk_indices=step_topk,
block_size=attn_metadata.block_size,
return_valid_counts=return_valid_counts,
)
if layer_index == 0:
if return_valid_counts:
step_indices, step_counts = step_result
valid_topk_counts.index_copy_(0, output_indices, step_counts)
else:
step_indices = step_result
physical_topk_indices.index_copy_(0, output_indices, step_indices)
physical_topk_indices = physical_topk_indices[:num_tokens]
if return_valid_counts:
return physical_topk_indices, valid_topk_counts[:num_tokens]
return physical_topk_indices
assert num_tokens == num_decodes * decode_query_len
logical_topk_by_request = logical_topk_indices.view(
num_decodes, decode_query_len, -1
)
physical_topk_by_request = self.physical_topk_indices[:num_tokens].view(
num_decodes, decode_query_len, -1
)
valid_topk_by_request = self.valid_topk_counts[:num_tokens].view(
num_decodes, decode_query_len
)
request_ids = self.request_ids[:num_decodes]
for step in range(decode_query_len):
cache = self.cache(layer_index)
source_block_table = cache.source_block_table
assert source_block_table is not None
cache.swap_in(
request_ids,
block_table=source_block_table,
logical_topk_indices=logical_topk_by_request[:, step],
block_size=attn_metadata.block_size,
return_valid_counts=return_valid_counts,
attention_indices_out=(
physical_topk_by_request[:, step] if layer_index == 0 else None
),
valid_counts_out=(
valid_topk_by_request[:, step]
if layer_index == 0 and return_valid_counts
else None
),
)
physical_topk_indices = physical_topk_by_request.view(num_tokens, -1)
if return_valid_counts:
return physical_topk_indices, valid_topk_by_request.view(num_tokens)
return physical_topk_indices
def stage_prefill_rows(
self, layer_index: int, kv_cache: torch.Tensor, attn_metadata: Any
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -477,7 +384,12 @@ class SparseMLAIndexGroupBuilder:
dtype=torch.int32,
device=self.logical_topk_indices.device,
),
request_ids=torch.arange(
row_indices=torch.arange(
workspace_rows,
dtype=torch.int32,
device=self.logical_topk_indices.device,
),
request_ids=torch.empty(
workspace_rows,
dtype=torch.int32,
device=self.logical_topk_indices.device,
vllm/v1/hisparse/runtime.py
+37 -1
View File
@@ -19,7 +19,8 @@ from vllm.distributed import get_tp_group
from vllm.forward_context import get_forward_context, in_piecewise_cudagraph
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.utils.math_utils import round_up
from vllm.utils.math_utils import cdiv, round_up
from vllm.utils.mem_utils import get_max_shared_memory_bytes
from vllm.utils.torch_utils import current_stream
from vllm.v1.attention.backend import max_decode_query_len
from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
@@ -96,6 +97,9 @@ class PagedCacheView:
class ResolvedHiSparseConfig:
top_k: int
device_buffer_size: int
# Most host rows one request's decode rows can select in a step: every
# verification row's top-k.
max_union_rows: int
eager_host_mirror: bool = True
@classmethod
@@ -151,10 +155,38 @@ class ResolvedHiSparseConfig:
return cls(
top_k=model_top_k,
device_buffer_size=device_buffer_size,
max_union_rows=min_device_buffer_size,
eager_host_mirror=config.eager_host_mirror,
)
def _check_residency_shared_memory(
device: torch.device, config: ResolvedHiSparseConfig
) -> None:
"""Reject a config whose residency resolver exceeds per-block shared memory."""
# Mirrors hisparse_resolve_residency's layout: union hash keys and values
# (plus an empty slot), per-chunk offsets, five counters, done bits, and
# the int16 compacted LRU.
hot_size = config.device_buffer_size
num_chunks = cdiv(hot_size, 32)
required = 4 * (2 * (config.max_union_rows + 1) + 3 * num_chunks + 7) + 2 * hot_size
device_index = (
device.index
if device.index is not None
else torch.accelerator.current_device_index()
)
available = get_max_shared_memory_bytes(device_index)
if required > available:
raise ValueError(
"HiSparse residency resolution needs "
f"{required} B of shared memory per block for "
f"device_buffer_size={config.device_buffer_size} and "
f"{config.max_union_rows} rows per request "
"(decode query length x index_topk), but the device allows "
f"{available} B. Lower device_buffer_size or num_speculative_tokens."
)
def check_hisparse_host_memory(pool_bytes: int) -> None:
"""Reject a physical host-pool allocation that cannot fit in RAM."""
mem = psutil.virtual_memory()
@@ -677,6 +709,9 @@ class HiSparseRuntime:
# Logical slots per request. Physical rows come from its ephemeral HMA
# block table and need not be contiguous in the shared slab.
self.region_stride = config.device_buffer_size
self.max_union_rows = config.max_union_rows
if self.device.type == "cuda":
_check_residency_shared_memory(self.device, config)
row_bytes = row_width * kv_dtype.itemsize
if row_bytes % 16 != 0:
@@ -919,6 +954,7 @@ class HiSparseRuntime:
group.lru_slots,
request_state_indices,
self.region_stride,
self.max_union_rows,
None,
group.swap_stats,
physical_topk_indices,