Select DeepSeek V4.1 index candidates with a radix select on Metal
This commit is contained in:
Tarjei Mandt committed
1 parent
4b1b1519ed
commit
af7c02b594
3 files changed
+433
-1
No files matched your search
+95
@@ -680,6 +680,9 @@ static id<MTLBuffer> g_router_selection_buffer;
|
||||
static id<MTLBuffer> g_router_weight_sum_buffer;
|
||||
static id<MTLBuffer> g_indexer_head_scores_buffer;
|
||||
static id<MTLBuffer> g_indexer_topk_buffer;
|
||||
static id<MTLBuffer> g_topk_select_buffer;
|
||||
static NSUInteger g_topk_select_bytes;
|
||||
static uint32_t g_topk_select_parity;
|
||||
static id<MTLBuffer> g_indexed_topk_buffer;
|
||||
static id<MTLBuffer> g_f16_round_scratch_buffer;
|
||||
static id<MTLBuffer> g_q8_prefill_scratch_buffer;
|
||||
@@ -6408,6 +6411,15 @@ typedef struct {
|
||||
uint32_t block_top_k;
|
||||
} ds4_gpu_kargs_argsort_merge;
|
||||
|
||||
typedef struct {
|
||||
uint32_t n;
|
||||
uint32_t rows;
|
||||
uint32_t top_k;
|
||||
uint32_t pass;
|
||||
uint32_t parity;
|
||||
uint32_t n_tg;
|
||||
} ds4_gpu_kargs_topk_select;
|
||||
|
||||
typedef struct {
|
||||
int64_t ne00;
|
||||
int64_t ne01;
|
||||
@@ -19151,6 +19163,84 @@ int ds4_gpu_indexer_scores_decode_batch_tensor(
|
||||
scale);
|
||||
}
|
||||
|
||||
/* Radix select, deterministic compaction and one in-place sort of the winners
|
||||
* (kernel_topk_select_*): every dispatch spans the GPU, where the merge ladder
|
||||
* ends on two or three threadgroups. Returns 0 without encoding anything when
|
||||
* the pipelines are unavailable, so the caller keeps the ladder. */
|
||||
static int ds4_gpu_indexer_topk_select(id<MTLBuffer> selbuf, NSUInteger sel_off,
|
||||
id<MTLBuffer> scorebuf, NSUInteger score_off,
|
||||
uint32_t n, uint32_t rows, uint32_t top_k) {
|
||||
id<MTLComputePipelineState> hist = ds4_gpu_get_pipeline("kernel_topk_select_hist");
|
||||
id<MTLComputePipelineState> count = ds4_gpu_get_pipeline("kernel_topk_select_count");
|
||||
id<MTLComputePipelineState> write = ds4_gpu_get_pipeline("kernel_topk_select_write");
|
||||
id<MTLComputePipelineState> sort = ds4_gpu_get_pipeline("kernel_topk_select_sort");
|
||||
id<MTLComputePipelineState> zero = ds4_gpu_get_pipeline("kernel_topk_select_zero");
|
||||
if (!hist || !count || !write || !sort || !zero ||
|
||||
hist.maxTotalThreadsPerThreadgroup < 1024u || count.maxTotalThreadsPerThreadgroup < 1024u ||
|
||||
write.maxTotalThreadsPerThreadgroup < 1024u || sort.maxTotalThreadsPerThreadgroup < 1024u) return 0;
|
||||
const uint32_t n_tg = (n + 2047u) / 2048u;
|
||||
/* sized for the widest row (2^23 keys), so the scratch is allocated once per row count */
|
||||
const NSUInteger state_off = (NSUInteger)rows * 4u * 2048u * sizeof(uint32_t);
|
||||
const NSUInteger counts_off = state_off + (NSUInteger)rows * 4u * sizeof(uint32_t);
|
||||
const NSUInteger bytes = counts_off + (NSUInteger)rows * 4096u * 2u * sizeof(uint32_t);
|
||||
id<MTLBuffer> before = g_topk_select_buffer;
|
||||
if (!ds4_gpu_ensure_scratch_buffer(&g_topk_select_buffer, &g_topk_select_bytes, bytes,
|
||||
"ds4_topk_select")) return 0;
|
||||
ds4_gpu_kargs_topk_select args = {n, rows, top_k, 0, g_topk_select_parity, n_tg};
|
||||
g_topk_select_parity ^= 1u;
|
||||
int owned = 0;
|
||||
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
|
||||
if (!cb) return 0;
|
||||
id<MTLComputeCommandEncoder> enc;
|
||||
if (g_topk_select_buffer != before) { /* a fresh scratch starts with empty histograms */
|
||||
const uint32_t words = (uint32_t)(g_topk_select_bytes / sizeof(uint32_t));
|
||||
enc = ds4_gpu_compute_encoder(cb);
|
||||
[enc setComputePipelineState:zero];
|
||||
[enc setBuffer:g_topk_select_buffer offset:0 atIndex:0];
|
||||
[enc setBytes:&words length:sizeof(words) atIndex:1];
|
||||
[enc dispatchThreadgroups:MTLSizeMake((words + 255u) / 256u, 1, 1)
|
||||
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
|
||||
ds4_gpu_end_compute_encoder(cb, enc);
|
||||
}
|
||||
for (args.pass = 0; args.pass < 3u; args.pass++) {
|
||||
enc = ds4_gpu_compute_encoder(cb);
|
||||
[enc setComputePipelineState:hist];
|
||||
[enc setBytes:&args length:sizeof(args) atIndex:0];
|
||||
[enc setBuffer:scorebuf offset:score_off atIndex:1];
|
||||
[enc setBuffer:g_topk_select_buffer offset:0 atIndex:2];
|
||||
[enc setThreadgroupMemoryLength:2048u * sizeof(uint32_t) atIndex:0];
|
||||
[enc dispatchThreadgroups:MTLSizeMake(n_tg, rows, 1) threadsPerThreadgroup:MTLSizeMake(1024, 1, 1)];
|
||||
ds4_gpu_end_compute_encoder(cb, enc);
|
||||
}
|
||||
enc = ds4_gpu_compute_encoder(cb);
|
||||
[enc setComputePipelineState:count];
|
||||
[enc setBytes:&args length:sizeof(args) atIndex:0];
|
||||
[enc setBuffer:scorebuf offset:score_off atIndex:1];
|
||||
[enc setBuffer:g_topk_select_buffer offset:0 atIndex:2];
|
||||
[enc setBuffer:g_topk_select_buffer offset:state_off atIndex:3];
|
||||
[enc setBuffer:g_topk_select_buffer offset:counts_off atIndex:4];
|
||||
[enc dispatchThreadgroups:MTLSizeMake(n_tg, rows, 1) threadsPerThreadgroup:MTLSizeMake(1024, 1, 1)];
|
||||
ds4_gpu_end_compute_encoder(cb, enc);
|
||||
enc = ds4_gpu_compute_encoder(cb);
|
||||
[enc setComputePipelineState:write];
|
||||
[enc setBytes:&args length:sizeof(args) atIndex:0];
|
||||
[enc setBuffer:scorebuf offset:score_off atIndex:1];
|
||||
[enc setBuffer:g_topk_select_buffer offset:state_off atIndex:2];
|
||||
[enc setBuffer:g_topk_select_buffer offset:counts_off atIndex:3];
|
||||
[enc setBuffer:selbuf offset:sel_off atIndex:4];
|
||||
[enc dispatchThreadgroups:MTLSizeMake(n_tg, rows, 1) threadsPerThreadgroup:MTLSizeMake(1024, 1, 1)];
|
||||
ds4_gpu_end_compute_encoder(cb, enc);
|
||||
enc = ds4_gpu_compute_encoder(cb);
|
||||
[enc setComputePipelineState:sort];
|
||||
[enc setBytes:&args length:sizeof(args) atIndex:0];
|
||||
[enc setBuffer:scorebuf offset:score_off atIndex:1];
|
||||
[enc setBuffer:selbuf offset:sel_off atIndex:2];
|
||||
[enc setThreadgroupMemoryLength:2u * 2048u * sizeof(uint32_t) atIndex:0];
|
||||
[enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) threadsPerThreadgroup:MTLSizeMake(1024, 1, 1)];
|
||||
ds4_gpu_end_compute_encoder(cb, enc);
|
||||
return ds4_gpu_finish_command_buffer(cb, owned, "indexer top-k select") ? 1 : 0;
|
||||
}
|
||||
|
||||
static int ds4_gpu_indexer_topk_tensor_impl(
|
||||
ds4_gpu_tensor *selected,
|
||||
const ds4_gpu_tensor *scores,
|
||||
@@ -19182,6 +19272,11 @@ static int ds4_gpu_indexer_topk_tensor_impl(
|
||||
fprintf(stderr, "ds4: Metal graph indexer top-k received undersized buffers\n");
|
||||
return 0;
|
||||
}
|
||||
/* wide rows: the merge ladder's last rounds run on two or three threadgroups */
|
||||
if (!causal_ratio && top_k <= 2048u && n_comp >= 16384u && n_comp <= (1u << 23) &&
|
||||
!getenv("DS4_METAL_DISABLE_V41_TOPK_SELECT") &&
|
||||
ds4_gpu_indexer_topk_select(selbuf, ds4_gpu_tensor_offset(selected), scorebuf,
|
||||
ds4_gpu_tensor_offset(scores), n_comp, n_tokens, top_k)) return 1;
|
||||
NSUInteger max_threads = g_argsort_f32_i32_desc_pipeline.maxTotalThreadsPerThreadgroup;
|
||||
if (max_threads == 0) max_threads = 256;
|
||||
int32_t nth = 1;
|
||||
|
||||
@@ -318,3 +318,270 @@ kernel void kernel_argsort_merge_f32_i32(
|
||||
template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<DS4_SORT_ORDER_DESC>;
|
||||
template [[host_name("kernel_argsort_merge_f32_i32_desc_causal")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<DS4_SORT_ORDER_DESC, true>;
|
||||
template [[host_name("kernel_argsort_merge_f32_i32_desc_causal_prefix")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<DS4_SORT_ORDER_DESC, true, true>;
|
||||
|
||||
// ---- Whole-GPU top-k for wide score rows -----------------------------------
|
||||
//
|
||||
// kernel_argsort_f32_i32 and its merge ladder finish a 131072-wide index row
|
||||
// with two or three threadgroups doing the last merges, so one owner's
|
||||
// selection costs several hundred microseconds at long context. This path
|
||||
// keeps every core busy: a radix select over the order-preserving key bits
|
||||
// finds the k-th largest score in three histogram passes, a deterministic
|
||||
// compaction writes the winners by ascending id, and one threadgroup sorts
|
||||
// them by (score descending, id ascending). Values by rank equal the merge
|
||||
// path's; only exactly tied keys can change places. All kernels run with
|
||||
// 1024 threads; the host falls back to the merge ladder otherwise.
|
||||
|
||||
struct ds4_metal_args_topk_select {
|
||||
uint n; // scores per row
|
||||
uint rows;
|
||||
uint top_k; // at most 2048
|
||||
uint pass; // histogram pass 0, 1, 2 over 11 + 11 + 10 key bits
|
||||
uint parity; // pass-0 histogram slot this call fills; the other is cleared
|
||||
uint n_tg; // element threadgroups per row
|
||||
};
|
||||
|
||||
#define DS4_TOPK_BINS 2048u
|
||||
#define DS4_TOPK_PER_TG 2048u
|
||||
#define DS4_TOPK_HIST_WORDS (4u * DS4_TOPK_BINS) // hist0[2], hist1, hist2
|
||||
|
||||
// Unsigned key with the float's order: larger score, larger key.
|
||||
static inline uint ds4_topk_key(float f) {
|
||||
const uint u = as_type<uint>(f);
|
||||
return (u & 0x80000000u) ? ~u : (u | 0x80000000u);
|
||||
}
|
||||
|
||||
// The bin holding the k-th largest key and the count of keys above that bin.
|
||||
static inline uint2 ds4_topk_scan(device atomic_uint *h, uint bins, uint k,
|
||||
threadgroup uint *scratch,
|
||||
uint tid, uint lane, uint simd) {
|
||||
const uint per = bins / 1024u;
|
||||
uint c[2] = {0u, 0u}, local = 0u;
|
||||
for (uint q = 0; q < per; q++) {
|
||||
c[q] = atomic_load_explicit(h + bins - 1u - (tid * per + q), memory_order_relaxed);
|
||||
local += c[q];
|
||||
}
|
||||
const uint excl = simd_prefix_exclusive_sum(local);
|
||||
if (lane == 31u) scratch[simd] = excl + local;
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tid == 0u) {
|
||||
uint run = 0u;
|
||||
for (uint s = 0; s < 32u; s++) { const uint t = scratch[s]; scratch[s] = run; run += t; }
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
uint before = scratch[simd] + excl;
|
||||
for (uint q = 0; q < per; q++) {
|
||||
if (before < k && before + c[q] >= k) {
|
||||
scratch[32] = bins - 1u - (tid * per + q);
|
||||
scratch[33] = before;
|
||||
}
|
||||
before += c[q];
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
const uint2 r = uint2(scratch[32], scratch[33]);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
return r;
|
||||
}
|
||||
|
||||
// Walk the finished passes: (key prefix, count still wanted inside it, keys
|
||||
// above it). After three passes that is (threshold, wanted ties, greater).
|
||||
static inline uint3 ds4_topk_prefix(device atomic_uint *hist, uint parity, uint passes,
|
||||
uint top_k, threadgroup uint *scratch,
|
||||
uint tid, uint lane, uint simd) {
|
||||
uint prefix = 0u, k = top_k, above = 0u;
|
||||
if (passes > 0u) {
|
||||
const uint2 r = ds4_topk_scan(hist + parity * DS4_TOPK_BINS, 2048u, k, scratch, tid, lane, simd);
|
||||
prefix = r.x; k -= r.y; above += r.y;
|
||||
}
|
||||
if (passes > 1u) {
|
||||
const uint2 r = ds4_topk_scan(hist + 2u * DS4_TOPK_BINS, 2048u, k, scratch, tid, lane, simd);
|
||||
prefix = (prefix << 11) | r.x; k -= r.y; above += r.y;
|
||||
}
|
||||
if (passes > 2u) {
|
||||
const uint2 r = ds4_topk_scan(hist + 3u * DS4_TOPK_BINS, 1024u, k, scratch, tid, lane, simd);
|
||||
prefix = (prefix << 10) | r.x; k -= r.y; above += r.y;
|
||||
}
|
||||
return uint3(prefix, k, above);
|
||||
}
|
||||
|
||||
kernel void kernel_topk_select_zero(
|
||||
device uint * words,
|
||||
constant uint & count,
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
if (gid < count) words[gid] = 0u;
|
||||
}
|
||||
|
||||
// One histogram pass. Pass 0 also clears the other pass-0 slot and hist1 for
|
||||
// the next call and the next pass; pass 1 clears hist2.
|
||||
kernel void kernel_topk_select_hist(
|
||||
constant ds4_metal_args_topk_select & args,
|
||||
device const float * scores,
|
||||
device atomic_uint * hist,
|
||||
threadgroup atomic_uint * bins [[threadgroup(0)]],
|
||||
uint2 tgpig [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
uint lane [[thread_index_in_simdgroup]],
|
||||
uint simd [[simdgroup_index_in_threadgroup]]) {
|
||||
threadgroup uint scratch[34];
|
||||
const uint row = tgpig.y;
|
||||
device atomic_uint * h = hist + row * DS4_TOPK_HIST_WORDS;
|
||||
const uint3 sel = ds4_topk_prefix(h, args.parity, args.pass, args.top_k, scratch, tid, lane, simd);
|
||||
if (args.pass == 0u) {
|
||||
for (uint i = tgpig.x * 1024u + tid; i < 2u * DS4_TOPK_BINS; i += args.n_tg * 1024u) {
|
||||
atomic_store_explicit(h + (i < DS4_TOPK_BINS ? (args.parity ^ 1u) * DS4_TOPK_BINS + i :
|
||||
DS4_TOPK_BINS + i), 0u, memory_order_relaxed);
|
||||
}
|
||||
} else if (args.pass == 1u) {
|
||||
for (uint i = tgpig.x * 1024u + tid; i < DS4_TOPK_BINS; i += args.n_tg * 1024u)
|
||||
atomic_store_explicit(h + 3u * DS4_TOPK_BINS + i, 0u, memory_order_relaxed);
|
||||
}
|
||||
for (uint i = tid; i < DS4_TOPK_BINS; i += 1024u) atomic_store_explicit(bins + i, 0u, memory_order_relaxed);
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
device const float * s = scores + (ulong)row * args.n;
|
||||
for (uint j = 0; j < DS4_TOPK_PER_TG / 1024u; j++) {
|
||||
const uint i = tgpig.x * DS4_TOPK_PER_TG + j * 1024u + tid;
|
||||
if (i >= args.n) break;
|
||||
const uint key = ds4_topk_key(s[i]);
|
||||
uint bin;
|
||||
if (args.pass == 0u) bin = key >> 21;
|
||||
else if (args.pass == 1u) { if ((key >> 21) != sel.x) continue; bin = (key >> 10) & 0x7ffu; }
|
||||
else { if ((key >> 10) != sel.x) continue; bin = key & 0x3ffu; }
|
||||
atomic_fetch_add_explicit(bins + bin, 1u, memory_order_relaxed);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
device atomic_uint * out = h + (args.pass == 0u ? args.parity : args.pass + 1u) * DS4_TOPK_BINS;
|
||||
for (uint i = tid; i < DS4_TOPK_BINS; i += 1024u) {
|
||||
const uint c = atomic_load_explicit(bins + i, memory_order_relaxed);
|
||||
if (c) atomic_fetch_add_explicit(out + i, c, memory_order_relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
// Threshold, greater count and wanted ties for the row, plus this
|
||||
// threadgroup's greater/equal counts for the compaction offsets.
|
||||
kernel void kernel_topk_select_count(
|
||||
constant ds4_metal_args_topk_select & args,
|
||||
device const float * scores,
|
||||
device atomic_uint * hist,
|
||||
device uint * state,
|
||||
device uint * counts,
|
||||
uint2 tgpig [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
uint lane [[thread_index_in_simdgroup]],
|
||||
uint simd [[simdgroup_index_in_threadgroup]]) {
|
||||
threadgroup uint scratch[34], part[64];
|
||||
const uint row = tgpig.y;
|
||||
const uint3 sel = ds4_topk_prefix(hist + row * DS4_TOPK_HIST_WORDS, args.parity, 3u, args.top_k,
|
||||
scratch, tid, lane, simd);
|
||||
if (tgpig.x == 0u && tid == 0u) {
|
||||
state[row * 4u] = sel.x;
|
||||
state[row * 4u + 1u] = sel.z;
|
||||
state[row * 4u + 2u] = sel.y;
|
||||
}
|
||||
device const float * s = scores + (ulong)row * args.n;
|
||||
uint gt = 0u, eq = 0u;
|
||||
for (uint j = 0; j < DS4_TOPK_PER_TG / 1024u; j++) {
|
||||
const uint i = tgpig.x * DS4_TOPK_PER_TG + j * 1024u + tid;
|
||||
if (i >= args.n) break;
|
||||
const uint key = ds4_topk_key(s[i]);
|
||||
gt += key > sel.x;
|
||||
eq += key == sel.x;
|
||||
}
|
||||
gt = simd_sum(gt); eq = simd_sum(eq);
|
||||
if (lane == 0u) { part[simd] = gt; part[32u + simd] = eq; }
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tid == 0u) {
|
||||
uint tg_gt = 0u, tg_eq = 0u;
|
||||
for (uint q = 0; q < 32u; q++) { tg_gt += part[q]; tg_eq += part[32u + q]; }
|
||||
counts[(row * args.n_tg + tgpig.x) * 2u] = tg_gt;
|
||||
counts[(row * args.n_tg + tgpig.x) * 2u + 1u] = tg_eq;
|
||||
}
|
||||
}
|
||||
|
||||
// Winners by ascending id: keys above the threshold first, then the wanted
|
||||
// number of keys equal to it.
|
||||
kernel void kernel_topk_select_write(
|
||||
constant ds4_metal_args_topk_select & args,
|
||||
device const float * scores,
|
||||
device const uint * state,
|
||||
device const uint * counts,
|
||||
device int * selected,
|
||||
uint2 tgpig [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_index_in_threadgroup]],
|
||||
uint lane [[thread_index_in_simdgroup]],
|
||||
uint simd [[simdgroup_index_in_threadgroup]]) {
|
||||
threadgroup uint part[66];
|
||||
const uint row = tgpig.y;
|
||||
const uint threshold = state[row * 4u], above = state[row * 4u + 1u], wanted = state[row * 4u + 2u];
|
||||
uint gt_before = 0u, eq_before = 0u;
|
||||
for (uint t = 0; t < tgpig.x; t++) {
|
||||
gt_before += counts[(row * args.n_tg + t) * 2u];
|
||||
eq_before += counts[(row * args.n_tg + t) * 2u + 1u];
|
||||
}
|
||||
device const float * s = scores + (ulong)row * args.n;
|
||||
device int * out = selected + (ulong)row * args.top_k;
|
||||
for (uint j = 0; j < DS4_TOPK_PER_TG / 1024u; j++) {
|
||||
const uint i = tgpig.x * DS4_TOPK_PER_TG + j * 1024u + tid;
|
||||
const bool live = i < args.n;
|
||||
const uint key = live ? ds4_topk_key(s[i]) : 0u;
|
||||
const uint gt = live && key > threshold, eq = live && key == threshold;
|
||||
const uint gt_x = simd_prefix_exclusive_sum(gt), eq_x = simd_prefix_exclusive_sum(eq);
|
||||
if (lane == 31u) { part[simd] = gt_x + gt; part[32u + simd] = eq_x + eq; }
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (tid == 0u) {
|
||||
uint run_gt = 0u, run_eq = 0u;
|
||||
for (uint q = 0; q < 32u; q++) {
|
||||
const uint a = part[q], b = part[32u + q];
|
||||
part[q] = run_gt; part[32u + q] = run_eq;
|
||||
run_gt += a; run_eq += b;
|
||||
}
|
||||
part[64] = run_gt; part[65] = run_eq;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
if (gt) out[gt_before + part[simd] + gt_x] = (int)i;
|
||||
if (eq) {
|
||||
const uint r = eq_before + part[32u + simd] + eq_x;
|
||||
if (r < wanted) out[above + r] = (int)i;
|
||||
}
|
||||
gt_before += part[64]; eq_before += part[65];
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
}
|
||||
|
||||
// Sort the row's winners in place: score descending, ties by ascending id.
|
||||
kernel void kernel_topk_select_sort(
|
||||
constant ds4_metal_args_topk_select & args,
|
||||
device const float * scores,
|
||||
device int * selected,
|
||||
threadgroup uint * keys [[threadgroup(0)]],
|
||||
uint row [[threadgroup_position_in_grid]],
|
||||
uint tid [[thread_index_in_threadgroup]]) {
|
||||
threadgroup uint * ids = keys + 2048u;
|
||||
device const float * s = scores + (ulong)row * args.n;
|
||||
device int * out = selected + (ulong)row * args.top_k;
|
||||
for (uint i = tid; i < 2048u; i += 1024u) {
|
||||
if (i < args.top_k) {
|
||||
const uint id = (uint)out[i];
|
||||
keys[i] = ds4_topk_key(s[id]);
|
||||
ids[i] = id;
|
||||
} else {
|
||||
keys[i] = 0u;
|
||||
ids[i] = 0xffffffffu;
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (uint k = 2u; k <= 2048u; k <<= 1) {
|
||||
for (uint j = k >> 1; j > 0u; j >>= 1) {
|
||||
for (uint i = tid; i < 2048u; i += 1024u) {
|
||||
const uint o = i ^ j;
|
||||
if (o <= i) continue;
|
||||
const bool first = keys[i] > keys[o] || (keys[i] == keys[o] && ids[i] < ids[o]);
|
||||
if (((i & k) == 0u) != first) {
|
||||
const uint tk = keys[i], ti = ids[i];
|
||||
keys[i] = keys[o]; ids[i] = ids[o];
|
||||
keys[o] = tk; ids[o] = ti;
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
}
|
||||
for (uint i = tid; i < args.top_k; i += 1024u) out[i] = (int)ids[i];
|
||||
}
|
||||
@@ -1085,6 +1085,66 @@ static int check_general_topk(void) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
/* The radix-select path against the merge ladder: same values by rank, same
|
||||
* ids and order when the selected values are unique, ties by ascending id. */
|
||||
static int check_topk_select(void) {
|
||||
const uint32_t widths[] = {16384, 65536, 131072, 262144};
|
||||
const uint32_t counts[] = {1, 512, 2048};
|
||||
for (size_t wi = 0; wi < sizeof(widths) / sizeof(*widths); wi++) {
|
||||
const uint32_t width = widths[wi], rows = 2;
|
||||
ds4_gpu_tensor *scores = upload(NULL, (size_t)width * rows * 4);
|
||||
ds4_gpu_tensor *fast = upload(NULL, (size_t)2048 * rows * 4);
|
||||
ds4_gpu_tensor *slow = upload(NULL, (size_t)2048 * rows * 4);
|
||||
topk_entry *reference = malloc((size_t)width * rows * sizeof(*reference));
|
||||
bool *seen = calloc(width, sizeof(*seen));
|
||||
CHECK(scores && fast && slow && reference && seen);
|
||||
float *s = ds4_gpu_tensor_contents(scores);
|
||||
for (uint32_t pattern = 0; pattern < 4; pattern++) {
|
||||
for (uint32_t t = 0; t < rows; t++) {
|
||||
for (uint32_t i = 0; i < width; i++) {
|
||||
/* 0: coarse random values with many ties; 1: unique; 2: unique
|
||||
* with seven of eight blocks masked like the candidate filter;
|
||||
* 3: fewer finite values than 2048, so ties at -inf are selected */
|
||||
const float unique = (float)((i * 7919u + t * 1009u) % 1048573u) - 500000.0f;
|
||||
s[(size_t)t * width + i] = pattern == 0 ? random_value() : pattern == 1 ? unique :
|
||||
pattern == 2 ? (i / 8 % 8 ? -INFINITY : unique) : (i % 1361 ? -INFINITY : unique);
|
||||
reference[(size_t)t * width + i] = (topk_entry){s[(size_t)t * width + i], i};
|
||||
}
|
||||
qsort(reference + (size_t)t * width, width, sizeof(*reference), compare_topk_entry);
|
||||
}
|
||||
for (size_t ki = 0; ki < sizeof(counts) / sizeof(*counts); ki++) {
|
||||
const uint32_t k = counts[ki];
|
||||
CHECK(unsetenv("DS4_METAL_DISABLE_V41_TOPK_SELECT") == 0);
|
||||
CHECK(ds4_gpu_indexer_topk_tensor(fast, scores, width, rows, k));
|
||||
CHECK(setenv("DS4_METAL_DISABLE_V41_TOPK_SELECT", "1", 1) == 0);
|
||||
CHECK(ds4_gpu_indexer_topk_tensor(slow, scores, width, rows, k));
|
||||
CHECK(unsetenv("DS4_METAL_DISABLE_V41_TOPK_SELECT") == 0);
|
||||
CHECK(ds4_gpu_synchronize());
|
||||
const int32_t *a = ds4_gpu_tensor_contents(fast), *b = ds4_gpu_tensor_contents(slow);
|
||||
for (uint32_t t = 0; t < rows; t++) {
|
||||
const float *row = s + (size_t)t * width;
|
||||
const topk_entry *ref = reference + (size_t)t * width;
|
||||
memset(seen, 0, width * sizeof(*seen));
|
||||
for (uint32_t i = 0; i < k; i++) {
|
||||
const int32_t id = a[t * k + i];
|
||||
CHECK(id >= 0 && (uint32_t)id < width && !seen[id]);
|
||||
seen[id] = true;
|
||||
CHECK(row[id] == ref[i].value);
|
||||
CHECK(row[id] == row[b[t * k + i]]);
|
||||
if (i && row[id] == row[a[t * k + i - 1]]) CHECK(a[t * k + i - 1] < id);
|
||||
}
|
||||
if (pattern == 1 || pattern == 2)
|
||||
CHECK(!memcmp(a + t * k, b + t * k, k * sizeof(*a)));
|
||||
}
|
||||
}
|
||||
}
|
||||
fprintf(stderr, "V4.1 radix top-k width=%u: values by rank, unique-ID and tie-order oracle PASS\n", width);
|
||||
ds4_gpu_tensor_free(scores); ds4_gpu_tensor_free(fast); ds4_gpu_tensor_free(slow);
|
||||
free(reference); free(seen);
|
||||
}
|
||||
return 1;
|
||||
}
|
||||
|
||||
static int check_causal_topk(void) {
|
||||
const uint32_t frontiers[] = {513, 514, 1023, 1024, 1025, 1535, 2047, 2048, 4095, 16383, 32767, 65535};
|
||||
const uint32_t counts[] = {1, 2, 31, 32, 33, 127, 128, 129};
|
||||
@@ -1111,6 +1171,10 @@ static int check_causal_topk(void) {
|
||||
}
|
||||
for (size_t ci = 0; ci < sizeof(counts) / sizeof(*counts); ci++) {
|
||||
const uint32_t rows = counts[ci];
|
||||
/* The causal kernels are the ladder's variants and agree with it
|
||||
* tie for tie; the radix path orders ties by id, so the reference
|
||||
* stays on the ladder. */
|
||||
CHECK(setenv("DS4_METAL_DISABLE_V41_TOPK_SELECT", "1", 1) == 0);
|
||||
CHECK(ds4_gpu_begin_commands());
|
||||
for (uint32_t t = 0; t < rows; t++) {
|
||||
const uint32_t visible = (start + t + 1u) / ratio;
|
||||
@@ -1120,6 +1184,7 @@ static int check_causal_topk(void) {
|
||||
ds4_gpu_tensor_free(out); ds4_gpu_tensor_free(sr);
|
||||
}
|
||||
CHECK(ds4_gpu_end_commands());
|
||||
CHECK(unsetenv("DS4_METAL_DISABLE_V41_TOPK_SELECT") == 0);
|
||||
ids[rows * 512u] = 123456;
|
||||
CHECK(ds4_gpu_dsv41_indexer_topk_batch(selected, scores, width, rows, start, ratio));
|
||||
CHECK(ds4_gpu_synchronize());
|
||||
@@ -1486,6 +1551,11 @@ int main(int argc, char **argv) {
|
||||
ds4_gpu_cleanup();
|
||||
return ok ? 0 : 1;
|
||||
}
|
||||
if (argc == 2 && !strcmp(argv[1], "--topk-select")) {
|
||||
const int ok = ds4_gpu_init() && check_topk_select();
|
||||
ds4_gpu_cleanup();
|
||||
return ok ? 0 : 1;
|
||||
}
|
||||
if (argc == 2 && !strcmp(argv[1], "--general-topk")) {
|
||||
const int ok = ds4_gpu_init() && check_general_topk();
|
||||
ds4_gpu_cleanup();
|
||||
@@ -1500,7 +1570,7 @@ int main(int argc, char **argv) {
|
||||
int ok = ds4_gpu_init() && check_router() && check_quantization() && check_engram() && check_rope_stride() && check_pool() &&
|
||||
check_candidates() && check_candidates_batch(16449, 17, 16400) &&
|
||||
check_candidates_batch(17017, 17, 17000) && check_sparse_gather() && check_indexer_batch() && check_index_score_wide() &&
|
||||
check_embedding() && check_index_projection() && check_general_topk() && check_causal_topk() && check_indexer_all() && check_compact_carry() && check_attention_output(false) &&
|
||||
check_embedding() && check_index_projection() && check_general_topk() && check_topk_select() && check_causal_topk() && check_indexer_all() && check_compact_carry() && check_attention_output(false) &&
|
||||
check_tp_attention();
|
||||
ds4_gpu_cleanup();
|
||||
return ok ? 0 : 1;
|
||||
|
||||
Reference in new issue
Block a user